pgLab/pglablib/SqlCompletionProvider.cpp

#include "SqlCompletionProvider.h"
#include "SqlFormattingUtils.h"
#include "catalog/PgDatabaseCatalog.h"
#include "catalog/PgAttributeContainer.h"
#include "catalog/PgClassContainer.h"
#include "catalog/PgNamespaceContainer.h"
#include "catalog/PgProcContainer.h"
#include "catalog/PgConstraintContainer.h"
#include <set>
#include <algorithm>
#include <map>

namespace {

bool matchesPrefix(const QString &name, const SqlContext &ctx)
{
	return name.startsWith(ctx.prefix, ctx.prefixWasQuoted ? Qt::CaseSensitive : Qt::CaseInsensitive);
}

bool isQueryable(char kind)
{
	return kind == 'r' || kind == 'v' || kind == 'm' || kind == 'f' || kind == 'p';
}

QString kindDescription(char kind)
{
	switch (kind) {
	case 'v': return QStringLiteral("view");
	case 'm': return QStringLiteral("materialized view");
	case 'f': return QStringLiteral("foreign table");
	default: return QStringLiteral("table");
	}
}

// The name used to refer to a relation in the query
QString referenceName(const RelationRef &ref)
{
	return ref.alias.isEmpty() ? ref.name : ref.alias;
}

} // namespace

CompletionCatalog buildCompletionCatalog(const PgDatabaseCatalog &catalog)
{
	CompletionCatalog cc;
	auto namespaces = catalog.namespaces();
	auto classes = catalog.classes();
	auto attributes = catalog.attributes();

	if (namespaces) {
		for (const auto &ns : *namespaces) {
			const QString name = ns.objectName();
			if (name == QLatin1String("pg_toast") || name.startsWith(QLatin1String("pg_temp")) || name.startsWith(QLatin1String("pg_toast_temp")))
				continue;
			cc.schemas.push_back(name);
		}
	}
	if (classes) {
		for (const auto &cls : *classes) {
			char kind = '\0';
			switch (cls.kind) {
			case RelKind::Table: kind = 'r'; break;
			case RelKind::View: kind = 'v'; break;
			case RelKind::MaterializedView: kind = 'm'; break;
			case RelKind::ForeignTable: kind = 'f'; break;
			case RelKind::PartitionedTable: kind = 'p'; break;
			default: break;
			}
			if (kind == '\0')
				continue;
			CompletionRelation rel;
			rel.schema = cls.nsName();
			rel.name = cls.objectName();
			rel.kind = kind;
			if (attributes) {
				for (const auto &att : attributes->getColumnsForRelation(cls.oid()))
					if (att.num > 0 && !att.isdropped)
						rel.columns.push_back(att.name);
			}
			cc.relations.push_back(std::move(rel));
		}
	}
	auto procs = catalog.procs();
	if (procs) {
		for (const auto &proc : *procs) {
			char kind = 'f';
			if (proc.isAggregate()) kind = 'a';
			else if (proc.isWindow()) kind = 'w';
			else if (!proc.isFunction()) continue;
			cc.functions.push_back({ proc.nsName(), proc.objectName(), kind });
		}
	}
	auto constraints = catalog.constraints();
	if (constraints && classes && attributes) {
		auto columnNames = [&](Oid relid, const auto &keys) {
			std::vector<QString> names;
			for (auto num : keys) {
				const PgAttribute *att = attributes->getByKey(PgAttribute::Key(relid, num));
				names.push_back(att ? att->name : QString());
			}
			return names;
		};
		for (const auto &con : *constraints) {
			if (con.type != ConstraintType::ForeignKey || con.key.size() != con.fkey.size())
				continue;
			const PgClass *child = classes->getByKey(con.relid);
			const PgClass *parent = classes->getByKey(con.frelid);
			if (!child || !parent)
				continue;
			CompletionForeignKey fk;
			fk.childSchema = child->nsName();
			fk.childTable = child->objectName();
			fk.parentSchema = parent->nsName();
			fk.parentTable = parent->objectName();
			fk.childColumns = columnNames(con.relid, con.key);
			fk.parentColumns = columnNames(con.frelid, con.fkey);
			if (std::find(fk.childColumns.begin(), fk.childColumns.end(), QString()) != fk.childColumns.end()
				|| std::find(fk.parentColumns.begin(), fk.parentColumns.end(), QString()) != fk.parentColumns.end())
				continue;
			cc.foreignKeys.push_back(std::move(fk));
		}
	}
	return cc;
}

SqlCompletionProvider::SqlCompletionProvider(CompletionCatalog catalog)
	: m_catalog(std::move(catalog))
{
}

const CompletionRelation* SqlCompletionProvider::resolve(const RelationRef &ref) const
{
	if (ref.derived)
		return nullptr;
	if (!ref.schema.isEmpty()) {
		for (const auto &r : m_catalog.relations)
			if (r.name == ref.name && r.schema == ref.schema)
				return &r;
		return nullptr;
	}
	for (const auto &schema : m_catalog.searchPath)
		for (const auto &r : m_catalog.relations)
			if (r.name == ref.name && r.schema == schema)
				return &r;
	for (const auto &r : m_catalog.relations)
		if (r.name == ref.name)
			return &r;
	return nullptr;
}

std::vector<CompletionItem> SqlCompletionProvider::complete(const SqlContext &ctx) const
{
	switch (ctx.kind) {
	case CompletionKind::Relation:
		return relationItems(ctx, ctx.qualifier);
	case CompletionKind::JoinTarget: {
		auto items = joinItems(ctx);
		auto rest = relationItems(ctx, QString());
		items.insert(items.end(), std::make_move_iterator(rest.begin()), std::make_move_iterator(rest.end()));
		return items;
	}
	case CompletionKind::Column:
		return columnItems(ctx);
	case CompletionKind::QualifiedColumn:
		return qualifiedColumnItems(ctx);
	default:
		return {};
	}
}

std::vector<CompletionItem> SqlCompletionProvider::relationItems(const SqlContext &ctx, const QString &schema) const
{
	std::vector<CompletionItem> items;
	auto addRelation = [&](const CompletionRelation &r, bool qualified) {
		if (!isQueryable(r.kind) || !matchesPrefix(r.name, ctx))
			return;
		CompletionItem item;
		item.display = r.name;
		item.detail = qualified ? kindDescription(r.kind) : r.schema + QStringLiteral(" ") + kindDescription(r.kind);
		item.insertText = quoteIdent(r.name);
		item.kind = CompletionKind::Relation;
		items.push_back(std::move(item));
	};

	if (!schema.isEmpty()) {
		for (const auto &r : m_catalog.relations)
			if (r.schema == schema)
				addRelation(r, true);
	}
	else {
		for (const auto &s : m_catalog.schemas) {
			if (!matchesPrefix(s, ctx))
				continue;
			CompletionItem item;
			item.display = s;
			item.detail = QStringLiteral("schema");
			item.insertText = quoteIdent(s);
			item.kind = CompletionKind::Schema;
			items.push_back(std::move(item));
		}
		for (const auto &sp : m_catalog.searchPath)
			for (const auto &r : m_catalog.relations)
				if (r.schema == sp)
					addRelation(r, false);
	}
	auto byName = [](const CompletionItem &a, const CompletionItem &b) {
		if (a.kind != b.kind)
			return a.kind == CompletionKind::Schema; // schemas first
		return QString::compare(a.display, b.display, Qt::CaseInsensitive) < 0;
	};
	std::stable_sort(items.begin(), items.end(), byName);
	return items;
}

std::vector<CompletionItem> SqlCompletionProvider::columnItems(const SqlContext &ctx) const
{
	struct Source {
		const RelationRef *ref;
		const CompletionRelation *rel;
	};
	std::vector<Source> sources;
	for (const auto &ref : ctx.relations) {
		const CompletionRelation *rel = resolve(ref);
		if (rel)
			sources.push_back({ &ref, rel });
	}

	std::map<QString, int> counts; // number of in scope relations having this column
	for (const auto &s : sources)
		for (const auto &c : s.rel->columns)
			++counts[c];

	std::vector<CompletionItem> items;
	for (const auto &s : sources) {
		const QString refName = referenceName(*s.ref);
		for (const auto &c : s.rel->columns) {
			if (!matchesPrefix(c, ctx))
				continue;
			CompletionItem item;
			item.display = c;
			item.detail = refName;
			item.ambiguous = counts[c] > 1;
			item.insertText = item.ambiguous ? quoteIdent(refName) + QLatin1Char('.') + quoteIdent(c) : quoteIdent(c);
			item.kind = CompletionKind::Column;
			items.push_back(std::move(item));
		}
	}

	addFunctionItems(ctx, items);

	// names usable as qualifier
	std::vector<QString> seen;
	for (const auto &ref : ctx.relations) {
		const QString refName = referenceName(ref);
		if (!matchesPrefix(refName, ctx) || std::find(seen.begin(), seen.end(), refName) != seen.end())
			continue;
		seen.push_back(refName);
		CompletionItem item;
		item.display = refName;
		item.detail = ref.alias.isEmpty() ? QStringLiteral("relation") : QStringLiteral("alias of ") + ref.name;
		item.insertText = quoteIdent(refName);
		item.kind = CompletionKind::Relation;
		items.push_back(std::move(item));
	}
	return items;
}

namespace {

QString initialsOf(const QString &name)
{
	QString result;
	bool startOfWord = true;
	for (QChar c : name) {
		if (c.isLetterOrNumber()) {
			if (startOfWord)
				result += c.toLower();
			startOfWord = false;
		}
		else
			startOfWord = true;
	}
	return result.isEmpty() ? QStringLiteral("t") : result;
}

QString uniqueAlias(const QString &tableName, const std::set<QString> &taken)
{
	const QString base = initialsOf(tableName);
	QString candidate = base;
	int n = 2;
	while (taken.count(candidate.toLower()) || identNeedsQuotes(candidate))
		candidate = base + QString::number(n++);
	return candidate;
}

} // namespace

std::vector<CompletionItem> SqlCompletionProvider::joinItems(const SqlContext &ctx) const
{
	std::vector<CompletionItem> items;
	std::set<QString> taken;
	for (const auto &ref : ctx.relations) {
		taken.insert(ref.name.toLower());
		if (!ref.alias.isEmpty())
			taken.insert(ref.alias.toLower());
	}

	auto isSearchPath = [&](const QString &schema) {
		return std::find(m_catalog.searchPath.begin(), m_catalog.searchPath.end(), schema) != m_catalog.searchPath.end();
	};

	for (const auto &ref : ctx.relations) {
		const CompletionRelation *rel = resolve(ref);
		if (!rel)
			continue;
		const QString refName = quoteIdent(referenceName(ref));

		for (const auto &fk : m_catalog.foreignKeys) {
			// both directions, a self referencing key gives both
			for (int dir = 0; dir < 2; ++dir) {
				const bool fromChild = dir == 0;
				const QString &fromSchema = fromChild ? fk.childSchema : fk.parentSchema;
				const QString &fromTable = fromChild ? fk.childTable : fk.parentTable;
				if (fromTable != rel->name || fromSchema != rel->schema)
					continue;
				const QString &toSchema = fromChild ? fk.parentSchema : fk.childSchema;
				const QString &toTable = fromChild ? fk.parentTable : fk.childTable;
				const auto &fromCols = fromChild ? fk.childColumns : fk.parentColumns;
				const auto &toCols = fromChild ? fk.parentColumns : fk.childColumns;
				if (!matchesPrefix(toTable, ctx))
					continue;

				const QString alias = uniqueAlias(toTable, taken);
				QString on;
				for (size_t i = 0; i < fromCols.size(); ++i) {
					if (i > 0)
						on += QStringLiteral(" AND ");
					on += refName + QLatin1Char('.') + quoteIdent(fromCols[i]) + QStringLiteral(" = ")
						+ alias + QLatin1Char('.') + quoteIdent(toCols[i]);
				}
				QString target = quoteIdent(toTable);
				if (!isSearchPath(toSchema))
					target = quoteIdent(toSchema) + QLatin1Char('.') + target;

				CompletionItem item;
				item.display = toTable + QLatin1Char(' ') + alias;
				item.detail = QStringLiteral("ON ") + on;
				item.insertText = target + QLatin1Char(' ') + alias + QStringLiteral(" ON ") + on;
				item.kind = CompletionKind::JoinTarget;
				const bool duplicate = std::any_of(items.begin(), items.end(), [&](const CompletionItem &i) {
					return i.insertText == item.insertText;
				});
				if (!duplicate)
					items.push_back(std::move(item));
			}
		}
	}
	return items;
}

void SqlCompletionProvider::addFunctionItems(const SqlContext &ctx, std::vector<CompletionItem> &items) const
{
	// the list of functions is long, wait until something has been typed
	if (ctx.prefix.isEmpty() || ctx.clause == SqlClause::InsertInto)
		return;
	std::set<QString> seen;
	for (const auto &f : m_catalog.functions) {
		if (std::find(m_catalog.searchPath.begin(), m_catalog.searchPath.end(), f.schema) == m_catalog.searchPath.end())
			continue;
		if (!matchesPrefix(f.name, ctx) || !seen.insert(f.name).second)
			continue;
		CompletionItem item;
		item.display = f.name;
		item.detail = f.kind == 'a' ? QStringLiteral("aggregate") : f.kind == 'w' ? QStringLiteral("window function") : QStringLiteral("function");
		item.insertText = quoteIdent(f.name) + QLatin1Char('(');
		item.kind = CompletionKind::Function;
		items.push_back(std::move(item));
	}
}

std::vector<CompletionItem> SqlCompletionProvider::qualifiedColumnItems(const SqlContext &ctx) const
{
	std::vector<CompletionItem> items;
	bool found = false;
	for (const auto &ref : ctx.relations) {
		bool match = false;
		if (!ref.alias.isEmpty())
			match = ref.alias == ctx.qualifier;
		else
			match = ref.name == ctx.qualifier && (ctx.qualifierSchema.isEmpty() || ref.schema == ctx.qualifierSchema);
		if (!match)
			continue;
		found = true;
		const CompletionRelation *rel = resolve(ref);
		if (!rel)
			continue;
		for (const auto &c : rel->columns) {
			if (!matchesPrefix(c, ctx))
				continue;
			CompletionItem item;
			item.display = c;
			item.detail = rel->name;
			item.insertText = quoteIdent(c);
			item.kind = CompletionKind::Column;
			items.push_back(std::move(item));
		}
	}
	if (!found) {
		// maybe the qualifier is a schema, or a table that is not mentioned in the FROM clause
		if (!ctx.qualifierSchema.isEmpty()) {
			for (const auto &r : m_catalog.relations) {
				if (r.schema != ctx.qualifierSchema || r.name != ctx.qualifier)
					continue;
				for (const auto &c : r.columns) {
					if (!matchesPrefix(c, ctx))
						continue;
					CompletionItem item;
					item.display = c;
					item.detail = r.name;
					item.insertText = quoteIdent(c);
					item.kind = CompletionKind::Column;
					items.push_back(std::move(item));
				}
			}
		}
		else
			return relationItems(ctx, ctx.qualifier);
	}
	return items;
}