From c67e464621ed2c87a7e42d66d7fe67b11e6cc3e1 Mon Sep 17 00:00:00 2001 From: Myles Scolnick Date: Thu, 23 Jul 2026 14:36:41 -0400 Subject: [PATCH] feat: sqlCompletion helper to register all completion sources at once Registering the CTE, alias-qualified, and unqualified column completion sources required three separate dialect.language.data.of({ autocomplete }) calls. Add a sqlCompletion({ dialect, schema, parser, contextAnalyzer }) helper that wires up all three, sharing a single parser and query-context analyzer, with per-source enable toggles. Update the demo to use it. --- demo/index.ts | 21 +--- src/__tests__/index.test.ts | 1 + src/index.ts | 4 + .../__tests__/completion-extension.test.ts | 56 +++++++++ src/sql/completion-extension.ts | 111 ++++++++++++++++++ 5 files changed, 178 insertions(+), 15 deletions(-) create mode 100644 src/sql/__tests__/completion-extension.test.ts create mode 100644 src/sql/completion-extension.ts diff --git a/demo/index.ts b/demo/index.ts index 7fdf65a..90cca90 100644 --- a/demo/index.ts +++ b/demo/index.ts @@ -4,15 +4,13 @@ import { Compartment, type EditorState, StateEffect, StateField } from "@codemir import { keymap } from "@codemirror/view"; import { basicSetup, EditorView } from "codemirror"; import { - aliasColumnCompletionSource, - createCteCompletionSource, DefaultSqlTooltipRenders, defaultSqlHoverTheme, NodeSqlParser, QueryContextAnalyzer, type SupportedDialects, + sqlCompletion, sqlExtension, - unqualifiedColumnCompletionSource, } from "../src/index.js"; import { tableTooltipRenderer } from "./custom-renderers.js"; import { defaultSqlDoc, schema } from "./data.js"; @@ -185,18 +183,11 @@ function initializeEditor() { parser, }, }), - defaultDialect.language.data.of({ - // Statement-scoped CTE names and their output columns - autocomplete: createCteCompletionSource({ parser }), - }), - defaultDialect.language.data.of({ - // Complete `u.` -> columns of `users` in `SELECT ... FROM users u` - autocomplete: aliasColumnCompletionSource({ schema, parser, contextAnalyzer }), - }), - defaultDialect.language.data.of({ - // Complete `SELECT e` -> `email` because `FROM users` is in the statement - autocomplete: unqualifiedColumnCompletionSource({ schema, parser, contextAnalyzer }), - }), + // Register all schema-aware completion sources at once: + // - CTE names and their output columns + // - `u.` -> columns of `users` in `SELECT ... FROM users u` + // - `SELECT e` -> `email` because `FROM users` is in the statement + sqlCompletion({ dialect: defaultDialect, schema, parser, contextAnalyzer }), // Custom theme for better SQL editing EditorView.theme({ "&": { diff --git a/src/__tests__/index.test.ts b/src/__tests__/index.test.ts index 2eab8ea..a784b86 100644 --- a/src/__tests__/index.test.ts +++ b/src/__tests__/index.test.ts @@ -23,6 +23,7 @@ describe("index.ts exports", () => { "gotoSqlDefinition", "renameSqlIdentifier", "resolveSqlSchema", + "sqlCompletion", "sqlExtension", "sqlGotoDefinition", "sqlHighlightReferences", diff --git a/src/index.ts b/src/index.ts index 3f924c0..e9fb836 100644 --- a/src/index.ts +++ b/src/index.ts @@ -6,6 +6,10 @@ export { type ColumnCompletionConfig, unqualifiedColumnCompletionSource, } from "./sql/column-completion-source.js"; +export { + type SqlCompletionConfig, + sqlCompletion, +} from "./sql/completion-extension.js"; export { createCteCompletionSource, type CteCompletionConfig, diff --git a/src/sql/__tests__/completion-extension.test.ts b/src/sql/__tests__/completion-extension.test.ts new file mode 100644 index 0000000..ea4741a --- /dev/null +++ b/src/sql/__tests__/completion-extension.test.ts @@ -0,0 +1,56 @@ +import type { CompletionSource } from "@codemirror/autocomplete"; +import { PostgreSQL, sql } from "@codemirror/lang-sql"; +import { EditorState } from "@codemirror/state"; +import { describe, expect, it } from "vitest"; +import { sqlCompletion } from "../completion-extension.js"; + +const schema = { users: ["id", "name", "email"] }; + +/** Collect the `autocomplete` completion sources registered on the language. */ +function autocompleteSources(...extensions: Parameters[0]["extensions"][]) { + const state = EditorState.create({ + doc: "SELECT ", + extensions: [sql({ dialect: PostgreSQL, schema }), ...extensions], + }); + return state + .languageDataAt("autocomplete", state.doc.length) + .filter((value) => typeof value === "function"); +} + +describe("sqlCompletion", () => { + it("registers all three completion sources by default", () => { + // Baseline: lang-sql registers its own schema-based autocomplete source + const baseline = autocompleteSources().length; + const withHelper = autocompleteSources( + sqlCompletion({ dialect: PostgreSQL, schema }), + ).length; + expect(withHelper - baseline).toBe(3); + }); + + it("omits sources that are disabled", () => { + const baseline = autocompleteSources().length; + const withHelper = autocompleteSources( + sqlCompletion({ + dialect: PostgreSQL, + schema, + enableAliasCompletion: false, + enableColumnCompletion: false, + }), + ).length; + expect(withHelper - baseline).toBe(1); + }); + + it("returns no extra sources when everything is disabled", () => { + const baseline = autocompleteSources().length; + const withHelper = autocompleteSources( + sqlCompletion({ + dialect: PostgreSQL, + schema, + enableCteCompletion: false, + enableAliasCompletion: false, + enableColumnCompletion: false, + }), + ).length; + expect(withHelper - baseline).toBe(0); + }); +}); diff --git a/src/sql/completion-extension.ts b/src/sql/completion-extension.ts new file mode 100644 index 0000000..dac8afc --- /dev/null +++ b/src/sql/completion-extension.ts @@ -0,0 +1,111 @@ +import type { SQLDialect } from "@codemirror/lang-sql"; +import type { Extension } from "@codemirror/state"; +import { aliasColumnCompletionSource } from "./alias-completion-source.js"; +import { unqualifiedColumnCompletionSource } from "./column-completion-source.js"; +import { createCteCompletionSource } from "./cte-completion-source.js"; +import { NodeSqlParser } from "./parser.js"; +import { QueryContextAnalyzer } from "./query-context.js"; +import type { SqlSchemaSource } from "./schema-facet.js"; +import type { SqlParser } from "./types.js"; + +/** + * Configuration for {@link sqlCompletion}, the convenience helper that + * registers every schema-aware SQL completion source at once. + */ +export interface SqlCompletionConfig { + /** + * The SQL dialect whose language the completion sources are registered on + * (e.g. `PostgreSQL`, or a dialect from `./dialects`). This must match the + * dialect passed to `sql({ dialect })` for the sources to activate. + */ + dialect: SQLDialect; + /** + * Database schema to complete columns from. Falls back to the shared + * `sqlSchemaFacet` when not provided. Not used by CTE completion, which + * derives columns from the statement itself. + */ + schema?: SqlSchemaSource; + /** + * Custom SQL parser shared by all completion sources. Defaults to a new + * `NodeSqlParser`. Pass the same instance used by the linter/hover so + * dialect-specific setups only configure the parser once. + */ + parser?: SqlParser; + /** + * Query-context analyzer shared by all completion sources, so each edit is + * analyzed once. Defaults to one built from `parser`. + */ + contextAnalyzer?: QueryContextAnalyzer; + + /** Whether to enable CTE name/column completion (default: true) */ + enableCteCompletion?: boolean; + /** Whether to enable alias-qualified column completion (default: true) */ + enableAliasCompletion?: boolean; + /** Whether to enable unqualified column completion (default: true) */ + enableColumnCompletion?: boolean; +} + +/** + * Registers every schema-aware SQL completion source in one call, so you don't + * have to wire up each `dialect.language.data.of({ autocomplete })` by hand: + * - {@link createCteCompletionSource} — CTE names and their output columns + * - {@link aliasColumnCompletionSource} — `u.` → columns of `users` in + * `SELECT ... FROM users u` + * - {@link unqualifiedColumnCompletionSource} — `SELECT e` → `email` from the + * statement's FROM/JOIN tables + * + * A single parser and query-context analyzer are shared across the sources so + * each edit is analyzed only once. This complements `sqlExtension`, which + * covers linting, hover, gutter, and navigation but not completion. + * + * @example + * ```ts + * import { sql, PostgreSQL } from '@codemirror/lang-sql'; + * import { sqlCompletion } from '@marimo-team/codemirror-sql'; + * + * const schema = { users: ['id', 'name', 'email'] }; + * const extensions = [ + * sql({ dialect: PostgreSQL, schema }), + * sqlCompletion({ dialect: PostgreSQL, schema }), + * ]; + * ``` + */ +export function sqlCompletion(config: SqlCompletionConfig): Extension[] { + const { + dialect, + schema, + parser = new NodeSqlParser(), + enableCteCompletion = true, + enableAliasCompletion = true, + enableColumnCompletion = true, + } = config; + const contextAnalyzer = config.contextAnalyzer ?? new QueryContextAnalyzer(parser); + + const extensions: Extension[] = []; + + if (enableCteCompletion) { + extensions.push( + dialect.language.data.of({ + autocomplete: createCteCompletionSource({ parser, contextAnalyzer }), + }), + ); + } + + if (enableAliasCompletion) { + extensions.push( + dialect.language.data.of({ + autocomplete: aliasColumnCompletionSource({ schema, parser, contextAnalyzer }), + }), + ); + } + + if (enableColumnCompletion) { + extensions.push( + dialect.language.data.of({ + autocomplete: unqualifiedColumnCompletionSource({ schema, parser, contextAnalyzer }), + }), + ); + } + + return extensions; +}