feat(vnext): authenticate CTE main query sites

Author: Light2Dark

src/vnext/__tests__/cte-layout.bench.ts

@@ -17,17 +17,19 @@ const dialect = DUCKDB_SQL_RELATION_DIALECT.cteLayout;
 
 function fixture(text: string): {
   readonly source: ReturnType<typeof createIdentitySqlSource>;
+  readonly index: ReturnType<typeof buildSqlStatementIndex>;
   readonly slot: ExactSqlStatementSlot;
 } {
   const source = createIdentitySqlSource(text);
-  const slot = buildSqlStatementIndex(
+  const index = buildSqlStatementIndex(
     source.analysisText,
     dialect.lexicalProfile,
-  ).slots[0];
+  );
+  const slot = index.slots[0];
   if (!slot || slot.boundaryQuality !== "exact") {
     throw new Error("CTE benchmark fixture requires an exact statement");
   }
-  return { slot, source };
+  return { index, slot, source };
 }
 
 const tenKibibytes = 10 * 1_024;
@@ -62,6 +64,7 @@ const bareFrameHeavyText = `SELECT ${Array.from(
 const bareFrameHeavyFixture = fixture(bareFrameHeavyText);
 const depthLayout = analyzeSqlCteLayout(
   depthHeavyFixture.source,
+  depthHeavyFixture.index,
   depthHeavyFixture.slot,
   dialect,
 );
@@ -70,6 +73,7 @@ if (depthLayout.status !== "ready") {
 }
 const projectedLayout = analyzeSqlCteLayout(
   declarationHeavyFixture.source,
+  declarationHeavyFixture.index,
   declarationHeavyFixture.slot,
   dialect,
 );
@@ -81,6 +85,7 @@ describe("CTE layout", () => {
   bench("ordinary 10 KiB statement", () => {
     analyzeSqlCteLayout(
       ordinaryFixture.source,
+      ordinaryFixture.index,
       ordinaryFixture.slot,
       dialect,
     );
@@ -89,6 +94,7 @@ describe("CTE layout", () => {
   bench("256 declarations", () => {
     analyzeSqlCteLayout(
       declarationHeavyFixture.source,
+      declarationHeavyFixture.index,
       declarationHeavyFixture.slot,
       dialect,
     );
@@ -97,6 +103,7 @@ describe("CTE layout", () => {
   bench("128-depth nested CTE", () => {
     analyzeSqlCteLayout(
       depthHeavyFixture.source,
+      depthHeavyFixture.index,
       depthHeavyFixture.slot,
       dialect,
     );
@@ -105,6 +112,7 @@ describe("CTE layout", () => {
   bench("256 sequential incomplete frames", () => {
     analyzeSqlCteLayout(
       bareFrameHeavyFixture.source,
+      bareFrameHeavyFixture.index,
       bareFrameHeavyFixture.slot,
       dialect,
     );

src/vnext/__tests__/cte-layout.test.ts

@@ -55,6 +55,7 @@ function analyze(
   expect(slot?.boundaryQuality).toBe("exact");
   const result = analyzeSqlCteLayout(
     source,
+    index,
     slot as ExactSqlStatementSlot,
     dialect,
   );
@@ -68,15 +69,18 @@ function analyze(
 function analyzeRaw(
   text: string,
   dialect: SqlCteLayoutDialect = postgres,
+  indexDialect: SqlCteLayoutDialect = dialect,
 ): SqlCteLayout {
   const source = createIdentitySqlSource(text);
-  const slot = buildSqlStatementIndex(
+  const index = buildSqlStatementIndex(
     source.analysisText,
-    dialect.lexicalProfile,
-  ).slots[0];
+    indexDialect.lexicalProfile,
+  );
+  const slot = index.slots[0];
   expect(slot?.boundaryQuality).toBe("exact");
   return analyzeSqlCteLayout(
     source,
+    index,
     slot as ExactSqlStatementSlot,
     dialect,
   );
@@ -400,6 +404,7 @@ describe("bounded CTE layout", () => {
     }
     const layout = analyzeSqlCteLayout(
       source,
+      index,
       slot,
       postgres,
     );
@@ -1022,7 +1027,9 @@ describe("bounded CTE layout", () => {
         },
       },
     };
-    expect(analyzeRaw("SELECT 1", getterLexicalProfile)).toEqual({
+    expect(
+      analyzeRaw("SELECT 1", getterLexicalProfile, postgres),
+    ).toEqual({
       reason: "resource-limit",
       status: "unavailable",
     });

src/vnext/__tests__/incremental-statement-index.test.ts

@@ -4,6 +4,7 @@ import {
   buildSqlStatementIndex,
   DREMIO_SQL_LEXICAL_PROFILE,
   DUCKDB_SQL_LEXICAL_PROFILE,
+  isExactSqlStatementSlotSnapshotFor,
   MAX_SQL_STATEMENT_SLOTS,
   POSTGRESQL_SQL_LEXICAL_PROFILE,
   type SqlLexicalProfile,
@@ -358,18 +359,93 @@ describe("incremental statement index reuse", () => {
     },
   );
 
+  it("rebuilds when empty change metadata does not match the text", () => {
+    const previousIndex = buildSqlStatementIndex(
+      "SELECT 1",
+      POSTGRESQL_SQL_LEXICAL_PROFILE,
+    );
+    const nextIndex = updateSqlStatementIndex(
+      previousIndex,
+      "SELECT 2",
+      [],
+      POSTGRESQL_SQL_LEXICAL_PROFILE,
+    );
+    expect(nextIndex).not.toBe(previousIndex);
+    expect(nextIndex).toEqual(
+      buildSqlStatementIndex(
+        "SELECT 2",
+        POSTGRESQL_SQL_LEXICAL_PROFILE,
+      ),
+    );
+  });
+
+  it("rebuilds when the lexical profile changes", () => {
+    const text = "SELECT 1";
+    const previousIndex = buildSqlStatementIndex(
+      text,
+      DUCKDB_SQL_LEXICAL_PROFILE,
+    );
+    const nextIndex = updateSqlStatementIndex(
+      previousIndex,
+      text,
+      [],
+      POSTGRESQL_SQL_LEXICAL_PROFILE,
+    );
+    expect(nextIndex).not.toBe(previousIndex);
+    expect(nextIndex).toEqual(
+      buildSqlStatementIndex(
+        text,
+        POSTGRESQL_SQL_LEXICAL_PROFILE,
+      ),
+    );
+  });
+
+  it("rebuilds when one lexical profile object mutates", () => {
+    const profile = { ...POSTGRESQL_SQL_LEXICAL_PROFILE };
+    const text = "SELECT # ;\n SELECT 2";
+    const previousIndex = buildSqlStatementIndex(text, profile);
+    expect(previousIndex.slots).toHaveLength(2);
+    profile.hashLineComments = true;
+    const nextIndex = updateSqlStatementIndex(
+      previousIndex,
+      text,
+      [],
+      profile,
+    );
+    expect(nextIndex).not.toBe(previousIndex);
+    expect(nextIndex).toEqual(buildSqlStatementIndex(text, profile));
+    expect(nextIndex.slots).toHaveLength(1);
+  });
+
   it("reuses unchanged prefix and zero-delta suffix slot identities", () => {
     const profile = DUCKDB_SQL_LEXICAL_PROFILE;
     const text = "SELECT 1; SELECT 2; SELECT 3";
     const change = replaceFirst(text, "2", "9");
-    const { nextIndex, previousIndex } = expectIncrementalMatchesOracle(
-      text,
-      [change],
-      profile,
-    );
-    expect(expectSlot(nextIndex.slots, 0)).toBe(
-      expectSlot(previousIndex.slots, 0),
-    );
+    const { nextIndex, nextText, previousIndex } =
+      expectIncrementalMatchesOracle(
+        text,
+        [change],
+        profile,
+      );
+    const sharedPrefix = expectSlot(nextIndex.slots, 0);
+    expect(sharedPrefix).toBe(expectSlot(previousIndex.slots, 0));
+    expect(sharedPrefix.boundaryQuality).toBe("exact");
+    expect(
+      isExactSqlStatementSlotSnapshotFor(
+        previousIndex,
+        sharedPrefix,
+        text,
+        profile,
+      ),
+    ).toBe(true);
+    expect(
+      isExactSqlStatementSlotSnapshotFor(
+        nextIndex,
+        sharedPrefix,
+        nextText,
+        profile,
+      ),
+    ).toBe(true);
     expect(expectSlot(nextIndex.slots, 1)).not.toBe(
       expectSlot(previousIndex.slots, 1),
     );

src/vnext/__tests__/query-site.bench.ts

@@ -1,8 +1,15 @@
 import { bench, describe } from "vitest";
+import {
+  analyzeSqlCteLayout,
+  MAX_CTE_DECLARATIONS,
+} from "../cte-layout.js";
 import {
   recognizeSqlRelationQuerySite,
 } from "../query-site.js";
 import { DUCKDB_SQL_RELATION_DIALECT } from "../relation-dialect.js";
+import {
+  recognizeSqlRelationQuerySiteWithCteLayout,
+} from "../relation-query-site.js";
 import { createIdentitySqlSource } from "../source.js";
 import {
   buildSqlStatementIndex,
@@ -61,6 +68,103 @@ const usingHeavySlot = findSqlStatementSlot(
   usingHeavyPosition,
   "left",
 );
+const cteQueryPrefix =
+  "WITH cte_name AS (SELECT 1) SELECT ";
+const cteQuerySuffix = " FROM schema_prefix";
+const cteProjectedListLength =
+  TEN_KIBIBYTES -
+  cteQueryPrefix.length -
+  cteQuerySuffix.length;
+const cteProjectedList = `${"x,".repeat(
+  Math.floor((cteProjectedListLength - 1) / 2),
+)}x`;
+const tenKilobyteCteQuery =
+  `${cteQueryPrefix}${cteProjectedList}` +
+  `${" ".repeat(
+    cteProjectedListLength - cteProjectedList.length,
+  )}${cteQuerySuffix}`;
+if (tenKilobyteCteQuery.length !== TEN_KIBIBYTES) {
+  throw new Error("CTE query benchmark fixture must be exactly 10 KiB");
+}
+const cteSource = createIdentitySqlSource(tenKilobyteCteQuery);
+const cteIndex = buildSqlStatementIndex(
+  cteSource.analysisText,
+  dialect.lexicalProfile,
+);
+const ctePosition = tenKilobyteCteQuery.length;
+const cteSlot = findSqlStatementSlot(
+  cteIndex,
+  ctePosition,
+  "left",
+);
+if (cteSlot.boundaryQuality === "opaque") {
+  throw new Error("CTE query benchmark requires an exact statement");
+}
+const maximumCteQuery =
+  `WITH ${Array.from(
+    { length: MAX_CTE_DECLARATIONS },
+    (_, index) => `cte_${index} AS (SELECT ${index})`,
+  ).join(", ")} SELECT * FROM target`;
+const maximumCteSource =
+  createIdentitySqlSource(maximumCteQuery);
+const maximumCteIndex = buildSqlStatementIndex(
+  maximumCteSource.analysisText,
+  dialect.lexicalProfile,
+);
+const maximumCtePosition = maximumCteQuery.length;
+const maximumCteSlot = findSqlStatementSlot(
+  maximumCteIndex,
+  maximumCtePosition,
+  "left",
+);
+if (maximumCteSlot.boundaryQuality === "opaque") {
+  throw new Error("Maximum CTE benchmark requires an exact statement");
+}
+const ctePreflightLayout = analyzeSqlCteLayout(
+  cteSource,
+  cteIndex,
+  cteSlot,
+  DUCKDB_SQL_RELATION_DIALECT.cteLayout,
+);
+const ctePreflightResult =
+  recognizeSqlRelationQuerySiteWithCteLayout(
+    cteSource,
+    cteSlot,
+    ctePosition,
+    DUCKDB_SQL_RELATION_DIALECT,
+    ctePreflightLayout,
+  );
+if (
+  ctePreflightLayout.status !== "ready" ||
+  ctePreflightResult.status !== "ready" ||
+  ctePreflightResult.anchor !== "from"
+) {
+  throw new Error("CTE query benchmark must exercise a ready FROM site");
+}
+const maximumCtePreflightLayout = analyzeSqlCteLayout(
+  maximumCteSource,
+  maximumCteIndex,
+  maximumCteSlot,
+  DUCKDB_SQL_RELATION_DIALECT.cteLayout,
+);
+const maximumCtePreflightResult =
+  recognizeSqlRelationQuerySiteWithCteLayout(
+    maximumCteSource,
+    maximumCteSlot,
+    maximumCtePosition,
+    DUCKDB_SQL_RELATION_DIALECT,
+    maximumCtePreflightLayout,
+  );
+if (
+  maximumCtePreflightLayout.status !== "ready" ||
+  maximumCtePreflightLayout.declarations.length !==
+    MAX_CTE_DECLARATIONS ||
+  maximumCtePreflightLayout.mainQueryEntrypoints.length !== 1 ||
+  maximumCtePreflightResult.status !== "ready" ||
+  maximumCtePreflightResult.prefix.value !== "target"
+) {
+  throw new Error("Maximum CTE benchmark must exercise all declarations");
+}
 
 describe("query-site recognizer", () => {
   bench("10 KiB active statement", () => {
@@ -84,4 +188,36 @@ describe("query-site recognizer", () => {
       dialect,
     );
   });
+
+  bench("10 KiB WITH main-query double scan", () => {
+    const layout = analyzeSqlCteLayout(
+      cteSource,
+      cteIndex,
+      cteSlot,
+      DUCKDB_SQL_RELATION_DIALECT.cteLayout,
+    );
+    recognizeSqlRelationQuerySiteWithCteLayout(
+      cteSource,
+      cteSlot,
+      ctePosition,
+      DUCKDB_SQL_RELATION_DIALECT,
+      layout,
+    );
+  });
+
+  bench("256 CTE main-query double scan", () => {
+    const layout = analyzeSqlCteLayout(
+      maximumCteSource,
+      maximumCteIndex,
+      maximumCteSlot,
+      DUCKDB_SQL_RELATION_DIALECT.cteLayout,
+    );
+    recognizeSqlRelationQuerySiteWithCteLayout(
+      maximumCteSource,
+      maximumCteSlot,
+      maximumCtePosition,
+      DUCKDB_SQL_RELATION_DIALECT,
+      layout,
+    );
+  });
 });