feat: sqlCompletion helper to register all completion sources at once

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({
       "&": {

src/__tests__/index.test.ts

@@ -23,6 +23,7 @@ describe("index.ts exports", () => {
         "gotoSqlDefinition",
         "renameSqlIdentifier",
         "resolveSqlSchema",
+        "sqlCompletion",
         "sqlExtension",
         "sqlGotoDefinition",
         "sqlHighlightReferences",

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,

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<typeof EditorState.create>[0]["extensions"][]) {
+  const state = EditorState.create({
+    doc: "SELECT ",
+    extensions: [sql({ dialect: PostgreSQL, schema }), ...extensions],
+  });
+  return state
+    .languageDataAt<CompletionSource>("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);
+  });
+});

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;
+}