From 4e76f6eeda5c5989f1d6e98a7e3434c05f03c662 Mon Sep 17 00:00:00 2001 From: Shahmir Varqha Date: Mon, 27 Jul 2026 12:45:13 +0800 Subject: [PATCH 1/2] feat: add bounded semantic SQL completion --- README.md | 9 +- demo/index.html | 28 + docs/capability-charter.md | 4 +- docs/marimo-sql-migration.md | 25 +- docs/session-primitives.md | 14 +- implementation.md | 59 ++ scripts/worker-placement.mjs | 8 +- src/__tests__/column-completion.test.ts | 419 ++++++++++++- src/__tests__/column-query-site.test.ts | 101 +++- src/__tests__/cte-layout.test.ts | 58 +- src/__tests__/local-relation-site.test.ts | 21 + src/__tests__/query-output.test.ts | 215 +++++++ src/__tests__/query-site.test.ts | 123 +++- .../semantic-completion-corpus.test.ts | 81 +++ .../semantic-completion-performance.test.ts | 99 +++ src/__tests__/session.test.ts | 567 +++++++++++++++++- .../sql-editor-completion-gate.test.ts | 37 ++ src/column-completion.ts | 206 ++++++- src/column-query-site.ts | 211 ++++++- src/cte-layout.ts | 97 ++- src/index.ts | 2 + src/local-relation-site.ts | 165 ++++- src/query-output.ts | 266 ++++++++ src/query-site.ts | 167 +++++- src/relation-completion-types.ts | 21 +- src/relation-dialect.ts | 1 + src/session.ts | 124 +++- test/corpora/semantic-completion/bigquery.ts | 35 ++ test/corpora/semantic-completion/duckdb.ts | 35 ++ .../corpora/semantic-completion/postgresql.ts | 35 ++ test/corpora/semantic-completion/types.ts | 14 + test/worker-placement/README.md | 12 +- 32 files changed, 3143 insertions(+), 116 deletions(-) create mode 100644 implementation.md create mode 100644 src/__tests__/query-output.test.ts create mode 100644 src/__tests__/semantic-completion-corpus.test.ts create mode 100644 src/__tests__/semantic-completion-performance.test.ts create mode 100644 src/query-output.ts create mode 100644 test/corpora/semantic-completion/bigquery.ts create mode 100644 test/corpora/semantic-completion/duckdb.ts create mode 100644 test/corpora/semantic-completion/postgresql.ts create mode 100644 test/corpora/semantic-completion/types.ts diff --git a/README.md b/README.md index a351410..18064e7 100644 --- a/README.md +++ b/README.md @@ -10,7 +10,7 @@ Published as [`@marimo-team/codemirror-sql`](https://www.npmjs.com/package/@mari - **Document sessions** — open SQL documents, apply atomic text and context updates, and track opaque revisions - **Built-in dialects** — PostgreSQL, DuckDB, BigQuery, and Dremio - **Embedded regions** — mask non-SQL spans (for example notebook interpolations) in document coordinates -- **Schema-aware completion** — complete relations, namespaces, physical columns, aliases, correlated scopes, and visible CTEs +- **Schema-aware completion** — complete relations, namespaces, physical columns, inferred CTE/derived outputs, aliases, correlated scopes, set-operation arms, and DML targets - **CodeMirror integration** — cancellation-safe completion, disposable detail panels, atomic context updates, and an optional virtualized statement gutter - **Composable completion** — package results coexist with SQL keywords, functions, embedded-language sources, and host sources - **Isolated parsing** — optional browser-worker parser execution kept off the public API surface @@ -72,9 +72,10 @@ pnpm install pnpm dev ``` -The playground exercises relation, namespace, physical-column, CTE, correlated -scope, and `LATERAL` completion against an in-memory catalog. It also exposes -provider latency and catalog invalidation controls. +The playground exercises relation, namespace, physical-column, CTE-output, +derived/set, `USING`, DML, correlated-scope, and `LATERAL` completion against +an in-memory catalog. It also exposes provider latency and catalog invalidation +controls. ## Development diff --git a/demo/index.html b/demo/index.html index 9197f05..0c7fd90 100644 --- a/demo/index.html +++ b/demo/index.html @@ -125,6 +125,13 @@

Completion scenarios

CTE visibility WITH recent AS (...) SELECT * FROM rec| + + + + diff --git a/docs/capability-charter.md b/docs/capability-charter.md index 21bed60..8e1614e 100644 --- a/docs/capability-charter.md +++ b/docs/capability-charter.md @@ -238,8 +238,8 @@ count and estimated retained bytes. Provisional compressed bundle budgets: -- Framework-independent core with relation, namespace, and column completion: - 50 KiB +- Framework-independent core with relation, namespace, column, and local query + output completion: 54 KiB - Core plus CodeMirror adapter, excluding parser and dialect data: 75 KiB - Each ordinary dialect module: 25 KiB - Optional `node-sql-parser` chunk: no regression from its recorded baseline diff --git a/docs/marimo-sql-migration.md b/docs/marimo-sql-migration.md index 995bc29..9808d75 100644 --- a/docs/marimo-sql-migration.md +++ b/docs/marimo-sql-migration.md @@ -133,14 +133,13 @@ columns. standard now provides all three through separate bounded providers, including qualified and unqualified query-site completion, ambiguity handling, stable provenance, cancellation, and batched column work. -The remaining feature gaps are: - -- the dialect coverage described above; and -- output-column inference for CTEs and derived relations. standard completes CTE - relation names but deliberately does not send a visible CTE name to the - physical column provider. Keep the legacy column source for those sites, or - defer full source replacement, until projection/output-column inference - lands. +The remaining cutover gap is the dialect coverage described above. standard +now infers provable output names for CTEs, derived relations, and the first arm +of set operations. Explicit CTE column lists are authoritative. Simple column +references and explicit aliases are inferred locally; stars and unaliased +expressions remain explicitly partial rather than receiving invented names. +Visible CTE and derived relation names are never sent to the physical column +provider. The fixture feeds marimo's immutable namespace projection—stable entity ID, scope, canonical identifier path, and namespace kind—through one public, @@ -158,8 +157,8 @@ scoped namespace provider on the shared service. 4. Compare relation and column results with the golden corpus, including quoted insert text, aliases, ambiguity, partial/loading/failure states, and cold epoch behavior. -5. Cut over supported physical-relation sites while preserving variable, - keyword, and legacy CTE/derived-output column sources. +5. Cut over supported physical and inferred local-relation sites while + preserving variable and keyword sources. 6. Add dialect coverage, expand the router, and remove completion-only legacy schema code. Keep legacy schema data while hover or diagnostics still use it. @@ -185,9 +184,9 @@ Library tests cover relation and physical-column completion for `FROM`, `JOIN`, `alias.`, unqualified projections and predicates, `USING`, correlated nested queries, quoted identifiers, ambiguous columns, provider loading/invalidations, bounded batching, and template barriers. CTE tests prove -declaration-order relation visibility and that CTE names are not incorrectly -resolved through the physical column provider; they do not prove CTE -output-column inference. +declaration-order relation visibility, conservative output-column inference +including partial stars and unaliased expressions, and that CTE names are not +incorrectly resolved through the physical column provider. Marimo integration tests cover: diff --git a/docs/session-primitives.md b/docs/session-primitives.md index 4d96333..4fa51a0 100644 --- a/docs/session-primitives.md +++ b/docs/session-primitives.md @@ -5,8 +5,18 @@ Import: `@marimo-team/codemirror-sql` This entry point provides document ownership, atomic text/context updates, opaque revisions, dialect registration, lifecycle management, and -parser-independent relation completion. Diagnostics, hover, navigation, and -general expression completion are not yet available. +parser-independent relation and column completion. Relation sites include +`SELECT` set-operation arms and bounded `INSERT`, `UPDATE`, `DELETE`, and +`MERGE` targets. Column completion combines physical catalog columns with +locally inferred CTE and derived-query outputs. Diagnostics, hover, navigation, +and general expression completion are not yet available. + +Local output inference only publishes names it can prove: explicit CTE column +lists, explicit projection aliases, and simple column references. Set +operations take their output names from the first arm. Stars, unaliased +expressions, opaque templates without an explicit alias, and resource-limited +shapes make the result incomplete; the service does not guess engine-generated +column names or query a physical provider with a CTE/derived relation name. CodeMirror consumers should use the separate [standard CodeMirror adapter](./codemirror-adapter.md). diff --git a/implementation.md b/implementation.md new file mode 100644 index 0000000..f887a83 --- /dev/null +++ b/implementation.md @@ -0,0 +1,59 @@ +# SQL editor completion plan + +Status: active +Updated: 2026-07-27 + +The standard language-service overhaul is delivered by PR #204. Two additional +vertical PRs finish the remaining product roadmap without coupling CodeMirror +to parser ASTs, database clients, language-server transports, or formatter +implementations. + +## PR 1: semantic completion + +This PR closes the remaining marimo schema-source cutover blocker and improves +the most common relation sites: + +- infer provable CTE and derived-query output names; +- make explicit CTE column lists authoritative; +- use the first set-operation arm for output names; +- complete relations in later set-operation arms; +- complete bounded DML target/source relation sites; +- restrict `JOIN ... USING` completion to shared columns on the immediate + relation pair; +- preserve explicit partial evidence for stars, unaliased expressions, + templates, malformed SQL, and resource limits; +- add owned PostgreSQL, BigQuery, and DuckDB completion corpora; +- exercise the behavior through sessions, CodeMirror browser tests, the demo, + and a warm 10 KiB p95 performance gate. + +Local relation names never cross the physical column-provider boundary. +Completion provenance distinguishes inferred query outputs from catalog +columns. + +## PR 2: language intelligence and release hardening + +The final PR adds the remaining feature methods and provider composition: + +- syntax, semantic, and host diagnostics; +- hover; +- definition, references, highlights, and rename; +- document symbols and folding; +- parameter, function, table-function, and snippet contracts; +- formatting and code-action providers; +- bounded native-engine and LSP composition seams; +- deterministic property/fuzz infrastructure and a mutation pilot; +- heap, listener, worker, and multi-editor leak gates; +- Chromium, Firefox, and WebKit evidence; +- public API and bundle regression checks; +- a packed runtime marimo fixture. + +Optional DuckDB, language-server, and formatter integrations remain providers. +They augment the fast local baseline and cannot delay or replace unrelated +local evidence. + +## Quality gates + +Each PR must keep repository and changed-code coverage above 95%, strict type +checking, zero-warning lint, package/SSR/worker placement, browser tests, +security scanning, and explicit latency/bundle budgets green. Exactly two +independent adversarial reviews run only after implementation is complete. diff --git a/scripts/worker-placement.mjs b/scripts/worker-placement.mjs index e92ed44..353ed8e 100644 --- a/scripts/worker-placement.mjs +++ b/scripts/worker-placement.mjs @@ -35,11 +35,11 @@ const PARSER_MARKERS = [ "tableList", ]; const BIGQUERY_GZIP_LIMIT = 50 * 1024; -const CORE_TOTAL_GZIP_LIMIT = 50 * 1024; -const CORE_TOTAL_RAW_LIMIT = 184 * 1024; +const CORE_TOTAL_GZIP_LIMIT = 54 * 1024; +const CORE_TOTAL_RAW_LIMIT = 200 * 1024; const POSTGRESQL_GZIP_LIMIT = 68 * 1024; -const WORKER_TOTAL_GZIP_LIMIT = 160 * 1024; -const WORKER_TOTAL_RAW_LIMIT = 700 * 1024; +const WORKER_TOTAL_GZIP_LIMIT = 164 * 1024; +const WORKER_TOTAL_RAW_LIMIT = 720 * 1024; const MIME_TYPES = new Map([ [".css", "text/css; charset=utf-8"], [".html", "text/html; charset=utf-8"], diff --git a/src/__tests__/column-completion.test.ts b/src/__tests__/column-completion.test.ts index f86655c..391a4fd 100644 --- a/src/__tests__/column-completion.test.ts +++ b/src/__tests__/column-completion.test.ts @@ -1,11 +1,15 @@ import { describe, expect, it } from "vitest"; import { composeSqlColumnCompletion, + composeSqlLocalQueryOutputCompletion, + filterSqlUsingCompletionList, prepareSqlColumnCatalogRelations, } from "../column-completion.js"; import { MAX_COLUMN_BATCH_RELATIONS, + MAX_COLUMNS_PER_BATCH, } from "../column-catalog-boundary.js"; +import { MAX_QUERY_OUTPUT_COLUMNS } from "../query-output.js"; import type { SqlColumnCatalogBatchOutcome, } from "../column-catalog-batch-coordinator.js"; @@ -32,6 +36,7 @@ function site( ): ReadySite { return Object.freeze({ coverage: "complete", + context: "select-list", issues: Object.freeze([]), prefix: Object.freeze({ quoted: false, value: prefix }), qualifier: Object.freeze(qualifier), @@ -94,6 +99,194 @@ function usable( } describe("column completion", () => { + it("composes bounded local output evidence and degrades missing output", () => { + const local = (output?: ReadySite["relations"][number]["local"]) => + Object.freeze({ + ...site([{ quoted: false, value: "c" }], "i"), + relations: Object.freeze([ + Object.freeze({ + alias: Object.freeze({ quoted: false, value: "c" }), + ...(output === undefined ? {} : { local: output }), + path: Object.freeze([]), + range: Object.freeze({ from: 20, to: 30 }), + }), + ]), + }); + expect( + composeSqlLocalQueryOutputCompletion( + local(), + POSTGRESQL_SQL_RELATION_DIALECT, + ), + ).toBeNull(); + const missing = local(Object.freeze({ + kind: "cte", + queryRange: Object.freeze({ from: 0, to: 10 }), + })); + expect( + composeSqlLocalQueryOutputCompletion( + missing, + POSTGRESQL_SQL_RELATION_DIALECT, + ), + ).toMatchObject({ + sources: [{ coverage: "partial" }], + value: { + isIncomplete: true, + items: [], + }, + }); + const column = Object.freeze({ + definition: Object.freeze({ from: 5, to: 7 }), + identifier: Object.freeze({ quoted: false, value: "id" }), + insertText: "id", + }); + const ready = local(Object.freeze({ + kind: "cte", + output: Object.freeze({ + columns: Object.freeze([column, column]), + coverage: "complete", + status: "ready", + }), + queryRange: Object.freeze({ from: 0, to: 10 }), + })); + expect( + composeSqlLocalQueryOutputCompletion( + ready, + POSTGRESQL_SQL_RELATION_DIALECT, + ), + ).toMatchObject({ + sources: [{ coverage: "complete" }], + value: { items: [{ label: "id" }] }, + }); + expect( + composeSqlLocalQueryOutputCompletion( + Object.freeze({ + ...ready, + prefix: Object.freeze({ quoted: false, value: "z" }), + }), + POSTGRESQL_SQL_RELATION_DIALECT, + )?.value.items, + ).toEqual([]); + }); + + it("bounds aggregate local output candidates", () => { + const output = Object.freeze({ + columns: Object.freeze(Array.from( + { length: MAX_QUERY_OUTPUT_COLUMNS }, + (_, index) => Object.freeze({ + definition: Object.freeze({ from: index, to: index + 1 }), + identifier: Object.freeze({ + quoted: false, + value: `column_${index}`, + }), + insertText: `column_${index}`, + }), + )), + coverage: "complete" as const, + status: "ready" as const, + }); + const current = Object.freeze({ + ...site(), + relations: Object.freeze(Array.from( + { + length: + Math.ceil(MAX_COLUMNS_PER_BATCH / MAX_QUERY_OUTPUT_COLUMNS) + + 1, + }, + (_, index) => Object.freeze({ + alias: Object.freeze({ + quoted: false, + value: `cte_${index}`, + }), + local: Object.freeze({ + kind: "cte" as const, + output, + queryRange: Object.freeze({ from: 0, to: 10 }), + }), + path: Object.freeze([]), + range: Object.freeze({ from: index, to: index + 1 }), + }), + )), + }); + + expect( + composeSqlLocalQueryOutputCompletion( + current, + POSTGRESQL_SQL_RELATION_DIALECT, + ), + ).toMatchObject({ + sources: [{ coverage: "partial" }], + value: { + isIncomplete: true, + items: { length: MAX_COLUMNS_PER_BATCH }, + }, + }); + }); + + it("reserves local USING work for the immediate relation pair", () => { + const unrelatedOutput = Object.freeze({ + columns: Object.freeze(Array.from( + { length: MAX_QUERY_OUTPUT_COLUMNS }, + (_, index) => Object.freeze({ + definition: Object.freeze({ from: index, to: index + 1 }), + identifier: Object.freeze({ + quoted: false, + value: `unrelated_${index}`, + }), + insertText: `unrelated_${index}`, + }), + )), + coverage: "complete" as const, + status: "ready" as const, + }); + const sharedOutput = Object.freeze({ + columns: Object.freeze([Object.freeze({ + definition: Object.freeze({ from: 0, to: 1 }), + identifier: Object.freeze({ quoted: false, value: "id" }), + insertText: "id", + })]), + coverage: "complete" as const, + status: "ready" as const, + }); + const relation = ( + index: number, + output: typeof unrelatedOutput | typeof sharedOutput, + ) => Object.freeze({ + alias: Object.freeze({ quoted: false, value: `r_${index}` }), + local: Object.freeze({ + kind: "cte" as const, + output, + queryRange: Object.freeze({ from: 0, to: 10 }), + }), + path: Object.freeze([]), + range: Object.freeze({ from: index, to: index + 1 }), + }); + const current = Object.freeze({ + ...site(), + context: "using" as const, + relations: Object.freeze([ + ...Array.from( + { length: 17 }, + (_, index) => relation(index, unrelatedOutput), + ), + relation(17, sharedOutput), + relation(18, sharedOutput), + ]), + }); + const composition = composeSqlLocalQueryOutputCompletion( + current, + POSTGRESQL_SQL_RELATION_DIALECT, + ); + expect(composition).not.toBeNull(); + expect( + composition && + filterSqlUsingCompletionList( + composition.value, + current, + POSTGRESQL_SQL_RELATION_DIALECT, + ).items.map((item) => item.label), + ).toEqual(["id"]); + }); + it("bounds relation batches and reports the omitted bindings", () => { const relations = Array.from( { length: MAX_COLUMN_BATCH_RELATIONS + 1 }, @@ -176,10 +369,188 @@ describe("column completion", () => { prepareSqlColumnCatalogRelations( current, POSTGRESQL_SQL_RELATION_DIALECT, - ).references, + ).references, ).toHaveLength(1); }); + it("offers only the intersection of the immediate USING relations", () => { + const current = Object.freeze({ + ...site([], "i"), + context: "using" as const, + relations: Object.freeze([ + Object.freeze({ + alias: Object.freeze({ quoted: false, value: "u" }), + path: Object.freeze([ + Object.freeze({ quoted: false, value: "users" }), + ]), + range: Object.freeze({ from: 0, to: 5 }), + }), + Object.freeze({ + alias: Object.freeze({ quoted: false, value: "o" }), + path: Object.freeze([ + Object.freeze({ quoted: false, value: "orders" }), + ]), + range: Object.freeze({ from: 6, to: 12 }), + }), + Object.freeze({ + alias: Object.freeze({ quoted: false, value: "p" }), + path: Object.freeze([ + Object.freeze({ quoted: false, value: "payments" }), + ]), + range: Object.freeze({ from: 13, to: 21 }), + }), + ]), + }); + const prepared = prepareSqlColumnCatalogRelations( + current, + POSTGRESQL_SQL_RELATION_DIALECT, + ); + expect(prepared.references.map((reference) => + reference.path[0]?.value + )).toEqual(["orders", "payments"]); + + const result = composeSqlColumnCompletion({ + dialect: POSTGRESQL_SQL_RELATION_DIALECT, + outcome: usable([ + { + columns: [ + column("id", "orders", 0), + column("internal_id", "orders", 1), + ], + coverage: "complete", + relationEntityId: "orders", + requestKey: "binding:1", + status: "ready", + }, + { + columns: [ + column("id", "payments", 0), + column("invoice_id", "payments", 1), + ], + coverage: "complete", + relationEntityId: "payments", + requestKey: "binding:2", + status: "ready", + }, + ]), + prepared, + providerId: "columns", + site: current, + }); + + expect(result?.value.items.map((item) => item.label)).toEqual(["id"]); + }); + + it("keeps physical self-join identifiers on both USING sides", () => { + const current = Object.freeze({ + ...site(), + context: "using" as const, + relations: Object.freeze([ + Object.freeze({ + alias: Object.freeze({ quoted: false, value: "left_users" }), + path: Object.freeze([ + Object.freeze({ quoted: false, value: "users" }), + ]), + range: Object.freeze({ from: 0, to: 5 }), + }), + Object.freeze({ + alias: Object.freeze({ quoted: false, value: "right_users" }), + path: Object.freeze([ + Object.freeze({ quoted: false, value: "users" }), + ]), + range: Object.freeze({ from: 6, to: 12 }), + }), + ]), + }); + const prepared = prepareSqlColumnCatalogRelations( + current, + POSTGRESQL_SQL_RELATION_DIALECT, + ); + const shared = column("id", "users", 0); + const result = composeSqlColumnCompletion({ + dialect: POSTGRESQL_SQL_RELATION_DIALECT, + outcome: usable([ + { + columns: [shared], + coverage: "complete", + relationEntityId: "users", + requestKey: "binding:0", + status: "ready", + }, + { + columns: [shared], + coverage: "complete", + relationEntityId: "users", + requestKey: "binding:1", + status: "ready", + }, + ]), + prepared, + providerId: "columns", + site: current, + }); + + expect(result?.value.items.map((item) => item.label)).toEqual(["id"]); + }); + + it("preserves quoted identifier semantics in USING intersections", () => { + const current = Object.freeze({ + ...site(), + context: "using" as const, + }); + const prepared = prepareSqlColumnCatalogRelations( + current, + POSTGRESQL_SQL_RELATION_DIALECT, + ); + const quoted = Object.freeze({ + ...column("quoted-foo", "orders", 1), + identifier: Object.freeze({ quoted: true, value: "Foo" }), + insertText: '"Foo"', + }); + const result = composeSqlColumnCompletion({ + dialect: POSTGRESQL_SQL_RELATION_DIALECT, + outcome: usable([ + { + columns: [ + Object.freeze({ + ...column("unquoted-foo", "users", 0), + identifier: Object.freeze({ quoted: false, value: "Foo" }), + insertText: "Foo", + }), + Object.freeze({ + ...quoted, + provenance: Object.freeze({ + ...quoted.provenance, + relationEntityId: "users", + }), + }), + ], + coverage: "complete", + relationEntityId: "users", + requestKey: "binding:0", + status: "ready", + }, + { + columns: [quoted], + coverage: "complete", + relationEntityId: "orders", + requestKey: "binding:1", + status: "ready", + }, + ]), + prepared, + providerId: "columns", + site: current, + }); + + expect(result?.value.items).toMatchObject([ + { + edit: { insert: '"Foo"' }, + label: "Foo", + }, + ]); + }); + it("composes deterministic exact edits and provenance", () => { const current = site([], "na"); const prepared = prepareSqlColumnCatalogRelations( @@ -234,6 +605,52 @@ describe("column completion", () => { }); }); + it("applies physical relation-alias column lists positionally", () => { + const current = Object.freeze({ + ...site([{ quoted: false, value: "u" }]), + relations: Object.freeze([ + Object.freeze({ + ...site().relations[0]!, + columnAliases: Object.freeze({ + columns: Object.freeze([ + Object.freeze({ + definition: Object.freeze({ from: 30, to: 37 }), + identifier: Object.freeze({ + quoted: false, + value: "renamed", + }), + insertText: "renamed", + }), + ]), + coverage: "complete" as const, + }), + }), + ]), + }); + const prepared = prepareSqlColumnCatalogRelations( + current, + POSTGRESQL_SQL_RELATION_DIALECT, + ); + const result = composeSqlColumnCompletion({ + dialect: POSTGRESQL_SQL_RELATION_DIALECT, + outcome: usable([{ + columns: [column("original", "users", 0)], + coverage: "complete", + relationEntityId: "users", + requestKey: "binding:0", + status: "ready", + }]), + prepared, + providerId: "columns", + site: current, + }); + + expect(result?.value.items).toMatchObject([{ + edit: { insert: "renamed" }, + label: "renamed", + }]); + }); + it("reports partial, loading, and failed relation evidence", () => { const current = Object.freeze({ ...site(), diff --git a/src/__tests__/column-query-site.test.ts b/src/__tests__/column-query-site.test.ts index abddfc7..64637d0 100644 --- a/src/__tests__/column-query-site.test.ts +++ b/src/__tests__/column-query-site.test.ts @@ -254,7 +254,7 @@ describe("recognizeSqlColumnQuerySite", () => { expect(position).toBeGreaterThan(0); }); - it("marks derived and table-function evidence partial", () => { + it("represents derived relations and marks table functions partial", () => { expect( ready( analyze( @@ -263,14 +263,103 @@ describe("recognizeSqlColumnQuerySite", () => { ), ).toMatchObject({ coverage: "partial", - issues: [ - "derived-relation", - "nested-query", - "table-function", - ], + issues: ["table-function"], }); }); + it("keeps a bounded derived query source and its explicit alias", () => { + const result = ready( + analyze( + "SELECT d.i| FROM (SELECT id FROM users) AS d", + ), + ); + expect(result.relations).toMatchObject([{ + alias: { value: "d" }, + local: { + kind: "derived", + queryRange: { + from: "SELECT d.i FROM (".length, + to: "SELECT d.i FROM (SELECT id FROM users".length, + }, + }, + path: [], + }]); + }); + + it("preserves authoritative relation-alias column lists", () => { + const result = ready( + analyze( + 'SELECT d.| FROM (SELECT original, other) AS d("Renamed", second)', + ), + ); + expect(result.relations).toMatchObject([{ + alias: { value: "d" }, + columnAliases: { + columns: [ + { + identifier: { quoted: true, value: "Renamed" }, + insertText: '"Renamed"', + }, + { + identifier: { quoted: false, value: "second" }, + insertText: "second", + }, + ], + coverage: "complete", + }, + local: { kind: "derived" }, + }]); + }); + + it.each([ + "SELECT d.| FROM (SELECT original) d()", + "SELECT d.| FROM (SELECT original) d(schema.name)", + "SELECT d.| FROM (SELECT original) d((nested))", + "SELECT d.| FROM (SELECT original) d(name,)", + "SELECT d.| FROM (SELECT original) d(/*comment*/ name", + ])("marks malformed relation-alias column lists partial in %s", (marked) => { + expect(analyze(marked)).toMatchObject({ + coverage: "partial", + issues: ["local-output-partial"], + status: "ready", + }); + }); + + it("bounds relation-alias column lists", () => { + const columns = Array.from( + { length: 257 }, + (_, index) => `column_${index}`, + ).join(", "); + const result = ready( + analyze(`SELECT d.| FROM (SELECT original) d(${columns})`), + ); + expect(result).toMatchObject({ + coverage: "partial", + issues: ["local-output-partial"], + relations: [{ + columnAliases: { + columns: { length: 256 }, + coverage: "partial", + }, + }], + }); + }); + + it.each([ + "SELECT i| FROM (SELECT id FROM users)", + "SELECT i| FROM (VALUES (1)) v", + "SELECT i| FROM (SELECT id FROM users AS", + ])("marks an unprovable derived source partial in %s", (marked) => { + const result = analyze(marked); + expect(result).toMatchObject({ + coverage: "partial", + status: "ready", + }); + if (result.status === "ready") { + expect(result.issues).toContain("derived-relation"); + } + }); + it.each([ "SELECT * FROM users WHERE na|", "SELECT * FROM users GROUP BY na|", diff --git a/src/__tests__/cte-layout.test.ts b/src/__tests__/cte-layout.test.ts index 213e861..b7838d6 100644 --- a/src/__tests__/cte-layout.test.ts +++ b/src/__tests__/cte-layout.test.ts @@ -29,6 +29,7 @@ import { findSqlStatementSlot, type ExactSqlStatementSlot, } from "../statement-index.js"; +import { MAX_QUERY_OUTPUT_COLUMNS } from "../query-output.js"; const postgres = POSTGRESQL_SQL_RELATION_DIALECT.cteLayout; const duckdb = DUCKDB_SQL_RELATION_DIALECT.cteLayout; @@ -494,8 +495,8 @@ describe("bounded CTE layout", () => { "ambiguous-cte-header", ], ["WITH a AS (WITH) SELECT 1", "ambiguous-cte-header"], - ["WITH a AS (SELECT 1) DELETE FROM t", "unsupported-cte-extension"], ["WITH a AS (SELECT 1),", "ambiguous-cte-header"], + ["WITH a AS (SELECT 1) + SELECT 1", "unsupported-cte-extension"], [ "WITH a AS (SELECT 1), SELECT", "ambiguous-cte-header", @@ -529,7 +530,7 @@ describe("bounded CTE layout", () => { ); }); - it("stops exact structural coverage at an embedded barrier", () => { + it("keeps structural coverage around an embedded query expression", () => { const text = "WITH a AS (SELECT {value}), b AS (SELECT 2) " + "SELECT * FROM b"; @@ -539,17 +540,19 @@ describe("bounded CTE layout", () => { ]); expect(layout.status).toBe("partial"); expect(layout.exactThrough).toBe(from); - expect(layout.declarations).toEqual([]); - expect(layout.draftDeclarations).toHaveLength(1); - expect(layout.draftDeclarations[0]).toMatchObject({ - bodyRange: { from: text.indexOf("SELECT"), to: from }, - name: { quoted: false, value: "a" }, - }); + expect(layout.declarations.map((item) => item.name.value)).toEqual([ + "a", + "b", + ]); + expect(layout.draftDeclarations).toEqual([]); expect(layout.issues).toContain("opaque-template-context"); expect( visibleSqlCtesAt(layout, text.lastIndexOf("FROM b") + 5), ).toMatchObject({ - ctes: [], + ctes: [ + { name: { value: "a" } }, + { name: { value: "b" } }, + ], quality: "recovered", shadowing: { coverage: "unknown" }, }); @@ -559,6 +562,28 @@ describe("bounded CTE layout", () => { shadowing: { coverage: "unknown" }, }); + const repeatedBarrierText = + "WITH a AS (SELECT {first} AS x, {second} AS y) SELECT 1"; + const firstBarrier = repeatedBarrierText.indexOf("{first}"); + const secondBarrier = repeatedBarrierText.indexOf("{second}"); + expect( + analyze(repeatedBarrierText, postgres, [ + { + from: firstBarrier, + language: "python", + to: firstBarrier + "{first}".length, + }, + { + from: secondBarrier, + language: "python", + to: secondBarrier + "{second}".length, + }, + ]), + ).toMatchObject({ + exactThrough: firstBarrier, + status: "partial", + }); + const laterBarrierText = "WITH a AS (SELECT 1), b AS (SELECT * FROM target {value}) " + "SELECT * FROM b"; @@ -857,6 +882,21 @@ describe("bounded CTE layout", () => { resource: "cte-declaration", }); + const columns = Array.from( + { length: MAX_QUERY_OUTPUT_COLUMNS + 1 }, + (_, index) => `column_${index}`, + ).join(", "); + const columnLayout = analyze( + `WITH c(${columns}) AS (SELECT 1) SELECT 1`, + ); + expect(columnLayout).toMatchObject({ + declarations: [{ + declaredColumns: { length: MAX_QUERY_OUTPUT_COLUMNS }, + }], + resource: "cte-column", + status: "partial", + }); + const acceptedDeclarations = Array.from( { length: MAX_CTE_DECLARATIONS }, (_, index) => `c${index} AS (SELECT ${index})`, diff --git a/src/__tests__/local-relation-site.test.ts b/src/__tests__/local-relation-site.test.ts index 1a1c56f..a0e21c9 100644 --- a/src/__tests__/local-relation-site.test.ts +++ b/src/__tests__/local-relation-site.test.ts @@ -2,6 +2,7 @@ import { describe, expect, it } from "vitest"; import { analyzeSqlLocalColumnSite, analyzeSqlLocalRelationSite, + applySqlQueryOutputAliases, prepareSqlLocalRelationStatement, type SqlLocalRelationSiteResult, } from "../local-relation-site.js"; @@ -23,6 +24,26 @@ import { type SqlStatementSlot, } from "../statement-index.js"; +it("retains bounded alias evidence when an underlying output is unavailable", () => { + expect( + applySqlQueryOutputAliases( + { reason: "unsupported-query", status: "unavailable" }, + { + columns: [{ + definition: { from: 0, to: 4 }, + identifier: { quoted: false, value: "name" }, + insertText: "name", + }], + coverage: "complete", + }, + ), + ).toMatchObject({ + columns: [{ identifier: { value: "name" } }], + coverage: "partial", + status: "ready", + }); +}); + const RUNTIMES = [ { dialect: POSTGRESQL_SQL_RELATION_DIALECT, diff --git a/src/__tests__/query-output.test.ts b/src/__tests__/query-output.test.ts new file mode 100644 index 0000000..819e33d --- /dev/null +++ b/src/__tests__/query-output.test.ts @@ -0,0 +1,215 @@ +import { describe, expect, it } from "vitest"; +import { inferSqlQueryOutput } from "../query-output.js"; +import { + BIGQUERY_SQL_RELATION_DIALECT, + POSTGRESQL_SQL_RELATION_DIALECT, +} from "../relation-dialect.js"; +import { + createIdentitySqlSource, + createMaskedSqlSource, +} from "../source.js"; + +function infer( + text: string, + dialect = POSTGRESQL_SQL_RELATION_DIALECT, +) { + const source = createIdentitySqlSource(text); + return inferSqlQueryOutput( + source, + { from: 0, to: text.length }, + dialect, + ); +} + +describe("query output inference", () => { + it("infers simple and explicitly aliased projection names", () => { + expect( + infer("SELECT users.id, upper(name) AS display_name FROM users"), + ).toMatchObject({ + columns: [ + { + identifier: { quoted: false, value: "id" }, + insertText: "id", + }, + { + identifier: { quoted: false, value: "display_name" }, + insertText: "display_name", + }, + ], + coverage: "complete", + status: "ready", + }); + }); + + it("keeps known names while marking unprovable expressions partial", () => { + expect(infer("SELECT id, upper(name), * FROM users")).toMatchObject({ + columns: [{ + identifier: { quoted: false, value: "id" }, + }], + coverage: "partial", + status: "ready", + }); + }); + + it.each([ + "SELECT DISTINCT id", + "SELECT ALL id", + ])("ignores a projection quantifier in %s", (text) => { + expect(infer(text)).toMatchObject({ + columns: [{ identifier: { value: "id" } }], + coverage: "complete", + }); + }); + + it.each([ + "SELECT , id", + "SELECT (id)", + "SELECT id.", + "SELECT {opaque}", + ])("keeps unsupported projection shape partial in %s", (text) => { + expect(infer(text)).toMatchObject({ + coverage: "partial", + status: "ready", + }); + }); + + it("uses the first set-operation arm for output names", () => { + expect( + infer( + "SELECT id AS first_name FROM users " + + "UNION ALL SELECT id AS later_name FROM archived", + ), + ).toMatchObject({ + columns: [{ + identifier: { quoted: false, value: "first_name" }, + }], + coverage: "complete", + status: "ready", + }); + }); + + it("uses a parenthesized first set-operation arm for output names", () => { + expect( + infer( + "(SELECT id AS first_name FROM users) " + + "UNION ALL SELECT id AS later_name FROM archived", + ), + ).toMatchObject({ + columns: [{ + identifier: { quoted: false, value: "first_name" }, + }], + coverage: "complete", + status: "ready", + }); + }); + + it("uses the first projection of a nested first set arm", () => { + expect( + infer( + "(SELECT a AS first_name UNION SELECT b AS second_name) " + + "UNION SELECT c AS later_name", + ), + ).toMatchObject({ + columns: [{ + identifier: { quoted: false, value: "first_name" }, + }], + coverage: "complete", + status: "ready", + }); + }); + + it("selects the outer projection after a nested WITH body", () => { + expect( + infer( + "WITH inner_cte AS (SELECT hidden FROM source) " + + "SELECT visible FROM inner_cte", + ), + ).toMatchObject({ + columns: [{ + identifier: { quoted: false, value: "visible" }, + }], + coverage: "complete", + status: "ready", + }); + }); + + it("preserves BigQuery quoted insertion spelling", () => { + expect( + infer("SELECT value AS `Display Name` FROM source", BIGQUERY_SQL_RELATION_DIALECT), + ).toMatchObject({ + columns: [{ + identifier: { quoted: true, value: "Display Name" }, + insertText: "`Display Name`", + }], + coverage: "complete", + status: "ready", + }); + }); + + it("keeps an aliased embedded projection explicitly partial", () => { + const text = "SELECT {opaque} AS metric"; + const from = text.indexOf("{opaque}"); + const source = createMaskedSqlSource(text, [{ + from, + language: "template", + to: from + "{opaque}".length, + }]); + expect(inferSqlQueryOutput( + source, + { from: 0, to: text.length }, + POSTGRESQL_SQL_RELATION_DIALECT, + )).toMatchObject({ + columns: [{ identifier: { value: "metric" } }], + coverage: "partial", + status: "ready", + }); + }); + + it("fails bounded and non-query ranges closed", () => { + expect(infer("VALUES (1)")).toEqual({ + reason: "unsupported-query", + status: "unavailable", + }); + const text = `SELECT ${"x".repeat(65_537)}`; + expect(infer(text)).toEqual({ + reason: "resource-limit", + status: "unavailable", + }); + }); + + it("bounds output count and lexical work", () => { + const wide = `SELECT ${Array.from( + { length: 257 }, + (_, index) => `column_${index}`, + ).join(", ")}`; + expect(infer(wide)).toMatchObject({ + columns: { length: 256 }, + coverage: "partial", + status: "ready", + }); + const lexical = `SELECT ${"a+".repeat(9_000)}a`; + expect(infer(lexical)).toEqual({ + reason: "resource-limit", + status: "unavailable", + }); + }); + + it.each([ + { from: -1, to: 1 }, + { from: 1, to: 1 }, + { from: 0.5, to: 1 }, + { from: 0, to: 100 }, + ])("rejects invalid range $from..$to", (range) => { + const source = createIdentitySqlSource("SELECT id"); + expect( + inferSqlQueryOutput( + source, + range, + POSTGRESQL_SQL_RELATION_DIALECT, + ), + ).toEqual({ + reason: "resource-limit", + status: "unavailable", + }); + }); +}); diff --git a/src/__tests__/query-site.test.ts b/src/__tests__/query-site.test.ts index 9f58b02..ea9d326 100644 --- a/src/__tests__/query-site.test.ts +++ b/src/__tests__/query-site.test.ts @@ -979,7 +979,6 @@ describe("fail-closed query-site behavior", () => { [", SELECT * FROM |", "inactive"], ["(SELECT 1) SELECT * FROM |", "inactive"], ["((SELECT 1)) SELECT * FROM |", "inactive"], - ["DELETE FROM |", "inactive"], ["COPY x FROM |", "inactive"], ["SELECT * FROM users alias |", "inactive"], ["SELECT * FROM|", "inactive"], @@ -1060,6 +1059,9 @@ describe("fail-closed query-site behavior", () => { ["SELECT * FROM a JOIN b USING(id other) JOIN |", "unavailable"], ["SELECT * FROM a JOIN b USING(id) other JOIN |", "unavailable"], ['SELECT * FROM a JOIN b USING(id) "other" JOIN |', "unavailable"], + ["SELECT * FROM a JOIN b USING(id) + JOIN |", "unavailable"], + ["SELECT * FROM a JOIN b USING(id) LEFT(1) JOIN |", "unavailable"], + ["SELECT * FROM a JOIN b USING(id) LEFT, |", "unavailable"], ["SELECT * FROM a JOIN b USING 'id' JOIN |", "unavailable"], ["SELECT * FROM a JOIN b USING FETCH JOIN |", "unavailable"], ["SELECT * FROM a JOIN b USING GROUP JOIN |", "unavailable"], @@ -1135,7 +1137,6 @@ describe("fail-closed query-site behavior", () => { ["SELECT * FROM 'not a relation' JOIN |", "unavailable"], ["SELECT * FROM , |", "unavailable"], ["SELECT * FROM schema..|", "unavailable"], - ["SELECT * FROM a UNION SELECT * FROM |", "unavailable"], ["SELECT * FROM a QUALIFY x JOIN |", "unavailable"], ] as const)("does not invent a site for %s", (marked, status) => { expect(recognize(marked).status).toBe(status); @@ -1565,8 +1566,6 @@ describe("fail-closed query-site behavior", () => { "SELECT * FROM users LEFT RIGHT JOIN |", "SELECT * FROM users LEFT potato JOIN |", "SELECT * FROM users a b JOIN |", - "SELECT * FROM users EXCEPT SELECT * FROM |", - "SELECT * FROM users INTERSECT SELECT * FROM |", "SELECT * FROM users WINDOW value JOIN |", "SELECT * FROM users alias . JOIN |", "SELECT * FROM users alias + JOIN |", @@ -1575,6 +1574,122 @@ describe("fail-closed query-site behavior", () => { expect(recognize(marked).status).toBe("unavailable"); }); + it.each([ + "SELECT * FROM users UNION SELECT * FROM |", + "SELECT * FROM users UNION ALL SELECT * FROM |", + "SELECT * FROM users UNION DISTINCT SELECT * FROM |", + "SELECT * FROM users INTERSECT SELECT * FROM |", + "SELECT * FROM users EXCEPT SELECT * FROM |", + ])("recognizes a relation site in each set-operation arm for %s", (marked) => { + expect(recognize(marked)).toMatchObject({ + anchor: "from", + prefix: { quoted: false, value: "" }, + qualifier: [], + status: "ready", + }); + }); + + it.each([ + ["INSERT INTO |", "from"], + ["UPDATE app.us| SET name = 'x'", "from"], + ["DELETE FROM | WHERE true", "from"], + ["MERGE INTO target USING | ON true", "join"], + ["/* leading */ UPDATE |", "from"], + ] as const)("recognizes the DML relation site in %s", (marked, anchor) => { + expect(recognize(marked)).toMatchObject({ + anchor, + status: "ready", + }); + }); + + it.each([ + "INSERT target|", + "MERGE target| USING source ON true", + ])("recognizes optional BigQuery INTO in %s", (marked) => { + expect(recognize(marked, { + dialect: bigQueryDialect, + })).toMatchObject({ + anchor: "from", + status: "ready", + }); + }); + + it.each([ + "INSERT INTO target SELECT * FROM |", + "UPDATE target SET value = (SELECT value FROM |)", + "DELETE FROM target WHERE EXISTS (SELECT 1 FROM |)", + "MERGE INTO target USING source ON EXISTS (SELECT 1 FROM |)", + ])("continues into nested SELECT relation sites for %s", (marked) => { + expect(recognize(marked).status).toBe("ready"); + }); + + it.each([ + "INSERT |", + "INSERT target |", + "DELETE |", + "DELETE target |", + "MERGE USING |", + "MERGE target |", + "UPDATE (SELECT 1) |", + "UPDATE 'target' |", + ])("fails malformed DML relation transitions closed in %s", (marked) => { + expect(recognize(marked).status).not.toBe("ready"); + }); + + it("keeps completed DML targets outside relation completion", () => { + expect(recognize("UPDATE target SET value = 1|")).toEqual({ + reason: "not-select-query", + status: "inactive", + }); + expect(recognize("MERGE INTO target USING source ON true|")).toEqual({ + reason: "not-select-query", + status: "inactive", + }); + }); + + it("classifies DML cursor barriers without inventing targets", () => { + expect(recognize("UP|DATE users").status).toBe("inactive"); + expect(recognize("UPDATE /* tar|get */").status).toBe("inactive"); + expect(recognize("UPDATE 'tar|get'").status).toBe("inactive"); + const marked = "UPDATE {py|thon}"; + const { position, text } = markedSource(marked); + const from = text.indexOf("{python}"); + expect(recognize(marked, { + regions: [{ + from, + language: "python", + to: from + "{python}".length, + }], + })).toEqual({ + reason: "cursor-in-embedded-region", + status: "inactive", + }); + expect(position).toBeGreaterThan(from); + }); + + it("fails closed when an embedded region precedes a DML target", () => { + const marked = "UPDATE {python} |"; + const { text } = markedSource(marked); + const from = text.indexOf("{python}"); + expect(recognize(marked, { + regions: [{ + from, + language: "python", + to: from + "{python}".length, + }], + })).toEqual({ + reason: "ambiguous-query-site", + status: "unavailable", + }); + }); + + it("does not interpret non-USING MERGE suffixes as source relations", () => { + expect(recognize("MERGE INTO target ON |")).toEqual({ + reason: "not-relation-position", + status: "inactive", + }); + }); + it("suppresses the three-word IS NOT DISTINCT FROM expression", () => { expect(recognize("SELECT x IS NOT DISTINCT FROM |").status).toBe( "inactive", diff --git a/src/__tests__/semantic-completion-corpus.test.ts b/src/__tests__/semantic-completion-corpus.test.ts new file mode 100644 index 0000000..2a54b7b --- /dev/null +++ b/src/__tests__/semantic-completion-corpus.test.ts @@ -0,0 +1,81 @@ +import { describe, expect, it } from "vitest"; +import { + bigQueryDialect, + createSqlLanguageService, + duckdbDialect, + postgresDialect, + type SqlDialect, +} from "../index.js"; +import { bigQuerySemanticCompletion } from "../../test/corpora/semantic-completion/bigquery.js"; +import { duckDbSemanticCompletion } from "../../test/corpora/semantic-completion/duckdb.js"; +import { postgresqlSemanticCompletion } from "../../test/corpora/semantic-completion/postgresql.js"; +import type { + SemanticCompletionCase, + SemanticCompletionCategory, +} from "../../test/corpora/semantic-completion/types.js"; + +const categories: readonly SemanticCompletionCategory[] = [ + "valid", + "invalid", + "incomplete", + "templated", + "multi-statement", +]; + +const corpora: readonly { + readonly cases: readonly SemanticCompletionCase[]; + readonly dialect: SqlDialect; +}[] = [ + { cases: postgresqlSemanticCompletion, dialect: postgresDialect() }, + { cases: bigQuerySemanticCompletion, dialect: bigQueryDialect() }, + { cases: duckDbSemanticCompletion, dialect: duckdbDialect() }, +]; + +describe.each(corpora)("$dialect.id semantic completion corpus", ({ + cases, + dialect, +}) => { + it("owns every required corpus category", () => { + expect(new Set(cases.map((item) => item.category))).toEqual( + new Set(categories), + ); + }); + + it.each(cases)("$category: $sql", async (fixture) => { + const position = fixture.sql.indexOf("|"); + expect(position).toBeGreaterThanOrEqual(0); + expect(fixture.sql.indexOf("|", position + 1)).toBe(-1); + const text = + fixture.sql.slice(0, position) + fixture.sql.slice(position + 1); + const templateFrom = fixture.template + ? text.indexOf(fixture.template) + : -1; + const service = createSqlLanguageService({ + dialects: [dialect], + }); + const session = service.openDocument({ + context: { dialect: dialect.id }, + embeddedRegions: templateFrom < 0 || !fixture.template + ? [] + : [{ + from: templateFrom, + language: "host", + to: templateFrom + fixture.template.length, + }], + text, + }); + + const result = await session.complete({ + position, + trigger: { kind: "invoked" }, + }); + expect(result.status).toBe("ready"); + if (result.status === "ready") { + expect(result.value.items.map((item) => item.label)).toEqual( + fixture.expectedLabels, + ); + expect(result.value.isIncomplete).toBe(fixture.expectedIncomplete); + } + service.dispose(); + }); +}); diff --git a/src/__tests__/semantic-completion-performance.test.ts b/src/__tests__/semantic-completion-performance.test.ts new file mode 100644 index 0000000..87a9d9c --- /dev/null +++ b/src/__tests__/semantic-completion-performance.test.ts @@ -0,0 +1,99 @@ +import { afterEach, describe, expect, it } from "vitest"; +import { + createSqlLanguageService, + duckdbDialect, + type SqlLanguageService, +} from "../index.js"; + +const services: SqlLanguageService<{ + readonly dialect: string; +}>[] = []; + +afterEach(() => { + for (const service of services.splice(0)) service.dispose(); +}); + +function percentile95(samples: readonly number[]): number { + const ordered = [...samples].sort((left, right) => left - right); + return ordered[Math.ceil(ordered.length * 0.95) - 1] ?? Infinity; +} + +describe("semantic completion performance gates", () => { + it("keeps warm 10 KiB CTE output completion below 50 ms p95", async () => { + const projections = Array.from( + { length: 200 }, + (_, index) => `source_${index} AS output_${index}`, + ).join(", "); + const prefix = + `WITH c AS (SELECT ${projections} FROM physical) ` + + "SELECT c.output_ FROM c "; + const padding = `/*${"x".repeat(10 * 1024 - prefix.length - 2)}*/`; + const text = prefix + padding; + const position = text.indexOf("output_ FROM") + "output_".length; + const service = createSqlLanguageService({ + dialects: [duckdbDialect()], + }); + services.push(service); + const session = service.openDocument({ + context: { dialect: "duckdb" }, + text, + }); + const request = { + position, + trigger: { kind: "invoked" as const }, + }; + await session.complete(request); + const samples: number[] = []; + for (let index = 0; index < 20; index += 1) { + const startedAt = performance.now(); + const result = await session.complete(request); + samples.push(performance.now() - startedAt); + expect(result.status).toBe("ready"); + if (result.status === "ready") { + expect(result.value.items).toHaveLength(200); + } + } + expect(percentile95(samples)).toBeLessThan(50); + }); + + it("keeps repeated CTE inference cached and aggregate output bounded", async () => { + const projections = Array.from( + { length: 64 }, + (_, index) => `source_${index} AS output_${index}`, + ).join(", "); + const relations = Array.from( + { length: 65 }, + (_, index) => `c AS c_${index}`, + ).join(", "); + const text = + `WITH c AS (SELECT ${projections} FROM physical) ` + + `SELECT FROM ${relations}`; + const position = + text.indexOf("SELECT FROM") + "SELECT ".length; + const service = createSqlLanguageService({ + dialects: [duckdbDialect()], + }); + services.push(service); + const session = service.openDocument({ + context: { dialect: "duckdb" }, + text, + }); + const request = { + position, + trigger: { kind: "invoked" as const }, + }; + await session.complete(request); + const samples: number[] = []; + for (let index = 0; index < 10; index += 1) { + const startedAt = performance.now(); + const result = await session.complete(request); + samples.push(performance.now() - startedAt); + expect(result.status).toBe("ready"); + if (result.status === "ready") { + expect(result.value.items).toHaveLength(4_096); + expect(result.value.isIncomplete).toBe(true); + } + } + expect(percentile95(samples)).toBeLessThan(50); + }); +}); diff --git a/src/__tests__/session.test.ts b/src/__tests__/session.test.ts index c8d3ead..11f69cf 100644 --- a/src/__tests__/session.test.ts +++ b/src/__tests__/session.test.ts @@ -263,6 +263,97 @@ describe("relation completion session integration", () => { service.dispose(); }); + it.each([ + "INSERT INTO us", + "UPDATE us SET active = true", + "DELETE FROM us WHERE active = false", + "MERGE INTO target USING us ON true", + "WITH c AS (SELECT 1) UPDATE us SET active = true", + "WITH c AS (SELECT 1) DELETE FROM us WHERE active = false", + "WITH c AS (SELECT 1) INSERT INTO us", + "WITH c AS (SELECT 1) MERGE INTO us", + ])("completes a DML relation target in %s", async (text) => { + const service = createSqlLanguageService({ + catalog: catalogProvider(async () => ({ + coverage: { kind: "complete" }, + epoch: { generation: 0, token: "dml" }, + relations: [{ + canonicalPath: [{ + quoted: false, + role: "relation", + value: "users", + }], + completionPathStart: 0, + entityId: "users", + matchQuality: "exact", + relationKind: "table", + }], + status: "ready", + })), + dialects: [duckdb], + }); + const position = text.lastIndexOf("us") + 2; + const session = service.openDocument({ + context: { + catalog: { scope: "connection:dml" }, + dialect: "duckdb", + engine: "local", + }, + text, + }); + + await expect(session.complete({ + position, + trigger: { kind: "invoked" }, + })).resolves.toMatchObject({ + status: "ready", + value: { + items: [{ edit: { insert: "users" }, label: "users" }], + }, + }); + service.dispose(); + }); + + it("completes relations in a later set-operation arm", async () => { + const text = "SELECT * FROM current UNION ALL SELECT * FROM us"; + const service = createSqlLanguageService({ + catalog: catalogProvider(async () => ({ + coverage: { kind: "complete" }, + epoch: { generation: 0, token: "set-arm" }, + relations: [{ + canonicalPath: [{ + quoted: false, + role: "relation", + value: "users", + }], + completionPathStart: 0, + entityId: "users", + matchQuality: "exact", + relationKind: "table", + }], + status: "ready", + })), + dialects: [duckdb], + }); + const session = service.openDocument({ + context: { + catalog: { scope: "connection:set-arm" }, + dialect: "duckdb", + engine: "local", + }, + text, + }); + + await expect(session.complete({ + position: text.length, + trigger: { kind: "invoked" }, + })).resolves.toMatchObject({ + status: "ready", + value: { items: [{ label: "users" }] }, + }); + service.dispose(); + }); + it("keeps completion edits in absolute UTF-16 document coordinates", async () => { const service = createSqlLanguageService({ catalog: catalogProvider(async () => ({ @@ -1263,6 +1354,473 @@ describe("column completion session integration", () => { }); } + it("completes proven CTE output columns without a physical provider", async () => { + const service = createSqlLanguageService({ + dialects: [duckdb], + }); + const text = + "WITH recent AS (SELECT id, name AS label FROM users) " + + "SELECT recent.la FROM recent"; + const session = service.openDocument({ + context: { dialect: "duckdb", engine: "local" }, + text, + }); + + await expect(session.complete({ + position: text.indexOf("la FROM") + 2, + trigger: { kind: "invoked" }, + })).resolves.toMatchObject({ + sources: [{ + coverage: "complete", + feature: "query-output", + outcome: "ready", + }], + status: "ready", + value: { + isIncomplete: false, + items: [{ + edit: { insert: "label" }, + kind: "column", + label: "label", + provenance: { kind: "query-output" }, + }], + }, + }); + service.dispose(); + }); + + it("uses an explicit CTE column list as the authoritative output shape", async () => { + const service = createSqlLanguageService({ + dialects: [postgres], + }); + const text = + 'WITH recent("User ID", label) AS (SELECT *, upper(name) FROM users) ' + + 'SELECT recent."U" FROM recent'; + const session = service.openDocument({ + context: { dialect: "postgresql", engine: "local" }, + text, + }); + + await expect(session.complete({ + position: text.indexOf('"U" FROM') + 2, + trigger: { kind: "invoked" }, + })).resolves.toMatchObject({ + sources: [{ + coverage: "complete", + feature: "query-output", + }], + status: "ready", + value: { + isIncomplete: false, + items: [{ + edit: { insert: '"User ID"' }, + label: "User ID", + provenance: { kind: "query-output" }, + }], + }, + }); + service.dispose(); + }); + + it("completes derived and set-operation output names from the first arm", async () => { + const service = createSqlLanguageService({ + dialects: [duckdb], + }); + const text = + "SELECT d.fi FROM (" + + "SELECT id AS first_name FROM users " + + "UNION ALL SELECT id AS later_name FROM archived" + + ") AS d"; + const session = service.openDocument({ + context: { dialect: "duckdb", engine: "local" }, + text, + }); + + await expect(session.complete({ + position: text.indexOf("fi FROM") + 2, + trigger: { kind: "invoked" }, + })).resolves.toMatchObject({ + sources: [{ feature: "query-output" }], + status: "ready", + value: { + items: [{ + edit: { insert: "first_name" }, + label: "first_name", + provenance: { kind: "query-output" }, + }], + }, + }); + service.dispose(); + }); + + it("uses derived and CTE relation-alias column lists", async () => { + const service = createSqlLanguageService({ + dialects: [postgres], + }); + const derivedText = + "SELECT d. FROM (SELECT original, other) AS d(renamed, second)"; + const derived = service.openDocument({ + context: { dialect: "postgresql", engine: "local" }, + text: derivedText, + }); + await expect(derived.complete({ + position: derivedText.indexOf(" FROM"), + trigger: { kind: "invoked" }, + })).resolves.toMatchObject({ + status: "ready", + value: { + items: [ + { label: "renamed" }, + { label: "second" }, + ], + }, + }); + derived.dispose(); + + const cteText = + "WITH c AS (SELECT original, other) " + + "SELECT x. FROM c AS x(renamed, second)"; + const cte = service.openDocument({ + context: { dialect: "postgresql", engine: "local" }, + text: cteText, + }); + await expect(cte.complete({ + position: cteText.indexOf(" FROM c"), + trigger: { kind: "invoked" }, + })).resolves.toMatchObject({ + status: "ready", + value: { + items: [ + { label: "renamed" }, + { label: "second" }, + ], + }, + }); + cte.dispose(); + service.dispose(); + }); + + it("keeps partial relation-alias evidence bounded and explicit", async () => { + const service = createSqlLanguageService({ + dialects: [postgres], + }); + const text = + "SELECT d. FROM (SELECT original, *) AS d(renamed, extra,)"; + const session = service.openDocument({ + context: { dialect: "postgresql", engine: "local" }, + text, + }); + await expect(session.complete({ + position: text.indexOf(" FROM"), + trigger: { kind: "invoked" }, + })).resolves.toMatchObject({ + status: "ready", + value: { + isIncomplete: true, + issues: [{ reason: "query-binding-partial" }], + items: [ + { label: "extra" }, + { label: "renamed" }, + ], + }, + }); + service.dispose(); + }); + + it("reuses one inferred CTE output across repeated aliases", async () => { + const service = createSqlLanguageService({ + dialects: [duckdb], + }); + const text = + "WITH c AS (SELECT id) SELECT FROM c a, c b, c d"; + const session = service.openDocument({ + context: { dialect: "duckdb", engine: "local" }, + text, + }); + await expect(session.complete({ + position: text.indexOf(" FROM"), + trigger: { kind: "invoked" }, + })).resolves.toMatchObject({ + status: "ready", + value: { + items: [ + { detail: "a", label: "id" }, + { detail: "b", label: "id" }, + { detail: "d", label: "id" }, + ], + }, + }); + service.dispose(); + }); + + it("keeps provable local columns while reporting unknown projections partial", async () => { + const service = createSqlLanguageService({ + dialects: [duckdb], + }); + const text = + "WITH recent AS (SELECT id, *, upper(name) FROM users) " + + "SELECT recent.i FROM recent"; + const session = service.openDocument({ + context: { dialect: "duckdb", engine: "local" }, + text, + }); + + await expect(session.complete({ + position: text.indexOf("i FROM") + 1, + trigger: { kind: "invoked" }, + })).resolves.toMatchObject({ + sources: [{ + coverage: "partial", + feature: "query-output", + }], + status: "ready", + value: { + isIncomplete: true, + issues: [{ reason: "query-binding-partial" }], + items: [{ label: "id" }], + }, + }); + service.dispose(); + }); + + it("merges inferred and physical columns in one completion result", async () => { + const service = serviceWithColumns(async (request) => ({ + epoch: { generation: 1, token: "mixed" }, + relations: request.relations.map((relation) => ({ + columns: [{ + columnEntityId: "physical:invoice_id", + identifier: { quoted: false, value: "invoice_id" }, + insertText: "invoice_id", + ordinal: 0, + }], + coverage: "complete", + relationEntityId: "physical", + requestKey: relation.requestKey, + status: "ready", + })), + })); + const text = + "WITH c AS (SELECT id AS local_id) " + + "SELECT FROM c JOIN physical p ON true"; + const session = service.openDocument({ + context: { + catalog: { scope: "connection:mixed" }, + dialect: "duckdb", + engine: "local", + }, + text, + }); + + await expect(session.complete({ + position: text.indexOf(" FROM"), + trigger: { kind: "invoked" }, + })).resolves.toMatchObject({ + sources: [ + { feature: "query-output" }, + { feature: "column-catalog" }, + ], + status: "ready", + value: { + items: [ + { label: "local_id" }, + { label: "invoice_id" }, + ], + }, + }); + service.dispose(); + }); + + it("retains local output when physical column completion is unavailable", async () => { + const service = createSqlLanguageService({ + dialects: [duckdb], + }); + const text = + "WITH c AS (SELECT id AS local_id) " + + "SELECT l FROM c JOIN physical p ON true"; + const session = service.openDocument({ + context: { dialect: "duckdb", engine: "local" }, + text, + }); + + await expect(session.complete({ + position: text.indexOf("l FROM") + 1, + trigger: { kind: "invoked" }, + })).resolves.toMatchObject({ + sources: [{ feature: "query-output" }], + status: "ready", + value: { + isIncomplete: true, + issues: [{ reason: "query-binding-partial" }], + items: [{ label: "local_id" }], + }, + }); + service.dispose(); + }); + + it("intersects local query outputs for JOIN USING completion", async () => { + const service = createSqlLanguageService({ + dialects: [duckdb], + }); + const text = + "WITH a AS (SELECT 1 AS id, 2 AS only_a), " + + "b AS (SELECT 1 AS id, 3 AS only_b) " + + "SELECT * FROM a JOIN b USING ()"; + const session = service.openDocument({ + context: { dialect: "duckdb", engine: "local" }, + text, + }); + + await expect(session.complete({ + position: text.lastIndexOf(")"), + trigger: { kind: "invoked" }, + })).resolves.toMatchObject({ + status: "ready", + value: { items: [{ label: "id" }] }, + }); + service.dispose(); + }); + + it("intersects local and physical outputs for JOIN USING completion", async () => { + const service = serviceWithColumns(async (request) => ({ + epoch: { generation: 1, token: "mixed-using" }, + relations: request.relations.map((relation) => ({ + columns: [ + { + columnEntityId: "physical:id", + identifier: { quoted: false, value: "id" }, + insertText: "id", + ordinal: 0, + }, + { + columnEntityId: "physical:only_physical", + identifier: { quoted: false, value: "only_physical" }, + insertText: "only_physical", + ordinal: 1, + }, + ], + coverage: "complete", + relationEntityId: "physical", + requestKey: relation.requestKey, + status: "ready", + })), + })); + const text = + "WITH c AS (SELECT 1 AS id, 2 AS only_local) " + + "SELECT * FROM c JOIN physical p USING ()"; + const session = service.openDocument({ + context: { + catalog: { scope: "connection:mixed-using" }, + dialect: "duckdb", + engine: "local", + }, + text, + }); + + await expect(session.complete({ + position: text.lastIndexOf(")"), + trigger: { kind: "invoked" }, + })).resolves.toMatchObject({ + sources: [ + { feature: "query-output" }, + { feature: "column-catalog" }, + ], + status: "ready", + value: { items: [{ label: "id" }] }, + }); + service.dispose(); + }); + + it("uses the left physical spelling for mixed JOIN USING completion", async () => { + const service = serviceWithColumns(async (request) => ({ + epoch: { generation: 1, token: "physical-left-using" }, + relations: request.relations.map((relation) => ({ + columns: [ + { + columnEntityId: "physical:id", + identifier: { quoted: false, value: "id" }, + insertText: "id", + ordinal: 0, + }, + { + columnEntityId: "physical:only_physical", + identifier: { quoted: false, value: "only_physical" }, + insertText: "only_physical", + ordinal: 1, + }, + ], + coverage: "complete", + relationEntityId: "physical", + requestKey: relation.requestKey, + status: "ready", + })), + })); + const text = + "WITH c AS (SELECT 1 AS id, 2 AS only_local) " + + "SELECT * FROM physical p JOIN c USING ()"; + const session = service.openDocument({ + context: { + catalog: { scope: "connection:physical-left-using" }, + dialect: "duckdb", + engine: "local", + }, + text, + }); + + await expect(session.complete({ + position: text.lastIndexOf(")"), + trigger: { kind: "invoked" }, + })).resolves.toMatchObject({ + status: "ready", + value: { + items: [{ + edit: { insert: "id" }, + label: "id", + provenance: { kind: "column-catalog" }, + }], + }, + }); + service.dispose(); + }); + + it("returns local output while a physical column request exceeds its budget", async () => { + const service = createSqlLanguageService({ + catalog: relationCatalog, + columns: { + id: "columns", + loadColumns: () => new Promise(() => {}), + }, + completion: { catalogResponseBudgetMs: 0 }, + dialects: [duckdb], + }); + const text = + "WITH c AS (SELECT id AS local_id) " + + "SELECT l FROM c JOIN physical p ON true"; + const session = service.openDocument({ + context: { + catalog: { scope: "connection:mixed-timeout" }, + dialect: "duckdb", + engine: "local", + }, + text, + }); + + await expect(session.complete({ + position: text.indexOf("l FROM") + 1, + trigger: { kind: "invoked" }, + })).resolves.toMatchObject({ + sources: [ + { feature: "query-output" }, + { feature: "column-catalog", outcome: "loading" }, + ], + status: "ready", + value: { + isIncomplete: true, + items: [{ label: "local_id" }], + }, + }); + service.dispose(); + }); + it("batches and completes a qualified alias through the session", async () => { const requests: Parameters< SqlColumnCatalogProvider["loadColumns"] @@ -1532,7 +2090,14 @@ describe("column completion session integration", () => { position, trigger: { kind: "invoked" }, })).resolves.toMatchObject({ - status: "unavailable", + sources: [{ feature: "query-output" }], + status: "ready", + value: { + items: [{ + edit: { insert: "local_id" }, + label: "local_id", + }], + }, }); expect(calls).toBe(0); service.dispose(); diff --git a/src/codemirror/browser_tests/sql-editor-completion-gate.test.ts b/src/codemirror/browser_tests/sql-editor-completion-gate.test.ts index abdfaa8..243ca0c 100644 --- a/src/codemirror/browser_tests/sql-editor-completion-gate.test.ts +++ b/src/codemirror/browser_tests/sql-editor-completion-gate.test.ts @@ -210,6 +210,43 @@ test("standard editor can accept the first completion immediately", async () => expect(view.state.doc.toString()).toBe("SELECT * FROM users"); }); +test("standard editor completes and applies a CTE output column", async () => { + const parent = document.createElement("div"); + document.body.append(parent); + onTestFinished(() => parent.remove()); + const service = createSqlLanguageService({ + dialects: [duckdbDialect()], + }); + onTestFinished(() => service.dispose()); + const support = sqlEditor({ + autocomplete: { selectOnOpen: true }, + initialContext: { dialect: "duckdb" }, + service, + }); + const documentText = + "WITH c AS (SELECT id AS user_id) SELECT c. FROM c"; + const position = documentText.indexOf("c. FROM") + 2; + const view = new EditorView({ + doc: documentText, + extensions: support.extension, + parent, + selection: { anchor: position }, + }); + onTestFinished(() => view.destroy()); + + view.focus(); + expect(startCompletion(view)).toBe(true); + await expect.poll(() => + currentCompletions(view.state).map((item) => item.label) + ).toEqual(["user_id"]); + await expect.poll(() => selectedCompletionIndex(view.state)).toBe(0); + await new Promise((resolve) => setTimeout(resolve, 100)); + expect(acceptCompletion(view)).toBe(true); + expect(view.state.doc.toString()).toBe( + "WITH c AS (SELECT id AS user_id) SELECT c.user_id FROM c", + ); +}); + test("standard editor preserves SQL language completions", async () => { const parent = document.createElement("div"); document.body.append(parent); diff --git a/src/column-completion.ts b/src/column-completion.ts index 8ea49e9..f0b433f 100644 --- a/src/column-completion.ts +++ b/src/column-completion.ts @@ -3,6 +3,7 @@ import type { } from "./column-catalog-batch-coordinator.js"; import { MAX_COLUMN_BATCH_RELATIONS, + MAX_COLUMNS_PER_BATCH, } from "./column-catalog-boundary.js"; import type { SqlColumnCatalogRelationReference, @@ -18,6 +19,8 @@ import type { SqlCompletionIssue, SqlCompletionItem, SqlCompletionList, + SqlCompletionProviderReport, + SqlQueryOutputProviderReport, } from "./relation-completion-types.js"; import type { SqlRelationDialectRuntime } from "./relation-dialect.js"; import type { @@ -40,10 +43,15 @@ export interface SqlPreparedColumnCatalogRelations { } export interface SqlColumnCompletionComposition { - readonly sources: readonly SqlColumnCatalogProviderReport[]; + readonly sources: readonly SqlCompletionProviderReport[]; readonly value: SqlCompletionList; } +const completionIdentifiers = new WeakMap< + SqlCompletionItem, + SqlIdentifierComponent +>(); + function sameIdentifier( dialect: SqlRelationDialectRuntime, left: SqlIdentifierComponent, @@ -94,10 +102,15 @@ export function prepareSqlColumnCatalogRelations( SqlColumnQueryRelation >(); let coverage: "complete" | "partial" = "complete"; - for (let index = 0; index < site.relations.length; index += 1) { + const firstRelation = + site.context === "using" + ? Math.max(0, site.relations.length - 2) + : 0; + for (let index = firstRelation; index < site.relations.length; index += 1) { const relation = site.relations[index]; if ( relation === undefined || + relation.local !== undefined || !relationMatchesQualifier( dialect, relation, @@ -129,6 +142,135 @@ function relationLabel(relation: SqlColumnQueryRelation): string { relation.path.map((component) => component.value).join("."); } +export function composeSqlLocalQueryOutputCompletion( + site: ReadyColumnSite, + dialect: SqlRelationDialectRuntime, +): SqlColumnCompletionComposition | null { + const items: SqlCompletionItem[] = []; + const seen = new Set(); + let found = false; + let partial = site.coverage === "partial"; + const firstRelation = site.context === "using" + ? Math.max(0, site.relations.length - 2) + : 0; + relations: for ( + let index = firstRelation; + index < site.relations.length; + index += 1 + ) { + const relation = site.relations[index]; + if ( + !relation?.local || + !relationMatchesQualifier(dialect, relation, site.qualifier) + ) { + continue; + } + found = true; + const output = relation.local.output; + if (!output || output.status !== "ready") { + partial = true; + continue; + } + if (output.coverage === "partial") partial = true; + for (const column of output.columns) { + if ( + dialect.completion.cteIdentifierMatchesPrefix( + column.identifier, + site.prefix, + ) !== "match" + ) { + continue; + } + const identity = [ + relation.range.from, + relation.range.to, + column.identifier.quoted, + column.identifier.value, + column.definition.from, + column.definition.to, + ].join("\u0000"); + if (seen.has(identity)) continue; + seen.add(identity); + if (items.length === MAX_COLUMNS_PER_BATCH) { + partial = true; + break relations; + } + const item = Object.freeze({ + detail: relationLabel(relation), + edit: Object.freeze({ + from: site.replacementRange.from, + insert: column.insertText, + to: site.replacementRange.to, + }), + kind: "column", + label: column.identifier.value, + provenance: Object.freeze({ + definition: column.definition, + kind: "query-output", + relation: relation.range, + }), + relationRequestKey: `local:${index}`, + }); + completionIdentifiers.set(item, column.identifier); + items.push(item); + } + } + if (!found) return null; + const source: SqlQueryOutputProviderReport = Object.freeze({ + coverage: partial ? "partial" : "complete", + feature: "query-output", + outcome: "ready", + }); + return Object.freeze({ + sources: Object.freeze([source]), + value: list( + Object.freeze(items.sort(compareItems)), + partial ? Object.freeze([issue("query-binding-partial")]) : Object.freeze([]), + ), + }); +} + +export function filterSqlUsingCompletionList( + value: SqlCompletionList, + site: ReadyColumnSite, + dialect: SqlRelationDialectRuntime, +): SqlCompletionList { + if (site.context !== "using" || site.relations.length < 2) return value; + const leftIndex = site.relations.length - 2; + const rightIndex = site.relations.length - 1; + const left = site.relations[leftIndex]; + const right = site.relations[rightIndex]; + if (!left || !right) return value; + const leftKey = left.local ? `local:${leftIndex}` : `binding:${leftIndex}`; + const rightKey = right.local + ? `local:${rightIndex}` + : `binding:${rightIndex}`; + const rightIdentifiers = value.items.flatMap((item) => { + if ( + item.kind !== "column" || + item.relationRequestKey !== rightKey + ) { + return []; + } + const identifier = completionIdentifiers.get(item); + return identifier ? [identifier] : []; + }); + const items = value.items.filter((item) => { + if ( + item.kind !== "column" || + item.relationRequestKey !== leftKey + ) { + return false; + } + const candidate = completionIdentifiers.get(item); + return candidate !== undefined && + rightIdentifiers.some((identifier) => + sameIdentifier(dialect, candidate, identifier) + ); + }); + return list(Object.freeze(items), value.issues); +} + function issue( reason: Exclude, ): SqlCompletionIssue { @@ -239,6 +381,7 @@ export function composeSqlColumnCompletion( let hasLoading = false; let hasFailure = false; const seen = new Set(); + const usingIdentifiers = new Map(); const failures: SqlColumnCatalogFailure[] = []; for (const result of input.outcome.relations) { issues.push(...resultIssues(result)); @@ -263,10 +406,26 @@ export function composeSqlColumnCompletion( issues.push(issue("column-catalog-malformed")); continue; } - for (const column of result.columns) { + if ( + relation.columnAliases?.coverage === "partial" || + (relation.columnAliases?.columns.length ?? 0) > + result.columns.length + ) { + hasPartial = true; + issues.push(issue("query-binding-partial")); + } + for ( + let columnIndex = 0; + columnIndex < result.columns.length; + columnIndex += 1 + ) { + const column = result.columns[columnIndex]!; + const alias = relation.columnAliases?.columns[columnIndex]; + const completionIdentifier = alias?.identifier ?? column.identifier; + const insertText = alias?.insertText ?? column.insertText; if ( input.dialect.completion.cteIdentifierMatchesPrefix( - column.identifier, + completionIdentifier, input.site.prefix, ) !== "match" ) { @@ -277,13 +436,16 @@ export function composeSqlColumnCompletion( column.provenance.scope, column.provenance.relationEntityId, column.provenance.columnEntityId, - column.insertText, + insertText, ].join("\u0000"); + const identifiers = usingIdentifiers.get(result.requestKey) ?? []; + identifiers.push(completionIdentifier); + usingIdentifiers.set(result.requestKey, identifiers); if (seen.has(identity)) continue; seen.add(identity); const relationName = relationLabel(relation); const metadata = column.detail ?? column.dataType; - items.push(Object.freeze({ + const item = Object.freeze({ ...(column.dataType === undefined ? {} : { dataType: column.dataType }), @@ -292,11 +454,11 @@ export function composeSqlColumnCompletion( : `${metadata} — ${relationName}`, edit: Object.freeze({ from: input.site.replacementRange.from, - insert: column.insertText, + insert: insertText, to: input.site.replacementRange.to, }), kind: "column", - label: column.identifier.value, + label: completionIdentifier.value, provenance: Object.freeze({ columnEntityId: column.provenance.columnEntityId, epoch: column.provenance.epoch, @@ -306,7 +468,33 @@ export function composeSqlColumnCompletion( scope: column.provenance.scope, }), relationRequestKey: result.requestKey, - })); + }); + completionIdentifiers.set(item, completionIdentifier); + items.push(item); + } + } + if ( + input.site.context === "using" && + input.prepared.references.length === 2 + ) { + const leftKey = input.prepared.references[0]?.requestKey; + const rightKey = input.prepared.references[1]?.requestKey; + if (leftKey && rightKey) { + const rightIdentifiers = usingIdentifiers.get(rightKey) ?? []; + const shared = items.filter((item) => { + if ( + item.kind !== "column" || + item.relationRequestKey !== leftKey + ) { + return false; + } + const left = completionIdentifiers.get(item); + return left !== undefined && + rightIdentifiers.some((right) => + sameIdentifier(input.dialect, left, right) + ); + }); + items.splice(0, items.length, ...shared); } } const deduplicatedIssues = Array.from( diff --git a/src/column-query-site.ts b/src/column-query-site.ts index dc8e8a6..e16de49 100644 --- a/src/column-query-site.ts +++ b/src/column-query-site.ts @@ -17,18 +17,118 @@ import type { SqlIdentifierPath, SqlTextRange, } from "./types.js"; +import { + MAX_QUERY_OUTPUT_COLUMNS, + type SqlQueryOutput, + type SqlQueryOutputColumn, +} from "./query-output.js"; export const MAX_COLUMN_QUERY_RELATIONS = 256; export interface SqlColumnQueryRelation { readonly alias: SqlIdentifierComponent | null; + readonly columnAliases?: { + readonly columns: readonly SqlQueryOutputColumn[]; + readonly coverage: "complete" | "partial"; + }; + readonly local?: { + readonly kind: "cte" | "derived"; + readonly output?: SqlQueryOutput; + readonly queryRange: SqlTextRange; + }; readonly path: SqlIdentifierPath; readonly range: SqlTextRange; } +function parseColumnAliases( + source: SqlSourceSnapshot, + dialect: SqlRelationDialectRuntime, + tokens: readonly Token[], + start: number, + depth: number, +): { + readonly aliases: + | SqlColumnQueryRelation["columnAliases"] + | undefined; + readonly next: number; +} { + const open = tokens[start]; + if (!punctuation(source, open, "(") || open?.depth !== depth) { + return { aliases: undefined, next: start }; + } + const columns: SqlQueryOutputColumn[] = []; + let expectIdentifier = true; + let coverage: "complete" | "partial" = "complete"; + let index = start + 1; + for (; index < tokens.length; index += 1) { + const token = tokens[index]!; + if ( + token.kind === "comment" || + token.kind === "line-comment" + ) { + continue; + } + if (punctuation(source, token, ")") && token.depth === depth) { + if (expectIdentifier) coverage = "partial"; + return { + aliases: Object.freeze({ + columns: Object.freeze(columns), + coverage, + }), + next: index + 1, + }; + } + if (token.depth !== depth + 1) { + coverage = "partial"; + continue; + } + if (expectIdentifier && isIdentifier(token)) { + const path = decodePath( + source, + dialect, + token.from, + token.to, + ); + const identifier = path?.length === 1 ? path[0] : undefined; + if (!identifier) { + coverage = "partial"; + } else if (columns.length < MAX_QUERY_OUTPUT_COLUMNS) { + columns.push(Object.freeze({ + definition: Object.freeze({ + from: token.from, + to: token.to, + }), + identifier, + insertText: source.originalText.slice(token.from, token.to), + })); + } else { + coverage = "partial"; + } + expectIdentifier = false; + continue; + } + if ( + !expectIdentifier && + punctuation(source, token, ",") + ) { + expectIdentifier = true; + continue; + } + coverage = "partial"; + } + return { + aliases: Object.freeze({ + columns: Object.freeze(columns), + coverage: "partial", + }), + next: index, + }; +} + export type SqlColumnQuerySiteIssue = | "derived-relation" | "incomplete-relation" + | "local-output-partial" | "nested-query" | "opaque-template-context" | "table-function"; @@ -53,6 +153,7 @@ export type SqlColumnQuerySiteResult = | { readonly status: "ready"; readonly coverage: "complete" | "partial"; + readonly context: SqlColumnCompletionContext; readonly issues: readonly SqlColumnQuerySiteIssue[]; readonly prefix: SqlIdentifierComponent; readonly qualifier: SqlIdentifierPath; @@ -64,7 +165,7 @@ interface Token extends BoundedSqlLexeme { readonly depth: number; } -type Clause = +export type SqlColumnCompletionContext = | "from" | "group" | "having" @@ -73,6 +174,7 @@ type Clause = | "order" | "qualify" | "select-list" + | "using" | "where"; const RELATION_END_WORDS: ReadonlySet = new Set([ @@ -197,6 +299,81 @@ function parseNamedRelation( readonly issue: SqlColumnQuerySiteIssue | null; } { const first = tokens[start]; + if (punctuation(source, first, "(") && first?.depth === depth) { + const closeIndex = tokens.findIndex((token, index) => + index > start && + token.depth === depth && + punctuation(source, token, ")") + ); + const close = tokens[closeIndex]; + const hasSelect = tokens.some((token, index) => + index > start && + index < closeIndex && + token.depth === depth + 1 && + word(source, token) === "select" + ); + if (closeIndex < 0 || !close || !hasSelect) { + return { + issue: "derived-relation" as const, + next: start + 1, + relation: null, + }; + } + let end = closeIndex + 1; + const maybeAs = word(source, tokens[end]); + if (maybeAs === "as") end += 1; + const aliasToken = tokens[end]; + const aliasPath = + isIdentifier(aliasToken) && aliasToken?.depth === depth + ? decodePath( + source, + dialect, + aliasToken.from, + aliasToken.to, + ) + : null; + const alias = aliasPath?.length === 1 + ? aliasPath[0] ?? null + : null; + if (alias !== null) end += 1; + const parsedAliases = alias === null + ? { aliases: undefined, next: end } + : parseColumnAliases( + source, + dialect, + tokens, + end, + depth, + ); + end = parsedAliases.next; + return { + issue: + alias === null + ? "derived-relation" as const + : parsedAliases.aliases?.coverage === "partial" + ? "local-output-partial" as const + : null, + next: end, + relation: Object.freeze({ + alias, + ...(parsedAliases.aliases === undefined + ? {} + : { columnAliases: parsedAliases.aliases }), + local: Object.freeze({ + kind: "derived" as const, + queryRange: Object.freeze({ + from: first.to, + to: close.from, + }), + }), + path: Object.freeze([]), + range: Object.freeze({ + from: first.from, + to: close.to, + }), + }), + }; + } if (!isIdentifier(first) || first?.depth !== depth) { return { issue: punctuation(source, first, "(") @@ -258,14 +435,29 @@ function parseNamedRelation( end += 1; } } + const parsedAliases = alias === null + ? { aliases: undefined, next: end } + : parseColumnAliases( + source, + dialect, + tokens, + end, + depth, + ); + end = parsedAliases.next; return { issue: maybeAs === "as" && alias === null ? "incomplete-relation" - : null, + : parsedAliases.aliases?.coverage === "partial" + ? "local-output-partial" + : null, next: end, relation: Object.freeze({ alias, + ...(parsedAliases.aliases === undefined + ? {} + : { columnAliases: parsedAliases.aliases }), path, range: Object.freeze({ from: first.from, @@ -281,8 +473,8 @@ function clauseAt( selectIndex: number, depth: number, position: number, -): Clause { - let clause: Clause = "select-list"; +): SqlColumnCompletionContext { + let clause: SqlColumnCompletionContext = "select-list"; for ( let index = selectIndex + 1; index < tokens.length && tokens[index]!.from < position; @@ -306,9 +498,11 @@ function clauseAt( clause = "limit"; break; case "on": - case "using": clause = "join-condition"; break; + case "using": + clause = "using"; + break; case "order": clause = "order"; break; @@ -488,6 +682,12 @@ function collectRelations( ); if (parsed.issue !== null) issues.add(parsed.issue); if (parsed.relation !== null) { + if ( + precedingOnly && + parsed.relation.range.to >= visibilityPosition + ) { + break; + } relations.push(parsed.relation); if (relations.length > MAX_COLUMN_QUERY_RELATIONS) return false; } @@ -671,6 +871,7 @@ export function recognizeSqlColumnQuerySite( const issueList = Object.freeze(Array.from(issues).sort()); return Object.freeze({ coverage: issueList.length === 0 ? "complete" : "partial", + context: clause, issues: issueList, prefix: path.prefix, qualifier: path.qualifier, diff --git a/src/cte-layout.ts b/src/cte-layout.ts index ac11cdf..3d83544 100644 --- a/src/cte-layout.ts +++ b/src/cte-layout.ts @@ -14,6 +14,7 @@ import { type SqlStatementIndex, } from "./statement-index.js"; import type { SqlIdentifierComponent } from "./types.js"; +import { MAX_QUERY_OUTPUT_COLUMNS } from "./query-output.js"; const cteRangeBrand: unique symbol = Symbol("SqlCteRange"); @@ -41,6 +42,7 @@ export type SqlCteLayoutIssue = export type SqlCteLayoutResource = | "active-statement" + | "cte-column" | "cte-declaration" | "cte-frame" | "identifier-segment" @@ -51,6 +53,11 @@ export interface SqlCteIdentifier { readonly component: SqlIdentifierComponent; } +export interface SqlCteDeclaredColumn { + readonly name: SqlIdentifierComponent; + readonly range: SqlCteRange; +} + export type SqlCteIdentifierResult = | { readonly status: "identifier"; @@ -84,6 +91,7 @@ export interface SqlCteLayoutDialect { export interface SqlCteDeclaration { readonly ambiguous: boolean; readonly bodyRange: SqlCteRange; + readonly declaredColumns: readonly SqlCteDeclaredColumn[]; readonly equivalenceClass: number; readonly frameIndex: number; readonly name: SqlIdentifierComponent; @@ -96,6 +104,7 @@ export interface SqlCteDeclaration { export interface SqlCteDraftDeclaration { readonly ambiguous: boolean; readonly bodyRange: SqlCteRange; + readonly declaredColumns: readonly SqlCteDeclaredColumn[]; readonly equivalenceClass: number; readonly frameIndex: number; readonly name: SqlIdentifierComponent; @@ -247,6 +256,11 @@ interface DraftDeclaration { bodyLead: "select" | "with" | null; bodyLeadFrameIndex: number | null; component: SqlIdentifierComponent; + declaredColumns: { + readonly from: number; + readonly name: SqlIdentifierComponent; + readonly to: number; + }[]; nameFrom: number; nameTo: number; sourceSpelling: string; @@ -256,6 +270,7 @@ interface MutableDeclaration { ambiguous: boolean; bodyFrom: number; bodyTo: number; + declaredColumns: DraftDeclaration["declaredColumns"]; identityIndex: number; frameIndex: number; name: SqlIdentifierComponent; @@ -269,6 +284,7 @@ interface MutableDraftDeclaration { ambiguous: boolean; bodyFrom: number; bodyTo: number; + declaredColumns: DraftDeclaration["declaredColumns"]; identityIndex: number; frameIndex: number; name: SqlIdentifierComponent; @@ -788,6 +804,7 @@ function processName( bodyLead: null, bodyLeadFrameIndex: null, component: identifier.component, + declaredColumns: [], nameFrom: token.from - statementFrom, nameTo: token.to - statementFrom, sourceSpelling: text.slice(token.from, token.to), @@ -817,6 +834,7 @@ function commitDeclaration( ambiguous: false, bodyFrom: current.bodyFrom, bodyTo, + declaredColumns: [...current.declaredColumns], frameIndex: frame.index, identityIndex, name: current.component, @@ -885,6 +903,7 @@ function collectActiveBodyEvidence( ambiguous, bodyFrom: current.bodyFrom, bodyTo: exactThrough, + declaredColumns: [...current.declaredColumns], frameIndex: frame.index, identityIndex, name: current.component, @@ -960,6 +979,14 @@ function freezeLayout( declaration.bodyFrom, declaration.bodyTo, ), + declaredColumns: Object.freeze( + declaration.declaredColumns.map((column) => + Object.freeze({ + name: column.name, + range: createRange(column.from, column.to), + }) + ), + ), equivalenceClass: relationRoot( relations, declaration.identityIndex, @@ -986,6 +1013,14 @@ function freezeLayout( declaration.bodyFrom, declaration.bodyTo, ), + declaredColumns: Object.freeze( + declaration.declaredColumns.map((column) => + Object.freeze({ + name: column.name, + range: createRange(column.from, column.to), + }) + ), + ), equivalenceClass: relationRoot( relations, declaration.identityIndex, @@ -1125,6 +1160,7 @@ export function analyzeSqlCteLayout( const declarationAttempts = { value: 0 }; let depth = 0; let exactThrough = statementLength; + let recoveredThrough: number | null = null; let resource: SqlCteLayoutResource | undefined; scan: while (true) { @@ -1152,6 +1188,16 @@ export function analyzeSqlCteLayout( frame.issues.add("opaque-template-context"); } } + if ( + bodyOwners.size > 0 || + frames.some((frame) => frame.state === "main") + ) { + recoveredThrough = recoveredThrough === null + ? exactThrough + : Math.min(recoveredThrough, exactThrough); + exactThrough = statementLength; + continue; + } break; } const code = punctuation(text, token); @@ -1528,7 +1574,13 @@ export function analyzeSqlCteLayout( } if ( token.kind === "word" && - wordEquals(text, token, "select") + ( + wordEquals(text, token, "select") || + wordEquals(text, token, "insert") || + wordEquals(text, token, "update") || + wordEquals(text, token, "delete") || + wordEquals(text, token, "merge") + ) ) { frame.mainQueryStart = token.from - statementFrom; frame.state = "main"; @@ -1547,14 +1599,13 @@ export function analyzeSqlCteLayout( const columnOwner = columnOwners.get(depth); if (columnOwner) { if (columnOwner.columnExpectIdentifier) { - if ( - !normalizeIdentifier( - validatedDialect, - text, - token, - "cte-column", - ) - ) { + const column = normalizeIdentifier( + validatedDialect, + text, + token, + "cte-column", + ); + if (!column) { exactThrough = token.from - statementFrom; markPartial( columnOwner, @@ -1563,6 +1614,23 @@ export function analyzeSqlCteLayout( ); break; } + if ( + (columnOwner.current?.declaredColumns.length ?? 0) < + MAX_QUERY_OUTPUT_COLUMNS + ) { + columnOwner.current?.declaredColumns.push({ + from: token.from - statementFrom, + name: column.component, + to: token.to - statementFrom, + }); + } else { + resource ??= "cte-column"; + markPartial( + columnOwner, + issues, + "ambiguous-cte-header", + ); + } columnOwner.columnCount += 1; columnOwner.columnExpectIdentifier = false; continue; @@ -1628,7 +1696,7 @@ export function analyzeSqlCteLayout( relations, exactThrough, ); - const layout = freezeLayout( + const structuralLayout = freezeLayout( frames, declarations, draftDeclarations, @@ -1638,6 +1706,15 @@ export function analyzeSqlCteLayout( issues, resource, ); + const layout = recoveredThrough === null + ? structuralLayout + : Object.freeze({ + ...structuralLayout, + exactThrough: Math.min( + structuralLayout.exactThrough, + recoveredThrough, + ), + }); return registerSqlCteLayout( layout, source, diff --git a/src/index.ts b/src/index.ts index 43d23ed..0439452 100644 --- a/src/index.ts +++ b/src/index.ts @@ -92,6 +92,8 @@ export type { SqlColumnCatalogProviderReport, SqlColumnCatalogFailure, SqlColumnCompletionProvenance, + SqlQueryOutputCompletionProvenance, + SqlQueryOutputProviderReport, SqlCompletionProviderReport, SqlCompletionIssue, SqlCompletionRequest, diff --git a/src/local-relation-site.ts b/src/local-relation-site.ts index 27ad6cd..6d8b439 100644 --- a/src/local-relation-site.ts +++ b/src/local-relation-site.ts @@ -6,8 +6,14 @@ import { } from "./cte-layout.js"; import { recognizeSqlColumnQuerySite, + type SqlColumnQueryRelation, type SqlColumnQuerySiteResult, } from "./column-query-site.js"; +import { + inferSqlQueryOutput, + MAX_QUERY_OUTPUT_COLUMNS, + type SqlQueryOutput, +} from "./query-output.js"; import { recognizeSqlRelationQuerySiteWithEntrypoints, type SqlQuerySiteResult, @@ -24,11 +30,44 @@ import { type SqlStatementIndex, type SqlStatementSlot, } from "./statement-index.js"; +import type { SqlTextRange } from "./types.js"; const localRelationStatementBrand: unique symbol = Symbol( "SqlLocalRelationStatement", ); +export function applySqlQueryOutputAliases( + output: SqlQueryOutput, + aliases: NonNullable, +): SqlQueryOutput { + if (output.status !== "ready") { + return Object.freeze({ + columns: aliases.columns, + coverage: "partial" as const, + status: "ready" as const, + }); + } + const columns = output.columns.map((column, index) => + aliases.columns[index] ?? column + ); + if (aliases.columns.length > columns.length) { + columns.push(...aliases.columns.slice(columns.length)); + } + return Object.freeze({ + columns: Object.freeze( + columns.slice(0, MAX_QUERY_OUTPUT_COLUMNS), + ), + coverage: + aliases.coverage === "partial" || + output.coverage === "partial" || + columns.length > MAX_QUERY_OUTPUT_COLUMNS || + aliases.columns.length > output.columns.length + ? "partial" as const + : "complete" as const, + status: "ready" as const, + }); +} + export interface SqlLocalRelationStatement { readonly [localRelationStatementBrand]: "SqlLocalRelationStatement"; } @@ -227,22 +266,116 @@ export function analyzeSqlLocalColumnSite( context.layout, position - context.slot.source.from, ); - const relations = result.relations.filter((relation) => { - const name = relation.path.length === 1 - ? relation.path[0] - : undefined; - return name === undefined || - !visibility.ctes.some((cte) => - context.dialect.completion.compareCteIdentifiers( - name, - cte.name, - ) === "equal" + const outputCache = new Map(); + const inferOutput = (range: SqlTextRange): SqlQueryOutput => { + const key = `${range.from}:${range.to}`; + const cached = outputCache.get(key); + if (cached) return cached; + const output = inferSqlQueryOutput( + context.source, + range, + context.dialect, + ); + outputCache.set(key, output); + return output; + }; + const applyAliases = ( + output: SqlQueryOutput, + relation: SqlColumnQueryRelation, + ): SqlQueryOutput => { + const aliases = relation.columnAliases; + if (!aliases) return output; + return applySqlQueryOutputAliases(output, aliases); + }; + const relations: SqlColumnQueryRelation[] = result.relations.map( + (relation) => { + if (relation.local?.kind === "derived") { + return Object.freeze({ + ...relation, + local: Object.freeze({ + ...relation.local, + output: applyAliases( + inferOutput(relation.local.queryRange), + relation, + ), + }), + }); + } + const name = relation.path.length === 1 + ? relation.path[0] + : undefined; + const visible = name === undefined + ? undefined + : visibility.ctes.find((cte) => + context.dialect.completion.compareCteIdentifiers( + name, + cte.name, + ) === "equal" + ); + if (!visible) return relation; + const declaration = context.layout.declarations.find((candidate) => + candidate.nameRange.from === visible.declarationPosition ); - }); - return relations.length === result.relations.length - ? result - : Object.freeze({ - ...result, - relations: Object.freeze(relations), + if (!declaration) return relation; + const queryRange = Object.freeze({ + from: context.slot.source.from + declaration.bodyRange.from, + to: context.slot.source.from + declaration.bodyRange.to, + }); + const inferredOutput = inferOutput(queryRange); + const output = declaration.declaredColumns.length === 0 + ? inferredOutput + : Object.freeze({ + columns: Object.freeze( + declaration.declaredColumns.map((column) => { + const definition = Object.freeze({ + from: context.slot.source.from + column.range.from, + to: context.slot.source.from + column.range.to, + }); + return Object.freeze({ + definition, + identifier: column.name, + insertText: context.source.originalText.slice( + definition.from, + definition.to, + ), + }); + }), + ), + coverage: + context.layout.status === "partial" + ? "partial" as const + : "complete" as const, + status: "ready" as const, + }); + return Object.freeze({ + ...relation, + local: Object.freeze({ + kind: "cte" as const, + output: applyAliases(output, relation), + queryRange, + }), }); + }, + ); + const localPartial = relations.some((relation) => + relation.local !== undefined && + ( + relation.local.output?.status !== "ready" || + relation.local.output.coverage === "partial" + ) + ) || visibility.shadowing.coverage === "unknown"; + const issues = localPartial && + !result.issues.includes("local-output-partial") + ? Object.freeze([ + ...result.issues, + "local-output-partial" as const, + ].sort()) + : result.issues; + return Object.freeze({ + ...result, + coverage: + result.coverage === "partial" || localPartial ? "partial" : "complete", + issues, + relations: Object.freeze(relations), + }); } diff --git a/src/query-output.ts b/src/query-output.ts new file mode 100644 index 0000000..61f536c --- /dev/null +++ b/src/query-output.ts @@ -0,0 +1,266 @@ +import { + BoundedSqlLexer, + type BoundedSqlLexeme, +} from "./bounded-sql-lexer.js"; +import type { SqlRelationDialectRuntime } from "./relation-dialect.js"; +import type { SqlSourceSnapshot } from "./source.js"; +import type { + SqlIdentifierComponent, + SqlTextRange, +} from "./types.js"; + +export const MAX_QUERY_OUTPUT_COLUMNS = 256; +export const MAX_QUERY_OUTPUT_LENGTH = 65_536; + +interface Token extends BoundedSqlLexeme { + readonly depth: number; +} + +export interface SqlQueryOutputColumn { + readonly definition: SqlTextRange; + readonly identifier: SqlIdentifierComponent; + readonly insertText: string; +} + +export type SqlQueryOutput = + | { + readonly columns: readonly SqlQueryOutputColumn[]; + readonly coverage: "complete" | "partial"; + readonly status: "ready"; + } + | { + readonly reason: "resource-limit" | "unsupported-query"; + readonly status: "unavailable"; + }; + +function punctuation( + source: SqlSourceSnapshot, + token: Token | undefined, + expected: string, +): boolean { + return token?.kind === "punctuation" && + source.analysisText.slice(token.from, token.to) === expected; +} + +function word( + source: SqlSourceSnapshot, + token: Token | undefined, +): string | null { + return token?.kind === "word" + ? source.analysisText.slice(token.from, token.to).toLowerCase() + : null; +} + +function identifier(token: Token | undefined): boolean { + return token?.kind === "word" || + token?.kind === "quoted-identifier"; +} + +function tokenize( + source: SqlSourceSnapshot, + range: SqlTextRange, + dialect: SqlRelationDialectRuntime, +): readonly Token[] | null { + const lexer = new BoundedSqlLexer( + source, + range.from, + range.to, + dialect.querySite.lexicalProfile, + ); + const tokens: Token[] = []; + let depth = 0; + for (let lexeme = lexer.next(); lexeme; lexeme = lexer.next()) { + if (punctuation(source, { ...lexeme, depth }, ")")) { + depth = Math.max(0, depth - 1); + } + tokens.push(Object.freeze({ ...lexeme, depth })); + if (punctuation(source, { ...lexeme, depth }, "(")) { + depth += 1; + } + } + return lexer.resource === null ? Object.freeze(tokens) : null; +} + +function decodeColumn( + source: SqlSourceSnapshot, + dialect: SqlRelationDialectRuntime, + token: Token, +): SqlQueryOutputColumn | null { + const raw = source.analysisText.slice(token.from, token.to); + const decoded = dialect.querySite.decodeRelationPath(raw, raw.length); + if ( + decoded.status !== "decoded" || + decoded.qualifier.length !== 0 || + decoded.prefix.value.length === 0 + ) { + return null; + } + return Object.freeze({ + definition: Object.freeze({ from: token.from, to: token.to }), + identifier: decoded.prefix, + insertText: source.originalText.slice(token.from, token.to), + }); +} + +function projectionColumn( + source: SqlSourceSnapshot, + dialect: SqlRelationDialectRuntime, + segment: readonly Token[], + depth: number, +): SqlQueryOutputColumn | null { + const significant = segment.filter((token) => + token.kind !== "comment" && token.kind !== "line-comment" + ); + if (significant.length === 0) return null; + for (let index = significant.length - 2; index >= 0; index -= 1) { + const token = significant[index]; + const alias = significant[index + 1]; + if ( + token?.depth === depth && + word(source, token) === "as" && + alias?.depth === depth && + identifier(alias) && + index + 2 === significant.length + ) { + return decodeColumn(source, dialect, alias); + } + } + if (significant.some((token) => token.kind === "barrier")) return null; + const projection = significant.filter((token, index) => + !( + index === 0 && + token.depth === depth && + (word(source, token) === "all" || + word(source, token) === "distinct") + ) + ); + if (projection.length === 0 || projection.some((token) => + token.depth !== depth + )) { + return null; + } + for (let index = 0; index < projection.length; index += 1) { + const token = projection[index]; + if ( + index % 2 === 0 + ? !identifier(token) + : !punctuation(source, token, ".") + ) { + return null; + } + } + if (projection.length % 2 === 0) return null; + const final = projection[projection.length - 1]; + return final ? decodeColumn(source, dialect, final) : null; +} + +export function inferSqlQueryOutput( + source: SqlSourceSnapshot, + range: SqlTextRange, + dialect: SqlRelationDialectRuntime, +): SqlQueryOutput { + if ( + !Number.isSafeInteger(range.from) || + !Number.isSafeInteger(range.to) || + range.from < 0 || + range.from >= range.to || + range.to > source.analysisText.length || + range.to - range.from > MAX_QUERY_OUTPUT_LENGTH + ) { + return Object.freeze({ reason: "resource-limit", status: "unavailable" }); + } + const tokens = tokenize(source, range, dialect); + if (tokens === null) { + return Object.freeze({ reason: "resource-limit", status: "unavailable" }); + } + const selects = tokens + .map((token, index) => ({ index, token })) + .filter(({ token }) => word(source, token) === "select"); + const minimumDepth = Math.min(...selects.map(({ token }) => token.depth)); + const firstTopLevelSet = tokens.findIndex((token) => + token.depth === minimumDepth && + ( + word(source, token) === "union" || + word(source, token) === "intersect" || + word(source, token) === "except" + ) + ); + const precedingArmSelects = firstTopLevelSet < 0 + ? [] + : selects.filter(({ index }) => index < firstTopLevelSet); + const firstArmDepth = Math.min( + ...precedingArmSelects.map(({ token }) => token.depth), + ); + const selected = firstTopLevelSet < 0 + ? selects.find(({ token }) => token.depth === minimumDepth) + : precedingArmSelects.find(({ token }) => + token.depth === firstArmDepth + ); + if (!selected || !Number.isFinite(minimumDepth)) { + return Object.freeze({ + reason: "unsupported-query", + status: "unavailable", + }); + } + const projectionDepth = selected.token.depth; + const segments: Token[][] = [[]]; + let stopped = false; + let complete = !tokens.some((token) => token.kind === "barrier"); + for ( + let index = selected.index + 1; + index < tokens.length; + index += 1 + ) { + const token = tokens[index]!; + const tokenWord = word(source, token); + if ( + ( + token.depth === projectionDepth && + tokenWord === "from" + ) || + ( + token.depth <= projectionDepth && + ( + tokenWord === "union" || + tokenWord === "intersect" || + tokenWord === "except" + ) + ) + ) { + stopped = true; + break; + } + if ( + token.depth === projectionDepth && + punctuation(source, token, ",") + ) { + segments.push([]); + continue; + } + segments[segments.length - 1]!.push(token); + } + const columns: SqlQueryOutputColumn[] = []; + complete &&= stopped || segments.some((segment) => segment.length > 0); + for (const segment of segments) { + const column = projectionColumn( + source, + dialect, + segment, + projectionDepth, + ); + if (column === null) { + complete = false; + continue; + } + if (columns.length === MAX_QUERY_OUTPUT_COLUMNS) { + complete = false; + break; + } + columns.push(column); + } + return Object.freeze({ + columns: Object.freeze(columns), + coverage: complete ? "complete" : "partial", + status: "ready", + }); +} diff --git a/src/query-site.ts b/src/query-site.ts index 33c9480..ea14d75 100644 --- a/src/query-site.ts +++ b/src/query-site.ts @@ -126,6 +126,7 @@ export interface SqlQuerySiteDialect { }; readonly lexicalProfile: SqlLexicalProfile; readonly maximumPathDepth: number; + readonly optionalDmlInto?: boolean; readonly supportsNaturalJoin: boolean; readonly decodeRelationPath: ( rawPath: string, @@ -502,7 +503,13 @@ function processFrameWord( return; } if (isSetOperation(word)) { - markUnavailable(frame, "unsupported-query-site"); + frame.anchor = "from"; + frame.blocksNestedQuery = false; + frame.joinConstraint = null; + frame.joinConstraintAllowed = false; + frame.joinPrefix = null; + frame.selectWords = [null, null, null]; + frame.state = "select-list"; return; } if (word === "lateral" || word === "qualify" || word === "window") { @@ -1073,6 +1080,136 @@ function resultAtGap( return inactive("not-relation-position"); } +function recognizeDmlRelationQuerySite( + source: SqlSourceSnapshot, + slot: ExactSqlStatementSlot, + start: number, + position: number, + dialect: SqlQuerySiteDialect, + lexicalProfile: SqlLexicalProfile, + maximumPathDepth: number, +): SqlQuerySiteResult | "continue-select" | null { + const lexer = new BoundedSqlLexer( + source, + start, + slot.source.to, + lexicalProfile, + ); + let first = lexer.next(); + while (first && isCommentLexeme(first)) { + if (cursorIsInComment(first, position)) { + return inactive("cursor-in-comment"); + } + first = lexer.next(); + } + if (!first || first.kind !== "word") return null; + const leader = wordValue(source.analysisText, first); + let expectedKeyword: "from" | "into" | null; + let keywordOptional = false; + let merge = false; + if (leader === "update") { + expectedKeyword = null; + } else if (leader === "insert") { + expectedKeyword = "into"; + keywordOptional = dialect.optionalDmlInto === true; + } else if (leader === "delete") { + expectedKeyword = "from"; + } else if (leader === "merge") { + expectedKeyword = "into"; + keywordOptional = dialect.optionalDmlInto === true; + merge = true; + } else { + return null; + } + if (first.from <= position && position < first.to) { + return inactive("not-relation-position"); + } + let expectRelation = expectedKeyword === null; + let targetComplete = false; + let anchor: QueryFrame["anchor"] = "from"; + while (true) { + const token = lexer.next(); + if (lexer.resource) { + return unavailable( + "resource-limit", + querySiteLexerResource(lexer.resource), + ); + } + if (!token || token.from > position) { + if (expectRelation) { + const frame = createFrame(0, false); + frame.anchor = anchor; + frame.state = "expect-relation"; + return readyEmpty(slot, frame, position); + } + return inactive("not-relation-position"); + } + if (isCommentLexeme(token)) { + if (cursorIsInComment(token, position)) { + return inactive("cursor-in-comment"); + } + if (!token.closed) return unavailable("ambiguous-query-site"); + continue; + } + if (token.kind === "barrier") { + return token.from <= position && position <= token.to + ? inactive("cursor-in-embedded-region") + : unavailable("ambiguous-query-site"); + } + if (token.from <= position && position < token.to) { + if (token.kind === "string") return inactive("cursor-in-string"); + if (!expectRelation) return inactive("not-relation-position"); + } + if (expectedKeyword !== null) { + const keywordMatches = + token.kind === "word" && + wordValue(source.analysisText, token) === expectedKeyword; + if (!keywordMatches && !keywordOptional) { + return unavailable("ambiguous-query-site"); + } + expectedKeyword = null; + expectRelation = true; + if (keywordMatches) continue; + } + if (expectRelation) { + if ( + token.kind !== "word" && + token.kind !== "quoted-identifier" + ) { + return unavailable("unsupported-query-site"); + } + const frame = createFrame(0, false); + frame.anchor = anchor; + frame.state = "expect-relation"; + const result = recognizePath( + lexer, + source, + slot, + frame, + dialect, + maximumPathDepth, + token, + position, + ); + if (result) return result; + expectRelation = false; + targetComplete = true; + if (!merge || anchor === "join") return "continue-select"; + continue; + } + if ( + merge && + targetComplete && + token.kind === "word" && + wordValue(source.analysisText, token) === "using" + ) { + anchor = "join"; + expectRelation = true; + targetComplete = false; + } + } +} + function recognizeSqlRelationQuerySiteInternal( source: SqlSourceSnapshot, slot: SqlStatementSlot, @@ -1120,6 +1257,29 @@ function recognizeSqlRelationQuerySiteInternal( ) { return unavailable("ambiguous-query-site"); } + const dmlStarts = [ + slot.source.from, + ...(authenticatedEntrypoints ?? []) + .filter((entrypoint) => entrypoint.depth === 0) + .map((entrypoint) => slot.source.from + entrypoint.from), + ]; + let dml: SqlQuerySiteResult | "continue-select" | null = null; + for (let index = dmlStarts.length - 1; index >= 0; index -= 1) { + const start = dmlStarts[index]; + if (start === undefined || start > position) continue; + dml = recognizeDmlRelationQuerySite( + source, + slot, + start, + position, + dialect, + lexicalProfile, + maximumPathDepth, + ); + if (dml !== null) break; + } + const continueAfterDml = dml === "continue-select"; + if (dml !== null && dml !== "continue-select") return dml; const lexer = new BoundedSqlLexer( source, slot.source.from, @@ -1445,11 +1605,11 @@ function recognizeSqlRelationQuerySiteInternal( queryCandidates.has(depth) || isAuthenticatedEntrypoint ) { - queryCandidates.delete(depth); if ( isSelect && !topFrame(frames)?.blocksNestedQuery ) { + queryCandidates.delete(depth); frames.push(createFrame(depth, statementTainted)); sawSelect = true; if (token.to === position) { @@ -1457,6 +1617,9 @@ function recognizeSqlRelationQuerySiteInternal( } continue; } + if (!(continueAfterDml && depth === 0)) { + queryCandidates.delete(depth); + } } const activeFrame = topFrame(frames); let openedRelation = false; diff --git a/src/relation-completion-types.ts b/src/relation-completion-types.ts index c6057fe..03e5448 100644 --- a/src/relation-completion-types.ts +++ b/src/relation-completion-types.ts @@ -3,6 +3,7 @@ import type { SqlIdentifierPath, SqlRevision, SqlTextChange, + SqlTextRange, } from "./types.js"; // Provisional package-private declarations until the vertical slice is proven. @@ -258,6 +259,12 @@ export interface SqlColumnCompletionProvenance { readonly columnEntityId: string; } +export interface SqlQueryOutputCompletionProvenance { + readonly definition: SqlTextRange; + readonly kind: "query-output"; + readonly relation: SqlTextRange; +} + export interface SqlNamespaceCompletionProvenance { readonly containerEntityId: string; readonly epoch: SqlCatalogEpoch; @@ -289,6 +296,11 @@ export type SqlCompletionItem = readonly provenance: SqlColumnCompletionProvenance; readonly relationRequestKey: string; }) + | (SqlCompletionItemBase & { + readonly kind: "column"; + readonly provenance: SqlQueryOutputCompletionProvenance; + readonly relationRequestKey: string; + }) | (SqlCompletionItemBase & { readonly kind: "namespace"; readonly provenance: SqlNamespaceCompletionProvenance; @@ -457,10 +469,17 @@ export type SqlNamespaceCatalogProviderReport = | "provider-failed"; }; +export interface SqlQueryOutputProviderReport { + readonly coverage: "complete" | "partial"; + readonly feature: "query-output"; + readonly outcome: "ready"; +} + export type SqlCompletionProviderReport = | SqlCatalogProviderReport | SqlColumnCatalogProviderReport - | SqlNamespaceCatalogProviderReport; + | SqlNamespaceCatalogProviderReport + | SqlQueryOutputProviderReport; export interface SqlServiceFailure { readonly code: "internal"; diff --git a/src/relation-dialect.ts b/src/relation-dialect.ts index 39c2fe5..f8b4c96 100644 --- a/src/relation-dialect.ts +++ b/src/relation-dialect.ts @@ -1415,6 +1415,7 @@ function createRuntime(spec: RelationDialectSpec): SqlRelationDialectRuntime { decodeRelationPath: createPathDecoder(spec, decodeIdentifier), lexicalProfile: spec.lexicalProfile, maximumPathDepth: spec.maximumPathDepth, + ...(spec.kind === "bigquery" ? { optionalDmlInto: true } : {}), supportsNaturalJoin: spec.supportsNaturalJoin, }); return registerSqlRelationDialectRuntime( diff --git a/src/session.ts b/src/session.ts index 628410a..0178be5 100644 --- a/src/session.ts +++ b/src/session.ts @@ -18,6 +18,8 @@ import { } from "./column-catalog-batch-coordinator.js"; import { composeSqlColumnCompletion, + composeSqlLocalQueryOutputCompletion, + filterSqlUsingCompletionList, prepareSqlColumnCatalogRelations, } from "./column-completion.js"; import { @@ -670,6 +672,22 @@ function mergeCompletionLists( }); } +function completionListWithIssue( + value: SqlCompletionList, + reason: "query-binding-partial", +): SqlCompletionList { + if (value.issues.some((issue) => issue.reason === reason)) return value; + const issues: [SqlCompletionIssue, ...SqlCompletionIssue[]] = [ + Object.freeze({ reason }), + ...value.issues, + ]; + return Object.freeze({ + isIncomplete: true, + issues: Object.freeze(issues), + items: value.items, + }); +} + function completionListWithLoadingLease( value: SqlCompletionList, reason: @@ -1892,8 +1910,25 @@ export class DefaultSqlDocumentSession columnSite, snapshot.dialect.relationDialect, ); + const localComposition = composeSqlLocalQueryOutputCompletion( + columnSite, + snapshot.dialect.relationDialect, + ); if (preparedRelations.references.length === 0) { cancelPrevious(); + if (localComposition) { + return Object.freeze({ + refreshToken: null, + revision: snapshot.revision, + sources: localComposition.sources, + status: "ready", + value: filterSqlUsingCompletionList( + localComposition.value, + columnSite, + snapshot.dialect.relationDialect, + ), + }); + } return Object.freeze({ reason: unavailableCompletionReason(localSite), retryable: false, @@ -1903,6 +1938,22 @@ export class DefaultSqlDocumentSession } if (!catalog || !columnOwner || !columnCoordinator) { cancelPrevious(); + if (localComposition) { + return Object.freeze({ + refreshToken: null, + revision: snapshot.revision, + sources: localComposition.sources, + status: "ready", + value: completionListWithIssue( + filterSqlUsingCompletionList( + localComposition.value, + columnSite, + snapshot.dialect.relationDialect, + ), + "query-binding-partial", + ), + }); + } return Object.freeze({ reason: "unsupported-query-site", retryable: false, @@ -1939,24 +1990,37 @@ export class DefaultSqlDocumentSession if (raced.kind === "timeout") { const remainingIntentLeaseMs = this.#retainAuxiliaryRefresh(active, ticket.result); + const loadingValue = Object.freeze({ + isIncomplete: true as const, + issues: Object.freeze([Object.freeze({ + reason: "column-catalog-loading" as const, + remainingIntentLeaseMs, + })] as const), + items: Object.freeze([]), + }); return Object.freeze({ refreshToken: active.token, revision: snapshot.revision, - sources: Object.freeze([Object.freeze({ - feature: "column-catalog" as const, - failures: Object.freeze([]), - outcome: "loading" as const, - providerId: columnCoordinator.providerId, - })]), + sources: Object.freeze([ + ...(localComposition?.sources ?? []), + Object.freeze({ + feature: "column-catalog" as const, + failures: Object.freeze([]), + outcome: "loading" as const, + providerId: columnCoordinator.providerId, + }), + ]), status: "ready", - value: Object.freeze({ - isIncomplete: true, - issues: Object.freeze([Object.freeze({ - reason: "column-catalog-loading" as const, - remainingIntentLeaseMs, - })] as const), - items: Object.freeze([]), - }), + value: localComposition + ? filterSqlUsingCompletionList( + mergeCompletionLists( + localComposition.value, + loadingValue, + ), + columnSite, + snapshot.dialect.relationDialect, + ) + : loadingValue, }); } const composition = composeSqlColumnCompletion({ @@ -1995,15 +2059,33 @@ export class DefaultSqlDocumentSession return Object.freeze({ refreshToken: retryLoading ? active.token : null, revision: snapshot.revision, - sources: composition.sources, + sources: Object.freeze([ + ...(localComposition?.sources ?? []), + ...composition.sources, + ]), status: "ready", - value: retryLoading - ? completionListWithLoadingLease( - composition.value, - "column-catalog-loading", - remainingIntentLeaseMs, + value: localComposition + ? filterSqlUsingCompletionList( + mergeCompletionLists( + localComposition.value, + retryLoading + ? completionListWithLoadingLease( + composition.value, + "column-catalog-loading", + remainingIntentLeaseMs, + ) + : composition.value, + ), + columnSite, + snapshot.dialect.relationDialect, ) - : composition.value, + : retryLoading + ? completionListWithLoadingLease( + composition.value, + "column-catalog-loading", + remainingIntentLeaseMs, + ) + : composition.value, }); } cancelPrevious(); diff --git a/test/corpora/semantic-completion/bigquery.ts b/test/corpora/semantic-completion/bigquery.ts new file mode 100644 index 0000000..b3860ac --- /dev/null +++ b/test/corpora/semantic-completion/bigquery.ts @@ -0,0 +1,35 @@ +import type { SemanticCompletionCase } from "./types.js"; + +export const bigQuerySemanticCompletion = [ + { + category: "valid", + expectedIncomplete: false, + expectedLabels: ["Display Name"], + sql: "WITH c AS (SELECT name AS `Display Name`) SELECT c.| FROM c", + }, + { + category: "invalid", + expectedIncomplete: true, + expectedLabels: ["safe_id"], + sql: "WITH c AS (SELECT id AS safe_id, STRUCT() AS) SELECT c.| FROM c", + }, + { + category: "incomplete", + expectedIncomplete: false, + expectedLabels: ["project_id"], + sql: "WITH c AS (SELECT id AS project_id) SELECT c.proj| FROM c", + }, + { + category: "templated", + expectedIncomplete: true, + expectedLabels: ["metric"], + sql: "WITH c AS (SELECT {python} AS metric) SELECT c.| FROM c", + template: "{python}", + }, + { + category: "multi-statement", + expectedIncomplete: false, + expectedLabels: ["event_id"], + sql: "SELECT 1; WITH c AS (SELECT id AS event_id) SELECT c.| FROM c", + }, +] as const satisfies readonly SemanticCompletionCase[]; diff --git a/test/corpora/semantic-completion/duckdb.ts b/test/corpora/semantic-completion/duckdb.ts new file mode 100644 index 0000000..dfbaf9b --- /dev/null +++ b/test/corpora/semantic-completion/duckdb.ts @@ -0,0 +1,35 @@ +import type { SemanticCompletionCase } from "./types.js"; + +export const duckDbSemanticCompletion = [ + { + category: "valid", + expectedIncomplete: false, + expectedLabels: ["first_id"], + sql: "SELECT d.| FROM (SELECT id AS first_id) d", + }, + { + category: "invalid", + expectedIncomplete: true, + expectedLabels: ["known_id"], + sql: "SELECT d.| FROM (SELECT id AS known_id, count() +) d", + }, + { + category: "incomplete", + expectedIncomplete: false, + expectedLabels: ["first_id"], + sql: "SELECT d.fi| FROM (SELECT id AS first_id) d", + }, + { + category: "templated", + expectedIncomplete: true, + expectedLabels: ["dynamic_id"], + sql: "SELECT d.| FROM (SELECT {python} AS dynamic_id) d", + template: "{python}", + }, + { + category: "multi-statement", + expectedIncomplete: false, + expectedLabels: ["first_name"], + sql: "SELECT 1; SELECT d.| FROM (SELECT id AS first_name UNION ALL SELECT id AS later_name) d", + }, +] as const satisfies readonly SemanticCompletionCase[]; diff --git a/test/corpora/semantic-completion/postgresql.ts b/test/corpora/semantic-completion/postgresql.ts new file mode 100644 index 0000000..a5b95ca --- /dev/null +++ b/test/corpora/semantic-completion/postgresql.ts @@ -0,0 +1,35 @@ +import type { SemanticCompletionCase } from "./types.js"; + +export const postgresqlSemanticCompletion = [ + { + category: "valid", + expectedIncomplete: false, + expectedLabels: ["user_id"], + sql: "WITH c AS (SELECT id AS user_id) SELECT c.| FROM c", + }, + { + category: "invalid", + expectedIncomplete: true, + expectedLabels: ["safe_name"], + sql: "WITH c AS (SELECT id AS safe_name, upper(name) FROM) SELECT c.| FROM c", + }, + { + category: "incomplete", + expectedIncomplete: false, + expectedLabels: ["user_id"], + sql: "WITH c AS (SELECT id AS user_id) SELECT c.us| FROM c", + }, + { + category: "templated", + expectedIncomplete: true, + expectedLabels: ["computed"], + sql: "WITH c AS (SELECT {python} AS computed) SELECT c.| FROM c", + template: "{python}", + }, + { + category: "multi-statement", + expectedIncomplete: false, + expectedLabels: ["second_id"], + sql: "SELECT 1; WITH c AS (SELECT id AS second_id) SELECT c.| FROM c", + }, +] as const satisfies readonly SemanticCompletionCase[]; diff --git a/test/corpora/semantic-completion/types.ts b/test/corpora/semantic-completion/types.ts new file mode 100644 index 0000000..611e58a --- /dev/null +++ b/test/corpora/semantic-completion/types.ts @@ -0,0 +1,14 @@ +export type SemanticCompletionCategory = + | "incomplete" + | "invalid" + | "multi-statement" + | "templated" + | "valid"; + +export interface SemanticCompletionCase { + readonly category: SemanticCompletionCategory; + readonly expectedIncomplete: boolean; + readonly expectedLabels: readonly string[]; + readonly sql: string; + readonly template?: string; +} diff --git a/test/worker-placement/README.md b/test/worker-placement/README.md index 564c035..d7076e7 100644 --- a/test/worker-placement/README.md +++ b/test/worker-placement/README.md @@ -9,23 +9,23 @@ package's dependency. The fixture workspace pins Vite's floating transitive versions to the exact versions in the root lock, so the nested frozen install can run offline after a clean root CI install. -The minified Vite 8 packed-consumer baseline is 40,121 gzip/143,265 raw +The minified Vite 8 packed-consumer baseline is 53,612 gzip/199,957 raw bytes for the complete parser-free core, 67,573 gzip bytes for the PostgreSQL transitive graph, 50,470 gzip bytes for the BigQuery transitive graph, and -150,985 gzip/669,106 raw bytes for the complete worker build output. The core +164,513 gzip/725,798 raw bytes for the complete worker build output. The core measurement includes the four authenticated relation-dialect runtimes, their reserved-word tables, and relation-completion/session orchestration. The PostgreSQL and BigQuery figures each include their transitive shared chunks; the report also identifies those shared chunks explicitly. -The fail-closed ceilings retain approximately 18% gzip and 22% raw headroom -for the core, 8% gzip and 7% raw headroom for the complete worker output, and +The fail-closed ceilings retain approximately 2–3% headroom for the core and +complete worker output, plus the existing tight dialect-graph headroom: -- Complete parser-free core: 48 KiB gzip and 180 KiB raw +- Complete parser-free core: 54 KiB gzip and 200 KiB raw - PostgreSQL transitive graph: 68 KiB gzip - BigQuery transitive graph: 50 KiB gzip -- Complete worker build output: 160 KiB gzip and 700 KiB raw +- Complete worker build output: 164 KiB gzip and 720 KiB raw These are provisional placement limits, not product bundle promises. The orchestration script fails closed when they are exceeded, when the dialects no From 4691df8ebbb02d519caca2106c8b984caf4988c8 Mon Sep 17 00:00:00 2001 From: Shahmir Varqha Date: Mon, 27 Jul 2026 18:39:52 +0800 Subject: [PATCH 2/2] fix: address semantic completion review findings --- src/__tests__/column-query-site.test.ts | 44 +++++++++++++++++++++++++ src/__tests__/query-output.test.ts | 11 +++++++ src/__tests__/query-site.test.ts | 13 ++++++++ src/__tests__/session.test.ts | 20 +++++++---- src/column-query-site.ts | 5 ++- src/query-output.ts | 4 ++- src/query-site.ts | 5 ++- src/session.ts | 9 ++++- 8 files changed, 101 insertions(+), 10 deletions(-) diff --git a/src/__tests__/column-query-site.test.ts b/src/__tests__/column-query-site.test.ts index 64637d0..f46c2fb 100644 --- a/src/__tests__/column-query-site.test.ts +++ b/src/__tests__/column-query-site.test.ts @@ -286,6 +286,50 @@ describe("recognizeSqlColumnQuerySite", () => { }]); }); + it.each([ + "WHERE", + "ON", + "GROUP", + "ORDER", + "UNION", + ])("does not consume %s as a derived relation alias", (keyword) => { + const result = ready( + analyze( + `SELECT i| FROM (SELECT id FROM users) ${keyword} other`, + ), + ); + const derived = result.relations.find((relation) => + relation.local?.kind === "derived" + ); + expect(derived).toMatchObject({ + alias: null, + local: { kind: "derived" }, + path: [], + }); + expect(result.issues).toContain("derived-relation"); + }); + + it("continues parsing a join after an unaliased derived relation", () => { + const result = ready( + analyze( + "SELECT o.| FROM (SELECT id FROM users) JOIN orders o ON true", + { dialect: DUCKDB_SQL_RELATION_DIALECT }, + ), + ); + expect(result.relations).toMatchObject([ + { + alias: null, + local: { kind: "derived" }, + path: [], + }, + { + alias: { value: "o" }, + path: [{ value: "orders" }], + }, + ]); + expect(result.issues).toContain("derived-relation"); + }); + it("preserves authoritative relation-alias column lists", () => { const result = ready( analyze( diff --git a/src/__tests__/query-output.test.ts b/src/__tests__/query-output.test.ts index 819e33d..4c7fa80 100644 --- a/src/__tests__/query-output.test.ts +++ b/src/__tests__/query-output.test.ts @@ -73,6 +73,17 @@ describe("query output inference", () => { }); }); + it.each([ + "SELECT id /* unterminated", + "SELECT id FROM users WHERE name = 'unterminated", + ])("marks output inference partial for %s", (text) => { + expect(infer(text)).toMatchObject({ + columns: [{ identifier: { value: "id" } }], + coverage: "partial", + status: "ready", + }); + }); + it("uses the first set-operation arm for output names", () => { expect( infer( diff --git a/src/__tests__/query-site.test.ts b/src/__tests__/query-site.test.ts index ea9d326..027133d 100644 --- a/src/__tests__/query-site.test.ts +++ b/src/__tests__/query-site.test.ts @@ -1589,6 +1589,19 @@ describe("fail-closed query-site behavior", () => { }); }); + it.each([ + "SELECT * FROM users WHERE active UNION SELECT * FROM |", + "SELECT * FROM users GROUP BY team_id HAVING count(*) > 1 INTERSECT SELECT * FROM |", + "SELECT * FROM users ORDER BY id UNION ALL SELECT * FROM |", + ])("reopens a closed first arm at a set operation in %s", (marked) => { + expect(recognize(marked)).toMatchObject({ + anchor: "from", + prefix: { quoted: false, value: "" }, + qualifier: [], + status: "ready", + }); + }); + it.each([ ["INSERT INTO |", "from"], ["UPDATE app.us| SET name = 'x'", "from"], diff --git a/src/__tests__/session.test.ts b/src/__tests__/session.test.ts index 11f69cf..2e5056d 100644 --- a/src/__tests__/session.test.ts +++ b/src/__tests__/session.test.ts @@ -314,8 +314,9 @@ describe("relation completion session integration", () => { service.dispose(); }); - it("completes relations in a later set-operation arm", async () => { - const text = "SELECT * FROM current UNION ALL SELECT * FROM us"; + it("completes relations after a closed set-operation arm", async () => { + const text = + "SELECT * FROM current WHERE active UNION ALL SELECT * FROM us"; const service = createSqlLanguageService({ catalog: catalogProvider(async () => ({ coverage: { kind: "complete" }, @@ -1583,7 +1584,7 @@ describe("column completion session integration", () => { service.dispose(); }); - it("merges inferred and physical columns in one completion result", async () => { + it("merges inferred and physical columns without duplicate issues", async () => { const service = serviceWithColumns(async (request) => ({ epoch: { generation: 1, token: "mixed" }, relations: request.relations.map((relation) => ({ @@ -1600,7 +1601,7 @@ describe("column completion session integration", () => { })), })); const text = - "WITH c AS (SELECT id AS local_id) " + + "WITH c AS (SELECT id AS local_id, *) " + "SELECT FROM c JOIN physical p ON true"; const session = service.openDocument({ context: { @@ -1611,22 +1612,29 @@ describe("column completion session integration", () => { text, }); - await expect(session.complete({ + const result = await session.complete({ position: text.indexOf(" FROM"), trigger: { kind: "invoked" }, - })).resolves.toMatchObject({ + }); + expect(result).toMatchObject({ sources: [ { feature: "query-output" }, { feature: "column-catalog" }, ], status: "ready", value: { + isIncomplete: true, items: [ { label: "local_id" }, { label: "invoice_id" }, ], }, }); + if (result.status === "ready") { + expect(result.value.issues.map((issue) => issue.reason)).toEqual([ + "query-binding-partial", + ]); + } service.dispose(); }); diff --git a/src/column-query-site.ts b/src/column-query-site.ts index e16de49..f0752d5 100644 --- a/src/column-query-site.ts +++ b/src/column-query-site.ts @@ -323,8 +323,11 @@ function parseNamedRelation( const maybeAs = word(source, tokens[end]); if (maybeAs === "as") end += 1; const aliasToken = tokens[end]; + const aliasWord = aliasToken ? word(source, aliasToken) : null; const aliasPath = - isIdentifier(aliasToken) && aliasToken?.depth === depth + isIdentifier(aliasToken) && + aliasToken?.depth === depth && + (aliasWord === null || !RELATION_END_WORDS.has(aliasWord)) ? decodePath( source, dialect, diff --git a/src/query-output.ts b/src/query-output.ts index 61f536c..d730e65 100644 --- a/src/query-output.ts +++ b/src/query-output.ts @@ -205,7 +205,9 @@ export function inferSqlQueryOutput( const projectionDepth = selected.token.depth; const segments: Token[][] = [[]]; let stopped = false; - let complete = !tokens.some((token) => token.kind === "barrier"); + let complete = !tokens.some((token) => + token.kind === "barrier" || !token.closed + ); for ( let index = selected.index + 1; index < tokens.length; diff --git a/src/query-site.ts b/src/query-site.ts index ea14d75..3ee6ccc 100644 --- a/src/query-site.ts +++ b/src/query-site.ts @@ -463,10 +463,13 @@ function processFrameWord( text: string, token: Lexeme, ): void { - if (frame.state === "unavailable" || frame.state === "closed") { + if (frame.state === "unavailable") { return; } const word = wordValue(text, token); + if (frame.state === "closed" && !isSetOperation(word)) { + return; + } if (frame.state === "expect-alias") { if (word.length === 0) { markUnavailable(frame, "ambiguous-query-site"); diff --git a/src/session.ts b/src/session.ts index 0178be5..bd25033 100644 --- a/src/session.ts +++ b/src/session.ts @@ -652,7 +652,14 @@ function mergeCompletionLists( right: SqlCompletionList, ): SqlCompletionList { const items = Object.freeze([...left.items, ...right.items]); - const issues = Object.freeze([...left.issues, ...right.issues]); + const issueReasons = new Set(); + const issues = Object.freeze( + [...left.issues, ...right.issues].filter((issue) => { + if (issueReasons.has(issue.reason)) return false; + issueReasons.add(issue.reason); + return true; + }), + ); const first = issues[0]; if (first === undefined) { return Object.freeze({