From 8663a8443686dd1dc405a6dea32e18e46c292235 Mon Sep 17 00:00:00 2001 From: anandgupta42 Date: Thu, 15 Oct 2026 10:16:39 -0700 Subject: [PATCH 1/2] fix(schema): batch DuckDB introspection and bound indexing - `duckdb`: restrict schema/table/column introspection to the current catalog and de-duplicate schemas, so attached catalogs no longer produce repeated schemas, duplicate tables or merged columns - `duckdb`: add optional `introspectSchema(signal)` connector capability returning schemas, tables and columns in two queries independent of table count - `schema cache`: prefer bulk metadata when available, dedupe schema/table names for other connectors, and replace cache rows plus the warehouse summary in a single transaction - `schema_index`: forward session cancellation, apply a 30s default deadline (including connection acquisition), report partial results and stopping reason, and prevent late cache writes after cancellation - tests: driver, cache, tool and real-DuckDB integration regressions (native tests gated on `ALTIMATE_DUCKDB_E2E=1`) Co-Authored-By: Claude Opus 5.5 --- packages/drivers/src/duckdb.ts | 77 +++++-- packages/drivers/src/index.ts | 2 +- packages/drivers/src/types.ts | 6 + packages/drivers/test/duckdb-schema.test.ts | 113 ++++++++++ .../src/altimate/native/schema/cache.ts | 198 +++++++++--------- .../altimate/native/schema/index-control.ts | 41 ++++ .../src/altimate/native/schema/register.ts | 25 ++- .../opencode/src/altimate/native/types.ts | 4 + .../src/altimate/tools/schema-index.ts | 11 +- .../test/altimate/schema-index-duckdb.test.ts | 54 +++++ .../test/altimate/schema-index.test.ts | 150 +++++++++++++ .../test/altimate/schema-indexing.test.ts | 154 ++++++++++++++ 12 files changed, 718 insertions(+), 117 deletions(-) create mode 100644 packages/drivers/test/duckdb-schema.test.ts create mode 100644 packages/opencode/src/altimate/native/schema/index-control.ts create mode 100644 packages/opencode/test/altimate/schema-index-duckdb.test.ts create mode 100644 packages/opencode/test/altimate/schema-index.test.ts create mode 100644 packages/opencode/test/altimate/schema-indexing.test.ts diff --git a/packages/drivers/src/duckdb.ts b/packages/drivers/src/duckdb.ts index a39b6fa53d..8856133ccd 100644 --- a/packages/drivers/src/duckdb.ts +++ b/packages/drivers/src/duckdb.ts @@ -3,7 +3,7 @@ */ import { assertStoreExists, requireStorePath } from "./file-store" -import type { ConnectionConfig, Connector, ConnectorResult, ExecuteOptions, SchemaColumn } from "./types" +import type { ConnectionConfig, Connector, ConnectorResult, ExecuteOptions, SchemaColumn, SchemaMetadata } from "./types" import { loadOptionalDriver } from "./resolve" // altimate_change start — configurable, generous open budget @@ -101,15 +101,34 @@ export async function connect(config: ConnectionConfig): Promise { } // altimate_change end - function query(sql: string): Promise { + function query(sql: string, signal?: AbortSignal): Promise { return new Promise((resolve, reject) => { - connection.all(sql, (err: Error | null, rows: any[]) => { - if (err) reject(wrapDuckDBError(err)) - else resolve(rows ?? []) - }) + signal?.throwIfAborted() + // The legacy DuckDB SDK's interrupt() is a no-op. Stop waiting and avoid + // subsequent work on abort; the native query still finishes in the SDK. + const onAbort = () => reject(signal!.reason) + signal?.addEventListener("abort", onAbort, { once: true }) + try { + connection.all(sql, (err: Error | null, rows: any[]) => { + signal?.removeEventListener("abort", onAbort) + if (err) reject(wrapDuckDBError(err)) + else resolve(rows ?? []) + }) + } catch (err) { + signal?.removeEventListener("abort", onAbort) + reject(err) + } }) } + async function schemaNames(signal?: AbortSignal): Promise { + const rows = await query( + "SELECT DISTINCT schema_name FROM information_schema.schemata WHERE catalog_name = current_database() ORDER BY schema_name", + signal, + ) + return rows.map((r) => r.schema_name as string) + } + function queryWithParams(sql: string, params: any[]): Promise { return new Promise((resolve, reject) => { connection.all(sql, ...params, (err: Error | null, rows: any[]) => { @@ -282,17 +301,14 @@ export async function connect(config: ConnectionConfig): Promise { }, async listSchemas(): Promise { - const rows = await query( - "SELECT schema_name FROM information_schema.schemata ORDER BY schema_name", - ) - return rows.map((r) => r.schema_name as string) + return schemaNames() }, async listTables( schema: string, ): Promise> { const rows = await queryWithParams( - `SELECT table_name, table_type FROM information_schema.tables WHERE table_schema = $1 ORDER BY table_name`, + `SELECT table_name, table_type FROM information_schema.tables WHERE table_catalog = current_database() AND table_schema = $1 ORDER BY table_name`, [schema], ) return rows.map((r) => ({ @@ -306,7 +322,7 @@ export async function connect(config: ConnectionConfig): Promise { table: string, ): Promise { const rows = await queryWithParams( - `SELECT column_name, data_type, is_nullable FROM information_schema.columns WHERE table_schema = $1 AND table_name = $2 ORDER BY ordinal_position`, + `SELECT column_name, data_type, is_nullable FROM information_schema.columns WHERE table_catalog = current_database() AND table_schema = $1 AND table_name = $2 ORDER BY ordinal_position`, [schema, table], ) return rows.map((r) => ({ @@ -316,6 +332,43 @@ export async function connect(config: ConnectionConfig): Promise { })) }, + async introspectSchema(signal?: AbortSignal): Promise { + const schemas = await schemaNames(signal) + const rows = await query( + `SELECT t.table_schema, t.table_name, t.table_type, + c.column_name, c.data_type, c.is_nullable + FROM information_schema.tables t + LEFT JOIN information_schema.columns c + ON c.table_catalog = t.table_catalog + AND c.table_schema = t.table_schema + AND c.table_name = t.table_name + WHERE t.table_catalog = current_database() + ORDER BY t.table_schema, t.table_name, c.ordinal_position`, + signal, + ) + const tables: SchemaMetadata["tables"] = [] + let table: SchemaMetadata["tables"][number] | undefined + for (const row of rows) { + if (!table || table.schema !== row.table_schema || table.name !== row.table_name) { + table = { + schema: row.table_schema, + name: row.table_name, + type: row.table_type === "VIEW" ? "view" : "table", + columns: [], + } + tables.push(table) + } + if (row.column_name !== null) { + table.columns.push({ + name: row.column_name, + data_type: row.data_type, + nullable: row.is_nullable === "YES", + }) + } + } + return { schemas, tables } + }, + async close() { if (db) { await new Promise((resolve) => { diff --git a/packages/drivers/src/index.ts b/packages/drivers/src/index.ts index d89aa73ced..5bc69ae955 100644 --- a/packages/drivers/src/index.ts +++ b/packages/drivers/src/index.ts @@ -1,5 +1,5 @@ // Re-export types -export type { Connector, ConnectorResult, SchemaColumn, ConnectionConfig } from "./types" +export type { Connector, ConnectorResult, SchemaColumn, SchemaMetadata, ConnectionConfig } from "./types" // Re-export config normalization export { normalizeConfig, sanitizeConnectionString } from "./normalize" diff --git a/packages/drivers/src/types.ts b/packages/drivers/src/types.ts index 3bc3760d6c..6bec217883 100644 --- a/packages/drivers/src/types.ts +++ b/packages/drivers/src/types.ts @@ -20,6 +20,11 @@ export interface SchemaColumn { nullable: boolean } +export interface SchemaMetadata { + schemas: string[] + tables: Array<{ schema: string; name: string; type: string; columns: SchemaColumn[] }> +} + export interface ExecuteOptions { /** Skip the default LIMIT injection and post-truncation. Use when the caller * needs the complete, untruncated result set (e.g. data-diff pipelines). */ @@ -32,5 +37,6 @@ export interface Connector { listSchemas(): Promise listTables(schema: string): Promise> describeTable(schema: string, table: string): Promise + introspectSchema?(signal?: AbortSignal): Promise close(): Promise } diff --git a/packages/drivers/test/duckdb-schema.test.ts b/packages/drivers/test/duckdb-schema.test.ts new file mode 100644 index 0000000000..6fbc227f60 --- /dev/null +++ b/packages/drivers/test/duckdb-schema.test.ts @@ -0,0 +1,113 @@ +import { afterEach, describe, expect, spyOn, test } from "bun:test" +import { connect } from "../src/duckdb" +import type { Connector } from "../src/types" + +// Run in a separate process because other suites mock the DuckDB module. +const duckdbTest = process.env.ALTIMATE_DUCKDB_E2E === "1" ? test : test.skip + +let connector: Connector | undefined +let allSpy: ReturnType | undefined +let duckdb: typeof import("duckdb") + +afterEach(async () => { + allSpy?.mockRestore() + allSpy = undefined + await connector?.close() + connector = undefined +}) + +async function open() { + duckdb = await import("duckdb").then((module) => module.default) + connector = await connect({ type: "duckdb", path: ":memory:" }) + await connector.connect() + return connector +} + +describe("DuckDB schema introspection", () => { + duckdbTest("lists each schema once and scopes tables and columns to the current catalog", async () => { + const driver = await open() + await driver.execute("CREATE TABLE shared (local_id INTEGER NOT NULL)") + await driver.execute("ATTACH ':memory:' AS other") + await driver.execute("CREATE TABLE other.main.shared (foreign_id VARCHAR)") + await driver.execute("CREATE TABLE other.main.foreign_only (id INTEGER)") + + expect(await driver.listSchemas()).toEqual(["main"]) + expect(await driver.listTables("main")).toEqual([{ name: "shared", type: "table" }]) + expect(await driver.describeTable("main", "shared")).toEqual([ + { name: "local_id", data_type: "INTEGER", nullable: false }, + ]) + }) + + duckdbTest("reads complete table and view metadata with a constant number of queries", async () => { + const driver = await open() + await driver.execute("CREATE SCHEMA empty_schema") + await driver.execute("CREATE TABLE events (id INTEGER NOT NULL, label VARCHAR)") + await driver.execute("CREATE VIEW event_ids AS SELECT id FROM events") + await driver.execute("ATTACH ':memory:' AS other") + await driver.execute("CREATE TABLE other.main.events (foreign_id BIGINT)") + for (let i = 0; i < 24; i++) { + await driver.execute(`CREATE TABLE item_${i} (id INTEGER)`) + } + allSpy = spyOn(duckdb.Connection.prototype, "all") + + const metadata = await driver.introspectSchema!() + + expect(allSpy).toHaveBeenCalledTimes(2) + expect(metadata.schemas).toEqual(["empty_schema", "main"]) + expect(metadata.tables).toHaveLength(26) + expect(metadata.tables.find((table) => table.name === "events")).toEqual({ + schema: "main", + name: "events", + type: "table", + columns: [ + { name: "id", data_type: "INTEGER", nullable: false }, + { name: "label", data_type: "VARCHAR", nullable: true }, + ], + }) + expect(metadata.tables.find((table) => table.name === "event_ids")).toMatchObject({ + schema: "main", + type: "view", + columns: [{ name: "id", data_type: "INTEGER", nullable: true }], + }) + }) + + duckdbTest("does not truncate a large metadata inventory", async () => { + const driver = await open() + await driver.execute( + `CREATE TABLE wide_table (${Array.from({ length: 1100 }, (_, i) => `field_${i} INTEGER`).join(", ")})`, + ) + + const metadata = await driver.introspectSchema!() + + expect(metadata.tables).toHaveLength(1) + expect(metadata.tables[0].columns).toHaveLength(1100) + expect(metadata.tables[0].columns.at(-1)?.name).toBe("field_1099") + }) + + duckdbTest("rejects an already aborted introspection without querying", async () => { + const driver = await open() + allSpy = spyOn(duckdb.Connection.prototype, "all") + const controller = new AbortController() + controller.abort(new Error("Index cancelled")) + + await expect(driver.introspectSchema!(controller.signal)).rejects.toThrow("Index cancelled") + expect(allSpy).not.toHaveBeenCalled() + }) + + duckdbTest("stops waiting on abort and starts no further metadata queries", async () => { + const driver = await open() + let complete: ((error: Error | null, rows: unknown[]) => void) | undefined + allSpy = spyOn(duckdb.Connection.prototype, "all").mockImplementation(function (_sql, ...args) { + complete = args.at(-1) as typeof complete + return this + }) + const controller = new AbortController() + const pending = driver.introspectSchema!(controller.signal) + controller.abort(new Error("Index cancelled")) + + await expect(pending).rejects.toThrow("Index cancelled") + complete!(null, [{ schema_name: "main" }]) + await Promise.resolve() + expect(allSpy).toHaveBeenCalledTimes(1) + }) +}) diff --git a/packages/opencode/src/altimate/native/schema/cache.ts b/packages/opencode/src/altimate/native/schema/cache.ts index 5f4750b778..8b66e4fc57 100644 --- a/packages/opencode/src/altimate/native/schema/cache.ts +++ b/packages/opencode/src/altimate/native/schema/cache.ts @@ -10,7 +10,8 @@ import { Database } from "bun:sqlite" import * as path from "path" import * as os from "os" import * as fs from "fs" -import type { Connector } from "@altimateai/drivers" +import type { Connector, SchemaMetadata } from "@altimateai/drivers" +import { awaitIndex, withIndexSignal } from "./index-control" import type { SchemaIndexResult, SchemaSearchResult, @@ -146,110 +147,113 @@ export class SchemaCache { warehouseName: string, warehouseType: string, connector: Connector, + options: { signal?: AbortSignal; timeoutMs?: number } = {}, ): Promise { - const now = new Date().toISOString() + return withIndexSignal(options, async (signal) => { + const now = new Date().toISOString() + const schemas = new Set() + const tables = new Map() + let note: string | undefined + const includeSchema = (name: string) => name.toUpperCase() !== "INFORMATION_SCHEMA" - // Clear existing data - this.db.prepare("DELETE FROM columns_cache WHERE warehouse = ?").run(warehouseName) - this.db.prepare("DELETE FROM tables_cache WHERE warehouse = ?").run(warehouseName) - - let totalSchemas = 0 - let totalTables = 0 - let totalColumns = 0 - const databaseName: string | null = null - - let schemas: string[] = [] - try { - schemas = await connector.listSchemas() - } catch { - // ignore - } - - const insertTable = this.db.prepare( - `INSERT OR REPLACE INTO tables_cache - (warehouse, database_name, schema_name, table_name, table_type, search_text) - VALUES (?, ?, ?, ?, ?, ?)`, - ) - - const insertColumn = this.db.prepare( - `INSERT OR REPLACE INTO columns_cache - (warehouse, database_name, schema_name, table_name, column_name, data_type, nullable, search_text) - VALUES (?, ?, ?, ?, ?, ?, ?, ?)`, - ) - - // Batch inserts per-table inside a transaction to avoid per-statement disk fsyncs. - // The async connector calls (listTables, describeTable) run outside the transaction; - // only the synchronous SQLite inserts are wrapped. - const insertTableBatch = this.db.transaction( - (tableArgs: any[], columnArgsBatch: any[][]) => { - insertTable.run(...tableArgs) - for (const colArgs of columnArgsBatch) { - insertColumn.run(...colArgs) - } - }, - ) - - for (const schemaName of schemas) { - if (schemaName.toUpperCase() === "INFORMATION_SCHEMA") continue - totalSchemas++ - - let tables: Array<{ name: string; type: string }> = [] try { - tables = await connector.listTables(schemaName) - } catch { - continue - } - - for (const tableInfo of tables) { - totalTables++ - const searchText = makeSearchText(databaseName, schemaName, tableInfo.name, tableInfo.type) - - let columns: Array<{ name: string; data_type: string; nullable: boolean }> = [] - try { - columns = await connector.describeTable(schemaName, tableInfo.name) - } catch { - // continue with empty columns - } - - // Build column insert args - const columnArgsBatch: any[][] = [] - for (const col of columns) { - totalColumns++ - const colSearch = makeSearchText( - databaseName, schemaName, tableInfo.name, col.name, col.data_type, - ) - columnArgsBatch.push([ - warehouseName, databaseName, schemaName, tableInfo.name, - col.name, col.data_type, col.nullable ? 1 : 0, colSearch, - ]) + if (connector.introspectSchema) { + const metadata = await awaitIndex(signal, () => connector.introspectSchema!(signal)) + for (const schema of metadata.schemas) { + if (includeSchema(schema)) schemas.add(schema) + } + for (const table of metadata.tables) { + if (!includeSchema(table.schema)) continue + schemas.add(table.schema) + tables.set(JSON.stringify([table.schema, table.name]), table) + } + } else { + const names = await awaitIndex(signal, () => connector.listSchemas()) + for (const schema of new Set(names)) { + if (!includeSchema(schema)) continue + signal.throwIfAborted() + schemas.add(schema) + let inventory: Awaited> + try { + inventory = await awaitIndex(signal, () => connector.listTables(schema)) + } catch { + signal.throwIfAborted() + note = "Some schema metadata could not be read." + continue + } + for (const table of inventory) { + const key = JSON.stringify([schema, table.name]) + if (tables.has(key)) continue + let columns: SchemaMetadata["tables"][number]["columns"] = [] + try { + columns = await awaitIndex(signal, () => connector.describeTable(schema, table.name)) + } catch { + signal.throwIfAborted() + note = "Some schema metadata could not be read." + } + tables.set(key, { ...table, schema, columns }) + } + } } + } catch (error) { + note = signal.aborted + ? signal.reason?.name === "TimeoutError" + ? "Schema indexing deadline exceeded." + : "Schema indexing aborted." + : `Schema metadata could not be read: ${error instanceof Error ? error.message : String(error)}` + } - // Insert table + all its columns in a single transaction - insertTableBatch( - [warehouseName, databaseName, schemaName, tableInfo.name, tableInfo.type, searchText], - columnArgsBatch, + let totalColumns = 0 + // Gather metadata outside SQLite, then replace the snapshot in one transaction. + // A cancelled refresh with no completed tables leaves the previous cache intact. + if (!note || tables.size > 0) { + const insertTable = this.db.prepare( + `INSERT INTO tables_cache + (warehouse, database_name, schema_name, table_name, table_type, search_text) + VALUES (?, ?, ?, ?, ?, ?)`, + ) + const insertColumn = this.db.prepare( + `INSERT INTO columns_cache + (warehouse, database_name, schema_name, table_name, column_name, data_type, nullable, search_text) + VALUES (?, ?, ?, ?, ?, ?, ?, ?)`, ) + this.db.transaction(() => { + this.db.prepare("DELETE FROM columns_cache WHERE warehouse = ?").run(warehouseName) + this.db.prepare("DELETE FROM tables_cache WHERE warehouse = ?").run(warehouseName) + for (const table of tables.values()) { + insertTable.run( + warehouseName, null, table.schema, table.name, table.type, + makeSearchText(table.schema, table.name, table.type), + ) + for (const column of table.columns) { + insertColumn.run( + warehouseName, null, table.schema, table.name, + column.name, column.data_type, column.nullable ? 1 : 0, + makeSearchText(table.schema, table.name, column.name, column.data_type), + ) + totalColumns++ + } + } + this.db.prepare( + `INSERT OR REPLACE INTO warehouses + (name, type, last_indexed, databases_count, schemas_count, tables_count, columns_count) + VALUES (?, ?, ?, ?, ?, ?, ?)`, + ).run(warehouseName, warehouseType, now, 0, schemas.size, tables.size, totalColumns) + })() + } else { + note += " No cached metadata was replaced." } - } - // Update warehouse summary - this.db.prepare( - `INSERT OR REPLACE INTO warehouses - (name, type, last_indexed, databases_count, schemas_count, tables_count, columns_count) - VALUES (?, ?, ?, ?, ?, ?, ?)`, - ).run( - warehouseName, warehouseType, now, - databaseName ? 1 : 0, totalSchemas, totalTables, totalColumns, - ) - - return { - warehouse: warehouseName, - type: warehouseType, - schemas_indexed: totalSchemas, - tables_indexed: totalTables, - columns_indexed: totalColumns, - timestamp: now, - } + return { + warehouse: warehouseName, + type: warehouseType, + schemas_indexed: schemas.size, + tables_indexed: tables.size, + columns_indexed: totalColumns, + timestamp: now, + ...(note ? { partial: true, note } : {}), + } + }) } /** diff --git a/packages/opencode/src/altimate/native/schema/index-control.ts b/packages/opencode/src/altimate/native/schema/index-control.ts new file mode 100644 index 0000000000..6aee317533 --- /dev/null +++ b/packages/opencode/src/altimate/native/schema/index-control.ts @@ -0,0 +1,41 @@ +export const DEFAULT_INDEX_TIMEOUT_MS = 30_000 + +export async function withIndexSignal( + options: { signal?: AbortSignal; timeoutMs?: number }, + work: (signal: AbortSignal) => Promise, +): Promise { + const controller = new AbortController() + const signal = options.signal + ? AbortSignal.any([options.signal, controller.signal]) + : controller.signal + const timer = setTimeout(() => { + controller.abort(new DOMException("Schema indexing deadline exceeded.", "TimeoutError")) + }, options.timeoutMs ?? DEFAULT_INDEX_TIMEOUT_MS) + try { + return await work(signal) + } finally { + clearTimeout(timer) + } +} + +export async function awaitIndex(signal: AbortSignal, work: () => Promise): Promise { + signal.throwIfAborted() + let abort!: () => void + const cancelled = new Promise((_, reject) => { + abort = () => reject(signal.reason) + signal.addEventListener("abort", abort, { once: true }) + }) + try { + const result = await Promise.race([ + Promise.resolve().then(() => { + signal.throwIfAborted() + return work() + }), + cancelled, + ]) + signal.throwIfAborted() + return result + } finally { + signal.removeEventListener("abort", abort) + } +} diff --git a/packages/opencode/src/altimate/native/schema/register.ts b/packages/opencode/src/altimate/native/schema/register.ts index 30675f796e..67c69b3e6d 100644 --- a/packages/opencode/src/altimate/native/schema/register.ts +++ b/packages/opencode/src/altimate/native/schema/register.ts @@ -4,6 +4,7 @@ import { register } from "../dispatcher" import { getCache } from "./cache" +import { awaitIndex, withIndexSignal } from "./index-control" import { detectPii } from "./pii-detector" import { getTags, listTags } from "./tags" import * as Registry from "../connections/registry" @@ -28,13 +29,29 @@ export function registerAll(): void { // --- schema.index --- register("schema.index", async (params: SchemaIndexParams): Promise => { const startTime = Date.now() - const connector = await Registry.get(params.warehouse) const config = Registry.getConfig(params.warehouse) const warehouseType = config?.type || "unknown" - const cache = await getCache() try { - const result = await cache.indexWarehouse(params.warehouse, warehouseType, connector) + const result = await withIndexSignal({ signal: params.signal, timeoutMs: params.timeout_ms }, async (signal) => { + try { + const connector = await awaitIndex(signal, () => Registry.get(params.warehouse)) + const cache = await awaitIndex(signal, () => getCache()) + return await cache.indexWarehouse(params.warehouse, warehouseType, connector, { signal, timeoutMs: params.timeout_ms }) + } catch (error) { + if (!signal.aborted) throw error + return { + warehouse: params.warehouse, + type: warehouseType, + schemas_indexed: 0, + tables_indexed: 0, + columns_indexed: 0, + timestamp: new Date().toISOString(), + partial: true, + note: signal.reason instanceof Error ? signal.reason.message : "Schema indexing cancelled.", + } + } + }) try { Telemetry.track({ type: "warehouse_introspection", @@ -42,7 +59,7 @@ register("schema.index", async (params: SchemaIndexParams): Promise { + await using tmp = await tmpdir() + const module = await loadOptionalDriver("duckdb", "duckdb") + const duckdb = module.default ?? module + const connector = await connect({ type: "duckdb", path: ":memory:" }) + const cache = SchemaCache.create(path.join(tmp.path, "schema-cache.db")) + let metadataSpy: ReturnType | undefined + + try { + await connector.connect() + const tableCount = 3000 + for (let i = 0; i < tableCount; i++) { + await connector.execute(`CREATE TABLE item_${i} (id INTEGER NOT NULL, label VARCHAR)`) + } + await connector.execute("ATTACH ':memory:' AS secondary") + await connector.execute("CREATE TABLE secondary.main.item_0 (foreign_field BOOLEAN)") + await connector.execute("CREATE TABLE secondary.main.foreign_only (id INTEGER)") + metadataSpy = spyOn(duckdb.Connection.prototype, "all") + const started = performance.now() + + const result = await cache.indexWarehouse("warehouse", "duckdb", connector) + + expect(performance.now() - started).toBeLessThan(20_000) + expect(metadataSpy.mock.calls.length).toBeGreaterThan(0) + expect(metadataSpy.mock.calls.length).toBeLessThanOrEqual(3) + expect(result.schemas_indexed).toBe(1) + expect(result.tables_indexed).toBe(tableCount) + expect(result.columns_indexed).toBe(tableCount * 2) + expect(cache.cacheStatus()).toMatchObject({ + total_tables: tableCount, + total_columns: tableCount * 2, + }) + expect(cache.search("foreign_field", "warehouse").columns).toEqual([]) + expect(cache.search("foreign_only", "warehouse").tables).toEqual([]) + expect(cache.search("item_0", "warehouse").columns).toEqual([ + expect.objectContaining({ table: "item_0", name: "id", data_type: "INTEGER", nullable: false }), + expect.objectContaining({ table: "item_0", name: "label", data_type: "VARCHAR", nullable: true }), + ]) + } finally { + metadataSpy?.mockRestore() + cache.close() + await connector.close() + } +}, 60_000) diff --git a/packages/opencode/test/altimate/schema-index.test.ts b/packages/opencode/test/altimate/schema-index.test.ts new file mode 100644 index 0000000000..9d8ea2ecb0 --- /dev/null +++ b/packages/opencode/test/altimate/schema-index.test.ts @@ -0,0 +1,150 @@ +import { afterEach, beforeEach, describe, expect, mock, spyOn, test } from "bun:test" +import type { Connector } from "@altimateai/drivers" +import * as Dispatcher from "../../src/altimate/native/dispatcher" +import * as Registry from "../../src/altimate/native/connections/registry" +import * as SchemaCache from "../../src/altimate/native/schema/cache" +import { registerAll } from "../../src/altimate/native/schema/register" +import { SchemaIndexTool } from "../../src/altimate/tools/schema-index" +import { PostConnectSuggestions } from "../../src/altimate/tools/post-connect-suggestions" +import { MessageID, SessionID } from "../../src/session/schema" +import { initTool } from "./tool-fixture" + +const completeResult = { + warehouse: "inventory", + type: "duckdb", + schemas_indexed: 1, + tables_indexed: 2, + columns_indexed: 4, + timestamp: "2026-01-01T00:00:00.000Z", +} + +function context(signal: AbortSignal) { + return { + sessionID: SessionID.make("ses_schema"), + messageID: MessageID.make("msg_schema"), + agent: "build", + abort: signal, + messages: [], + } +} + +beforeEach(() => { + Dispatcher.setRegistrationHook(async () => registerAll()) + PostConnectSuggestions.resetShownSuggestions() +}) + +afterEach(() => { + mock.restore() + PostConnectSuggestions.resetShownSuggestions() +}) + +describe("schema_index tool", () => { + test("forwards session cancellation to native indexing", async () => { + const call = spyOn(Dispatcher, "call").mockResolvedValue(completeResult) + const controller = new AbortController() + const tool = await initTool(SchemaIndexTool) + + await tool.execute({ warehouse: "inventory" }, context(controller.signal)) + + expect(call).toHaveBeenCalledWith("schema.index", { + warehouse: "inventory", + signal: controller.signal, + }) + }) + + test("reports partial results without suggesting a completed cache", async () => { + spyOn(Dispatcher, "call").mockResolvedValue({ + ...completeResult, + partial: true, + note: "Schema indexing deadline exceeded.", + }) + const tool = await initTool(SchemaIndexTool) + + const result = await tool.execute({ warehouse: "inventory" }, context(new AbortController().signal)) + + expect(result.title).toContain("Partial") + expect(result.metadata.partial).toBe(true) + expect(result.metadata.note).toBe("Schema indexing deadline exceeded.") + expect(result.output).toContain("Schema indexing deadline exceeded.") + expect(result.output).not.toContain("Schema cache is now ready") + expect(result.output).not.toContain("sql_analyze") + }) +}) + +describe("schema.index native handler", () => { + test("passes the cancellation signal into schema caching", async () => { + const connector = {} as Connector + spyOn(Registry, "get").mockResolvedValue(connector) + spyOn(Registry, "getConfig").mockReturnValue({ type: "duckdb" }) + const started = Promise.withResolvers() + const indexWarehouse = mock(async (_name: string, _type: string, _connector: Connector, options: { signal: AbortSignal }) => { + started.resolve() + return new Promise((resolve) => { + options.signal.addEventListener("abort", () => resolve({ + ...completeResult, + partial: true, + note: "Session cancelled.", + }), { once: true }) + }) + }) + spyOn(SchemaCache, "getCache").mockResolvedValue({ indexWarehouse } as any) + const controller = new AbortController() + + const pending = Dispatcher.call("schema.index", { warehouse: "inventory", signal: controller.signal }) + await started.promise + controller.abort() + const result = await pending + + const args = indexWarehouse.mock.calls[0]! + expect(args[0]).toBe("inventory") + expect(args[1]).toBe("duckdb") + expect(args[2]).toBe(connector) + expect(args[3]?.signal).toBeInstanceOf(AbortSignal) + expect(result).toMatchObject({ ...completeResult, partial: true, note: "Session cancelled." }) + }) + + test("cancels while waiting for a warehouse connection", async () => { + const started = Promise.withResolvers() + const connection = Promise.withResolvers() + spyOn(Registry, "get").mockImplementation(() => { + started.resolve() + return connection.promise + }) + spyOn(Registry, "getConfig").mockReturnValue({ type: "duckdb" }) + const cache = spyOn(SchemaCache, "getCache") + const controller = new AbortController() + const pending = Dispatcher.call("schema.index", { + warehouse: "inventory", + signal: controller.signal, + }) + await started.promise + + controller.abort(new Error("Session cancelled.")) + const result = await pending + + expect(result).toMatchObject({ + warehouse: "inventory", + partial: true, + tables_indexed: 0, + note: "Session cancelled.", + }) + expect(cache).not.toHaveBeenCalled() + connection.resolve({} as Connector) + await connection.promise + await Promise.resolve() + expect(cache).not.toHaveBeenCalled() + }, 1_000) + + test("applies the indexing deadline to warehouse connection acquisition", async () => { + spyOn(Registry, "get").mockImplementation(() => new Promise(() => {})) + spyOn(Registry, "getConfig").mockReturnValue({ type: "duckdb" }) + + const result = await Dispatcher.call("schema.index", { + warehouse: "inventory", + timeout_ms: 5, + }) + + expect(result).toMatchObject({ partial: true, tables_indexed: 0 }) + expect(result.note).toContain("deadline") + }, 1_000) +}) diff --git a/packages/opencode/test/altimate/schema-indexing.test.ts b/packages/opencode/test/altimate/schema-indexing.test.ts new file mode 100644 index 0000000000..5b8a23a67b --- /dev/null +++ b/packages/opencode/test/altimate/schema-indexing.test.ts @@ -0,0 +1,154 @@ +import { afterEach, expect, mock, test } from "bun:test" +import type { Connector } from "@altimateai/drivers" +import { SchemaCache } from "../../src/altimate/native/schema/cache" + +const caches: SchemaCache[] = [] +afterEach(() => { + for (const cache of caches.splice(0)) cache.close() +}) + +function createCache() { + const cache = SchemaCache.createInMemory() + caches.push(cache) + return cache +} + +function connector(): Connector { + return { + connect: async () => {}, + close: async () => {}, + execute: async () => ({ columns: [], rows: [], row_count: 0, truncated: false }), + listSchemas: mock(async () => ["main"]), + listTables: mock(async () => [{ name: "orders", type: "table" }]), + describeTable: mock(async () => [{ name: "id", data_type: "INTEGER", nullable: false }]), + } +} + +test("uses bulk metadata for a large schema without per-table requests", async () => { + const cache = createCache() + const source = connector() + source.introspectSchema = mock(async () => ({ + schemas: ["main", "empty"], + tables: Array.from({ length: 3000 }, (_, i) => ({ + schema: "main", + name: `item_${i}`, + type: i === 0 ? "view" : "table", + columns: [{ name: "id", data_type: "INTEGER", nullable: false }], + })), + })) + + const result = await cache.indexWarehouse("warehouse", "duckdb", source) + expect(result.tables_indexed).toBe(3000) + expect(result.columns_indexed).toBe(3000) + expect(result.schemas_indexed).toBe(2) + expect(source.introspectSchema).toHaveBeenCalledTimes(1) + expect(source.listSchemas).not.toHaveBeenCalled() + expect(source.listTables).not.toHaveBeenCalled() + expect(source.describeTable).not.toHaveBeenCalled() + expect(cache.search("item_0").tables[0].type).toBe("view") + expect(cache.listColumns("warehouse")[0].nullable).toBe(false) + expect(cache.cacheStatus().total_tables).toBe(3000) +}) + +test("deduplicates fallback schema and table names before describing them", async () => { + const cache = createCache() + const source = connector() + source.listSchemas = mock(async () => ["main", "main", "main"]) + source.listTables = mock(async () => [ + { name: "orders", type: "table" }, + { name: "orders", type: "table" }, + ]) + + const result = await cache.indexWarehouse("warehouse", "other", source) + expect(result.schemas_indexed).toBe(1) + expect(result.tables_indexed).toBe(1) + expect(result.columns_indexed).toBe(1) + expect(source.listTables).toHaveBeenCalledTimes(1) + expect(source.describeTable).toHaveBeenCalledTimes(1) + expect(cache.cacheStatus().total_tables).toBe(1) +}) + +test("aborting a pending description saves completed tables and stops the crawl", async () => { + const cache = createCache() + const source = connector() + const controller = new AbortController() + const started = Promise.withResolvers() + const pending = Promise.withResolvers>>() + source.listTables = async () => ["first", "second", "third"].map((name) => ({ name, type: "table" })) + source.describeTable = mock(async (_schema: string, table: string) => { + if (table === "second") { + started.resolve() + return pending.promise + } + return [{ name: "id", data_type: "INTEGER", nullable: false }] + }) + const indexing = cache.indexWarehouse("warehouse", "other", source, { signal: controller.signal }) + await started.promise + controller.abort() + try { + const result = await Promise.race([ + indexing, + Bun.sleep(100).then(() => { throw new Error("indexing did not stop on abort") }), + ]) + expect(result.partial).toBe(true) + expect(result.note).toMatch(/abort/i) + expect(result.tables_indexed).toBe(1) + expect(source.describeTable).toHaveBeenCalledTimes(2) + } finally { + pending.resolve([{ name: "late_column", data_type: "TEXT", nullable: true }]) + await indexing + } + expect(cache.search("late_column").match_count).toBe(0) + expect(cache.cacheStatus().total_tables).toBe(1) +}) + +test("a deadline stops waiting for bulk metadata and signals the connector", async () => { + const cache = createCache() + const source = connector() + const pending = Promise.withResolvers<{ schemas: string[]; tables: [] }>() + let signal: AbortSignal | undefined + source.introspectSchema = async (input) => { + signal = input + return pending.promise + } + const result = await cache.indexWarehouse("warehouse", "duckdb", source, { timeoutMs: 5 }) + pending.resolve({ schemas: [], tables: [] }) + expect(result.partial).toBe(true) + expect(result.note).toMatch(/deadline|timed out/i) + expect(result.tables_indexed).toBe(0) + expect(signal?.aborted).toBe(true) +}) + +test("a pre-aborted index leaves the existing cache available without calling the connector", async () => { + const cache = createCache() + await cache.indexWarehouse("warehouse", "other", connector()) + const source = connector() + const result = await cache.indexWarehouse("warehouse", "other", source, { signal: AbortSignal.abort() }) + expect(result.partial).toBe(true) + expect(result.tables_indexed).toBe(0) + expect(source.listSchemas).not.toHaveBeenCalled() + expect(cache.search("orders").tables).toHaveLength(1) +}) + +test("keeps the previous snapshot readable until replacement metadata is collected", async () => { + const cache = createCache() + await cache.indexWarehouse("warehouse", "other", connector()) + const source = connector() + const started = Promise.withResolvers() + const pending = Promise.withResolvers>>() + source.listTables = async () => [{ name: "customers", type: "table" }] + source.describeTable = async () => { + started.resolve() + return pending.promise + } + const indexing = cache.indexWarehouse("warehouse", "other", source) + await started.promise + try { + expect(cache.search("orders").tables).toHaveLength(1) + } finally { + pending.resolve([{ name: "id", data_type: "INTEGER", nullable: false }]) + await indexing + } + expect(cache.search("orders").tables).toHaveLength(0) + expect(cache.search("customers").tables).toHaveLength(1) +}) From adb3450cd64049d985678d5d3f226d23e670853f Mon Sep 17 00:00:00 2001 From: anandgupta42 Date: Thu, 15 Oct 2026 10:24:21 -0700 Subject: [PATCH 2/2] fix(schema): keep complete cache on partial index; mark first partial stale - a partial (timed-out, cancelled or soft-failed) run no longer replaces an existing snapshot - with no prior snapshot, the partial result is saved with an epoch `last_indexed` so freshness checks treat it as stale and re-index - tests cover both cases Co-Authored-By: Claude Opus 5.5 --- .../src/altimate/native/schema/cache.ts | 10 +++-- .../test/altimate/schema-indexing.test.ts | 39 +++++++++++++++++++ 2 files changed, 45 insertions(+), 4 deletions(-) diff --git a/packages/opencode/src/altimate/native/schema/cache.ts b/packages/opencode/src/altimate/native/schema/cache.ts index 8b66e4fc57..a71208eaf7 100644 --- a/packages/opencode/src/altimate/native/schema/cache.ts +++ b/packages/opencode/src/altimate/native/schema/cache.ts @@ -205,8 +205,10 @@ export class SchemaCache { let totalColumns = 0 // Gather metadata outside SQLite, then replace the snapshot in one transaction. - // A cancelled refresh with no completed tables leaves the previous cache intact. - if (!note || tables.size > 0) { + // A partial run never replaces an existing snapshot. With no prior snapshot it is saved + // with an epoch timestamp so every freshness check treats it as stale and re-indexes. + const hasPrior = Boolean(this.db.prepare("SELECT 1 FROM warehouses WHERE name = ?").get(warehouseName)) + if (!note || (tables.size > 0 && !hasPrior)) { const insertTable = this.db.prepare( `INSERT INTO tables_cache (warehouse, database_name, schema_name, table_name, table_type, search_text) @@ -238,10 +240,10 @@ export class SchemaCache { `INSERT OR REPLACE INTO warehouses (name, type, last_indexed, databases_count, schemas_count, tables_count, columns_count) VALUES (?, ?, ?, ?, ?, ?, ?)`, - ).run(warehouseName, warehouseType, now, 0, schemas.size, tables.size, totalColumns) + ).run(warehouseName, warehouseType, note ? new Date(0).toISOString() : now, 0, schemas.size, tables.size, totalColumns) })() } else { - note += " No cached metadata was replaced." + note += hasPrior ? " The previous cached metadata was kept." : " No cached metadata was replaced." } return { diff --git a/packages/opencode/test/altimate/schema-indexing.test.ts b/packages/opencode/test/altimate/schema-indexing.test.ts index 5b8a23a67b..1e0b32fcda 100644 --- a/packages/opencode/test/altimate/schema-indexing.test.ts +++ b/packages/opencode/test/altimate/schema-indexing.test.ts @@ -152,3 +152,42 @@ test("keeps the previous snapshot readable until replacement metadata is collect expect(cache.search("orders").tables).toHaveLength(0) expect(cache.search("customers").tables).toHaveLength(1) }) + +function abortedAfterFirstTable() { + const source = connector() + const controller = new AbortController() + source.listTables = async () => ["first", "second"].map((name) => ({ name, type: "table" })) + source.describeTable = mock(async (_schema: string, table: string) => { + if (table === "second") { + controller.abort() + return new Promise(() => {}) + } + return [{ name: "id", data_type: "INTEGER", nullable: false }] + }) + return { source, signal: controller.signal } +} + +test("a partial run never replaces an existing complete cache", async () => { + const cache = createCache() + const full = connector() + full.listTables = async () => ["first", "second", "third"].map((name) => ({ name, type: "table" })) + const complete = await cache.indexWarehouse("warehouse", "other", full) + expect(complete.partial).toBeUndefined() + const before = cache.cacheStatus().warehouses[0].last_indexed + + const { source, signal } = abortedAfterFirstTable() + const result = await cache.indexWarehouse("warehouse", "other", source, { signal }) + expect(result.partial).toBe(true) + expect(cache.cacheStatus().total_tables).toBe(3) + expect(cache.cacheStatus().warehouses[0].last_indexed).toBe(before) +}) + +test("a partial first index is saved but recorded as stale", async () => { + const cache = createCache() + const { source, signal } = abortedAfterFirstTable() + const result = await cache.indexWarehouse("warehouse", "other", source, { signal }) + expect(result.partial).toBe(true) + expect(cache.cacheStatus().total_tables).toBe(1) + const indexed = cache.cacheStatus().warehouses[0].last_indexed! + expect(Date.now() - new Date(indexed).getTime()).toBeGreaterThan(365 * 24 * 3600 * 1000) +})