The editor now has context aware completions

pglablib/CMakeLists.txt

@@ -31,6 +31,7 @@ add_library(pglablib STATIC
     ui/catalog/tables/TableTreeBuilder.cpp
     util.cpp
     SqlFormattingUtils.cpp
+    SqlCompletionProvider.cpp
     catalog/PgKeywordList.cpp
     QueryGenerator.cpp
     ResultSourceMapping.cpp
@@ -107,6 +108,7 @@ add_library(pglablib STATIC
     ui/catalog/tables/TableTreeBuilder.h
     util.h
     SqlFormattingUtils.h
+    SqlCompletionProvider.h
     catalog/PgCatalogTypes.h
     catalog/PgKeywordList.h
     QueryGenerator.h

pglablib/SqlCompletionProvider.cpp

@@ -0,0 +1,430 @@
+#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;
+}

pglablib/SqlCompletionProvider.h

@@ -0,0 +1,69 @@
+#ifndef SQLCOMPLETIONPROVIDER_H
+#define SQLCOMPLETIONPROVIDER_H
+
+#include "SqlContext.h"
+#include <QString>
+#include <memory>
+#include <vector>
+
+class PgDatabaseCatalog;
+
+struct CompletionItem {
+	QString display;     // text shown in the list
+	QString detail;      // dimmed extra information, for example the relation a column belongs to
+	QString insertText;  // replaces the typed prefix when accepted
+	CompletionKind kind = CompletionKind::None;
+	bool ambiguous = false; // name occurs in more than one relation in scope
+};
+
+/// Plain snapshot of the parts of the catalog needed for completion, so the provider does not need a database.
+struct CompletionRelation {
+	QString schema;
+	QString name;
+	char kind = 'r'; // pg_class.relkind
+	std::vector<QString> columns; // in column order, without dropped and system columns
+};
+
+struct CompletionFunction {
+	QString schema;
+	QString name;
+	char kind = 'f'; // pg_proc.prokind: f function, a aggregate, w window
+};
+
+struct CompletionForeignKey {
+	QString childSchema, childTable;   // table that has the constraint
+	QString parentSchema, parentTable; // referenced table
+	std::vector<QString> childColumns;
+	std::vector<QString> parentColumns; // same length as childColumns
+};
+
+struct CompletionCatalog {
+	std::vector<QString> schemas;
+	std::vector<CompletionRelation> relations;
+	std::vector<CompletionFunction> functions;
+	std::vector<CompletionForeignKey> foreignKeys;
+	/// Schemas that are searched when a relation is not schema qualified, in order
+	std::vector<QString> searchPath = { QStringLiteral("public"), QStringLiteral("pg_catalog") };
+};
+
+CompletionCatalog buildCompletionCatalog(const PgDatabaseCatalog &catalog);
+
+class SqlCompletionProvider {
+public:
+	SqlCompletionProvider() = default;
+	explicit SqlCompletionProvider(CompletionCatalog catalog);
+
+	std::vector<CompletionItem> complete(const SqlContext &ctx) const;
+
+private:
+	CompletionCatalog m_catalog;
+
+	const CompletionRelation* resolve(const RelationRef &ref) const;
+	std::vector<CompletionItem> relationItems(const SqlContext &ctx, const QString &schema) const;
+	std::vector<CompletionItem> columnItems(const SqlContext &ctx) const;
+	std::vector<CompletionItem> qualifiedColumnItems(const SqlContext &ctx) const;
+	std::vector<CompletionItem> joinItems(const SqlContext &ctx) const;
+	void addFunctionItems(const SqlContext &ctx, std::vector<CompletionItem> &items) const;
+};
+
+#endif // SQLCOMPLETIONPROVIDER_H