feat(completion): add scoped metadata assistant

apps/desktop/src/components/editor/QueryEditor.vue

@@ -46,7 +46,7 @@ import * as api from "@/lib/api";
 import { areSqlSemanticDiagnosticsEqual, buildSqlParserErrorDiagnostic, buildSqlSemanticDiagnostics, shouldRunSqlSemanticDiagnostics, type SqlSemanticDiagnostic } from "@/lib/sqlSemanticDiagnostics";
 import { buildRedisSyntaxDiagnostics, shouldRunRedisDiagnostics } from "@/lib/redisSyntaxDiagnostics";
 import { buildRedisCompletionItemsFromContext, getRedisCompletionContext, getRedisCompletionResultValidFor, shouldAutoOpenRedisCompletion, takesKeyArgument, type RedisCompletionItem } from "@/lib/redisCompletion";
-import type { SqlCompletionColumn, SqlCompletionForeignKey, SqlCompletionItem, SqlCompletionObject } from "@/lib/sqlCompletion";
+import type { SqlCompletionColumn, SqlCompletionForeignKey, SqlCompletionItem, SqlCompletionObject, SqlCompletionTable } from "@/lib/sqlCompletion";
 import type { DatabaseType, SqlReferenceAnalysis, SqlTableReference, SqlTextSpan } from "@/types/database";
 
 const props = defineProps<{
@@ -65,6 +65,8 @@ const props = defineProps<{
   initialSelection?: { anchor: number; head: number };
 }>();
 
+const COMPLETION_REMOTE_LATENCY_BUDGET_MS = 120;
+
 const emit = defineEmits<{
   "update:modelValue": [value: string];
   selectionChange: [value: string];
@@ -1490,6 +1492,20 @@ function mergeCompletionTables(existing: Array<{ name: string; schema?: string;
   return merged;
 }
 
+function withCompletionLatencyBudget<T>(remote: Promise<T>, local: T): Promise<T> {
+  return Promise.race([remote, new Promise<T>((resolve) => setTimeout(() => resolve(local), COMPLETION_REMOTE_LATENCY_BUDGET_MS))]);
+}
+
+function listCompletionTablesWithLatencyBudget(connectionId: string, database: string, filter: string, limit: number, schema?: string): Promise<SqlCompletionTable[]> {
+  const local = connectionStore.lookupLocalCompletionTables(connectionId, database, filter, limit, schema);
+  const remote = connectionStore.listCompletionTables(connectionId, database, filter, limit, schema).then((tables) => {
+    cachedTables = mergeCompletionTables(cachedTables, tables);
+    return tables;
+  });
+  if (local.length === 0) return remote;
+  return withCompletionLatencyBudget(remote, local);
+}
+
 async function performAsyncCompletionWithResult(epoch: number, completionContext: ReturnType<typeof getSqlCompletionContext>, fullDoc: string, position: number) {
   const localOnlyMetadata = usesLocalOnlyCompletionMetadata();
   const onDemandOnlyColumns = usesOnDemandOnlyCompletionColumns();
@@ -1526,7 +1542,7 @@ async function performAsyncCompletionWithResult(epoch: number, completionContext
   let tables = shouldLoadTables
     ? localOnlyMetadata
       ? connectionStore.lookupLocalCompletionTables(props.connectionId!, tableLookupDatabase, tableLookupFilter, MAX_COMPLETION_TABLES, tableLookupSchema)
-      : await connectionStore.listCompletionTables(props.connectionId!, tableLookupDatabase, tableLookupFilter, MAX_COMPLETION_TABLES, tableLookupSchema)
+      : await listCompletionTablesWithLatencyBudget(props.connectionId!, tableLookupDatabase, tableLookupFilter, MAX_COMPLETION_TABLES, tableLookupSchema)
     : cachedTables;
   if (epoch !== completionEpoch) return null;
 
@@ -1568,7 +1584,7 @@ async function performAsyncCompletionWithResult(epoch: number, completionContext
   if (completionContext.qualifier && !qualifierDatabase && !isReferencedTableQualifier(completionContext) && tables.length === 0 && (completionContext.suggestTables || completionContext.exclusiveColumnSuggestions)) {
     const schemaTables = localOnlyMetadata
       ? connectionStore.lookupLocalCompletionTables(props.connectionId!, props.database!, completionContext.prefix, MAX_COMPLETION_TABLES, completionContext.qualifier)
-      : await connectionStore.listCompletionTables(props.connectionId!, props.database!, completionContext.prefix, MAX_COMPLETION_TABLES, completionContext.qualifier);
+      : await listCompletionTablesWithLatencyBudget(props.connectionId!, props.database!, completionContext.prefix, MAX_COMPLETION_TABLES, completionContext.qualifier);
     if (schemaTables.length > 0) {
       tables = schemaTables;
       qualifierIsSchema = true;

apps/desktop/src/lib/__tests__/sqlCompletion.context.spec.ts

@@ -27,3 +27,100 @@ describe("sqlCompletion quoted schema qualifiers", () => {
     expect(items.some((item) => item.label === "shipments" && item.type === "table")).toBe(true);
   });
 });
+
+describe("sqlCompletion scoped context classification", () => {
+  it("classifies JOIN table contexts", () => {
+    const sql = "SELECT * FROM users u JOIN ";
+    const context = getSqlCompletionContext(sql, sql.length);
+
+    expect(context.contextKind).toBe("join");
+    expect(context.suggestTables).toBe(true);
+    expect(context.exclusiveTableSuggestions).toBe(true);
+  });
+
+  it("classifies alias-qualified column contexts", () => {
+    const sql = "SELECT * FROM users u WHERE u.";
+    const context = getSqlCompletionContext(sql, sql.length);
+
+    expect(context.contextKind).toBe("alias_column");
+    expect(context.qualifier).toBe("u");
+    expect(context.suggestColumns).toBe(true);
+  });
+
+  it("classifies CALL routine contexts", () => {
+    const sql = "CALL usp_";
+    const context = getSqlCompletionContext(sql, sql.length);
+
+    expect(context.contextKind).toBe("exec");
+    expect(context.suggestRoutines).toBe(true);
+    expect(context.exclusiveRoutineSuggestions).toBe(true);
+  });
+
+  it("classifies INSERT column-list contexts", () => {
+    const sql = "INSERT INTO dbo.Users (";
+    const context = getSqlCompletionContext(sql, sql.length);
+
+    expect(context.contextKind).toBe("column");
+    expect(context.insertSchema).toBe("dbo");
+    expect(context.insertTable).toBe("users");
+    expect(context.exclusiveColumnSuggestions).toBe(true);
+  });
+
+  it("classifies UPDATE SET column contexts", () => {
+    const sql = "UPDATE dbo.Users SET ";
+    const context = getSqlCompletionContext(sql, sql.length);
+
+    expect(context.contextKind).toBe("column");
+    expect(context.updateTarget).toEqual({ schema: "dbo", table: "Users" });
+    expect(context.suggestColumns).toBe(true);
+  });
+
+  it("extracts statement-local table aliases", () => {
+    const sql = "SELECT * FROM dbo.Users u JOIN Orders AS o ON o.user_id = u.id WHERE u.";
+    const context = getSqlCompletionContext(sql, sql.length);
+
+    expect(context.referencedTables).toEqual(expect.arrayContaining([expect.objectContaining({ schema: "dbo", name: "Users", alias: "u" }), expect.objectContaining({ name: "Orders", alias: "o" })]));
+  });
+
+  it("exposes CTEs as table-like referenced tables", () => {
+    const sql = "WITH recent_orders(id, total) AS (SELECT id, total FROM orders) SELECT * FROM recent_orders ro WHERE ro.";
+    const context = getSqlCompletionContext(sql, sql.length);
+
+    expect(context.referencedTables).toEqual(expect.arrayContaining([expect.objectContaining({ name: "recent_orders", columns: ["id", "total"] }), expect.objectContaining({ name: "recent_orders", alias: "ro" })]));
+  });
+
+  it("extracts subquery aliases and projected columns", () => {
+    const sql = "SELECT * FROM (SELECT id, name AS user_name FROM users) sq WHERE sq.";
+    const context = getSqlCompletionContext(sql, sql.length);
+
+    expect(context.referencedTables).toEqual(expect.arrayContaining([expect.objectContaining({ name: "sq", alias: "sq", columns: ["id", "user_name"] })]));
+  });
+});
+
+describe("sqlCompletion scoped metadata ranking", () => {
+  it("ranks exact and prefix table matches ahead of contains/fuzzy matches", () => {
+    const sql = "SELECT * FROM Temp";
+    const items = buildSqlCompletionItems(sql, sql.length, {
+      dialect: "sqlserver",
+      tables: [
+        { name: "ArchiveTempTable", schema: "dbo", type: "table" },
+        { name: "TempAudit", schema: "dbo", type: "table" },
+        { name: "Temp", schema: "dbo", type: "table" },
+        { name: "Template", schema: "dbo", type: "table" },
+      ],
+      columnsByTable: new Map(),
+    }).filter((item) => item.type === "table");
+
+    expect(items.map((item) => item.label).slice(0, 3)).toEqual(["Temp", "Template", "TempAudit"]);
+    expect(items.some((item) => item.label === "ArchiveTempTable")).toBe(true);
+  });
+
+  it("keeps large table catalogs bounded", () => {
+    const tables = Array.from({ length: 500 }, (_, index) => ({ name: `TempTable_${String(index).padStart(3, "0")}`, schema: "dbo", type: "table" as const }));
+    const sql = "SELECT * FROM Temp";
+    const items = buildSqlCompletionItems(sql, sql.length, { dialect: "sqlserver", tables, columnsByTable: new Map() }).filter((item) => item.type === "table");
+
+    expect(items.length).toBeLessThanOrEqual(200);
+    expect(items[0]?.label).toBe("TempTable_000");
+  });
+});

apps/desktop/src/lib/api.ts

@@ -118,6 +118,7 @@ export const getTableComment = forward("getTableComment");
 export const listObjects = forward("listObjects");
 export const listObjectStatistics = forward("listObjectStatistics");
 export const listCompletionObjects = forward("listCompletionObjects");
+export const completionAssistantSearch = forward("completionAssistantSearch");
 export const getObjectSource = forward("getObjectSource");
 export const getColumns = forward("getColumns");
 export const listIndexes = forward("listIndexes");

apps/desktop/src/lib/http.ts

@@ -4,6 +4,8 @@ import type {
   LinkedServerInfo,
   TableInfo,
   ObjectInfo,
+  CompletionAssistantRequest,
+  CompletionAssistantResponse,
   ObjectStatistics,
   ObjectSource,
   ObjectSourceKind,
@@ -482,6 +484,10 @@ export async function listCompletionObjects(connectionId: string, database: stri
   return get(`/api/schema/completion-objects?${qs({ connection_id: connectionId, database, schema })}`);
 }
 
+export async function completionAssistantSearch(request: CompletionAssistantRequest): Promise<CompletionAssistantResponse> {
+  return post("/api/schema/completion-assistant", request);
+}
+
 export async function getObjectSource(connectionId: string, database: string, schema: string, name: string, objectType: ObjectSourceKind): Promise<ObjectSource> {
   return get(`/api/schema/object-source?${qs({ connection_id: connectionId, database, schema, table: name, object_type: objectType })}`);
 }

apps/desktop/src/lib/sqlCompletion.ts

@@ -1063,6 +1063,8 @@ export interface SqlCompletionReferencedTable {
 
 export type SqlStatementKind = "select" | "insert" | "update" | "delete" | "create" | "alter" | "drop" | "unknown";
 
+export type SqlCompletionContextKind = "table" | "schema" | "catalog" | "routine" | "column" | "alias_column" | "insert_target" | "update_target" | "exec" | "join" | "keyword";
+
 export interface SqlCompletionContext {
   prefix: string;
   qualifier?: string;
@@ -1090,6 +1092,7 @@ export interface SqlCompletionContext {
   updateTarget?: { table: string; schema?: string };
   deleteTarget?: { table: string; schema?: string };
   oracleTableFunctionContext?: boolean;
+  contextKind: SqlCompletionContextKind;
 }
 
 export interface SqlFunctionSignatureHelp {
@@ -1553,6 +1556,19 @@ export function getSqlCompletionContext(sql: string, cursor: number): SqlComplet
 
   const statementKind = detectStatementKind(beforeCursor || fullStatement);
   const preferredKeywords = preferredKeywordsForCompletion(updateInfo, deleteInfo);
+  const contextKind = detectCompletionContextKind({
+    qualifier,
+    exclusiveTableSuggestions,
+    exclusiveColumnSuggestions,
+    insertInfo,
+    updateInfo,
+    inCallRoutineContext,
+    oracleTableFunctionContext,
+    afterTableTrigger,
+    lastWord,
+    suggestColumns,
+    suggestRoutines,
+  });
 
   return {
     prefix,
@@ -1581,9 +1597,33 @@ export function getSqlCompletionContext(sql: string, cursor: number): SqlComplet
     updateTarget: updateInfo?.target,
     deleteTarget: deleteInfo?.target,
     oracleTableFunctionContext,
+    contextKind,
   };
 }
 
+function detectCompletionContextKind(options: {
+  qualifier?: string;
+  exclusiveTableSuggestions: boolean;
+  exclusiveColumnSuggestions: boolean;
+  insertInfo: ReturnType<typeof detectInsertColumnListContext>;
+  updateInfo: ReturnType<typeof detectUpdateCompletionContext>;
+  inCallRoutineContext: boolean;
+  oracleTableFunctionContext: boolean;
+  afterTableTrigger: boolean;
+  lastWord: string;
+  suggestColumns: boolean;
+  suggestRoutines: boolean;
+}): SqlCompletionContextKind {
+  if (options.insertInfo) return "column";
+  if (options.updateInfo?.inSetClause) return "column";
+  if (options.inCallRoutineContext) return "exec";
+  if (options.qualifier && options.exclusiveColumnSuggestions) return "alias_column";
+  if (options.oracleTableFunctionContext || options.suggestRoutines) return "routine";
+  if (options.exclusiveTableSuggestions || options.afterTableTrigger) return options.lastWord === "join" ? "join" : "table";
+  if (options.suggestColumns) return options.qualifier ? "alias_column" : "column";
+  return "keyword";
+}
+
 function parseTrailingIdentifierContext(input: string): { start: number; prefix: string; qualifier?: string; qualifierParts?: string[] } | null {
   if (/\s$/.test(input)) return null;
   let i = input.length - 1;

apps/desktop/src/lib/tauri.ts

@@ -6,6 +6,8 @@ import type {
   LinkedServerInfo,
   TableInfo,
   ObjectInfo,
+  CompletionAssistantRequest,
+  CompletionAssistantResponse,
   ObjectStatistics,
   ObjectSource,
   ObjectSourceKind,
@@ -533,6 +535,10 @@ export async function listCompletionObjects(connectionId: string, database: stri
   return invoke("list_completion_objects", { connectionId, database, schema });
 }
 
+export async function completionAssistantSearch(request: CompletionAssistantRequest): Promise<CompletionAssistantResponse> {
+  return invoke("completion_assistant_search", { request });
+}
+
 export async function getObjectSource(connectionId: string, database: string, schema: string, name: string, objectType: ObjectSourceKind): Promise<ObjectSource> {
   return invoke("get_object_source", { connectionId, database, schema, name, objectType });
 }

apps/desktop/src/stores/__tests__/connectionStore.completion.spec.ts

@@ -0,0 +1,86 @@
+import { createPinia, setActivePinia } from "pinia";
+import { beforeEach, describe, expect, it, vi } from "vitest";
+import type { ConnectionConfig } from "@/types/database";
+
+function installLocalStorage() {
+  const data = new Map<string, string>();
+  vi.stubGlobal("localStorage", {
+    getItem: vi.fn((key: string) => data.get(key) ?? null),
+    setItem: vi.fn((key: string, value: string) => data.set(key, value)),
+    removeItem: vi.fn((key: string) => data.delete(key)),
+  });
+}
+
+function postgresConnection(): ConnectionConfig {
+  return {
+    id: "pg-1",
+    name: "Postgres",
+    db_type: "postgres",
+    host: "127.0.0.1",
+    port: 5432,
+    username: "postgres",
+    password: "",
+    database: "app",
+    read_only: false,
+  } as ConnectionConfig;
+}
+
+describe("connectionStore completion assistant", () => {
+  beforeEach(() => {
+    vi.resetModules();
+    vi.unstubAllGlobals();
+    installLocalStorage();
+    setActivePinia(createPinia());
+  });
+
+  it("deduplicates in-flight assistant table requests", async () => {
+    const completionAssistantSearch = vi.fn().mockResolvedValue({
+      candidates: [{ name: "accounts", kind: "table", schema: "public" }],
+      incomplete: false,
+      fallback_used: false,
+    });
+
+    vi.doMock("@/lib/tauriRuntime", () => ({ isTauriRuntime: () => false }));
+    vi.doMock("@/lib/api", () => ({
+      checkConnectionHealth: vi.fn().mockResolvedValue(undefined),
+      completionAssistantSearch,
+      listSchemas: vi.fn().mockResolvedValue(["public"]),
+      listTables: vi.fn().mockResolvedValue([]),
+    }));
+
+    const { useConnectionStore } = await import("@/stores/connectionStore");
+    const store = useConnectionStore();
+    store.connections = [postgresConnection()];
+    store.connectedIds.add("pg-1");
+
+    const [first, second] = await Promise.all([store.listCompletionTables("pg-1", "app", "acc", 20, "public"), store.listCompletionTables("pg-1", "app", "acc", 20, "public")]);
+
+    expect(completionAssistantSearch).toHaveBeenCalledTimes(1);
+    expect(first).toEqual(second);
+    expect(first[0]).toMatchObject({ name: "accounts", schema: "public", type: "table" });
+  });
+
+  it("returns fallback metadata when assistant table search fails", async () => {
+    const completionAssistantSearch = vi.fn().mockRejectedValue(new Error("assistant unavailable"));
+    const listTables = vi.fn().mockResolvedValue([{ name: "accounts", table_type: "BASE TABLE", comment: null }]);
+
+    vi.doMock("@/lib/tauriRuntime", () => ({ isTauriRuntime: () => false }));
+    vi.doMock("@/lib/api", () => ({
+      checkConnectionHealth: vi.fn().mockResolvedValue(undefined),
+      completionAssistantSearch,
+      listSchemas: vi.fn().mockResolvedValue(["public"]),
+      listTables,
+    }));
+
+    const { useConnectionStore } = await import("@/stores/connectionStore");
+    const store = useConnectionStore();
+    store.connections = [postgresConnection()];
+    store.connectedIds.add("pg-1");
+
+    const tables = await store.listCompletionTables("pg-1", "app", "acc", 20, "public");
+
+    expect(completionAssistantSearch).toHaveBeenCalledTimes(1);
+    expect(listTables).toHaveBeenCalledWith("pg-1", "app", "public", "acc", 20);
+    expect(tables).toEqual([{ name: "accounts", schema: "public", type: "table" }]);
+  });
+});

apps/desktop/src/stores/connectionStore.ts

@@ -1,7 +1,7 @@
 import { defineStore } from "pinia";
 import { uuid } from "@/lib/utils";
 import { ref, computed, watch } from "vue";
-import type { ColumnInfo, ConnectionConfig, ForeignKeyInfo, ObjectInfo, SidebarLayout, TableInfo, TreeNode } from "@/types/database";
+import type { ColumnInfo, CompletionAssistantCandidate, CompletionAssistantObjectKind, CompletionAssistantRequest, ConnectionConfig, ForeignKeyInfo, ObjectInfo, SidebarLayout, TableInfo, TreeNode } from "@/types/database";
 import { applyPinnedTreeNodeState, updatePinnedTreeNodeInPlace } from "@/lib/pinnedItems";
 import {
   reconcileLayout,
@@ -2143,6 +2143,87 @@ export const useConnectionStore = defineStore("connection", () => {
     return promise;
   }
 
+  function completionAssistantRequestKey(request: CompletionAssistantRequest): string {
+    return JSON.stringify({
+      connection_id: request.connection_id,
+      database: request.database,
+      schema: request.schema ?? "",
+      object_kinds: [...(request.object_kinds ?? [])].sort(),
+      mask: request.mask ?? "",
+      case_sensitive: !!request.case_sensitive,
+      global_search: !!request.global_search,
+      max_results: request.max_results ?? null,
+      search_in_comments: !!request.search_in_comments,
+      search_in_definitions: !!request.search_in_definitions,
+      parent_schema: request.parent_schema ?? "",
+      parent_name: request.parent_name ?? "",
+      match_mode: request.match_mode ?? "prefix",
+    });
+  }
+
+  async function completionAssistantSearch(request: CompletionAssistantRequest) {
+    return withCompletionInFlight(`assistant:${completionAssistantRequestKey(request)}`, async () => {
+      await ensureConnected(request.connection_id);
+      return api.completionAssistantSearch(request);
+    });
+  }
+
+  function completionAssistantTables(candidates: CompletionAssistantCandidate[]): SqlCompletionTable[] {
+    return candidates
+      .filter((candidate) => candidate.kind === "table" || candidate.kind === "view")
+      .map((candidate) => ({
+        name: candidate.name,
+        schema: candidate.schema ?? undefined,
+        type: candidate.kind === "view" ? ("view" as const) : ("table" as const),
+      }));
+  }
+
+  function completionAssistantColumns(candidates: CompletionAssistantCandidate[], table: string, schema?: string): SqlCompletionColumn[] {
+    return candidates
+      .filter((candidate) => candidate.kind === "column")
+      .map((candidate) => ({
+        name: candidate.name,
+        table: candidate.parent_name ?? table,
+        schema: candidate.parent_schema ?? candidate.schema ?? schema,
+        dataType: candidate.data_type ?? undefined,
+        comment: candidate.comment ?? null,
+      }));
+  }
+
+  async function listCompletionAssistantTables(connectionId: string, database: string, filter: string, limit?: number, schema?: string): Promise<SqlCompletionTable[]> {
+    const objectKinds: CompletionAssistantObjectKind[] = ["table", "view"];
+    const response = await completionAssistantSearch({
+      connection_id: connectionId,
+      database,
+      schema: schema ?? null,
+      object_kinds: objectKinds,
+      mask: filter.trim(),
+      max_results: limit ?? 200,
+      parent_schema: schema ?? null,
+      match_mode: "prefix",
+    });
+    const tables = completionAssistantTables(response.candidates);
+    indexCompletionTables(connectionId, database, schema, tables);
+    return tables;
+  }
+
+  async function listCompletionAssistantColumns(connectionId: string, database: string, table: string, schema?: string): Promise<SqlCompletionColumn[]> {
+    const response = await completionAssistantSearch({
+      connection_id: connectionId,
+      database,
+      schema: schema ?? null,
+      object_kinds: ["column"],
+      mask: "",
+      max_results: 500,
+      parent_schema: schema ?? null,
+      parent_name: table,
+      match_mode: "prefix",
+    });
+    const columns = completionAssistantColumns(response.candidates, table, schema);
+    if (columns.length > 0) indexCompletionColumns(connectionId, database, table, schema, columns);
+    return columns;
+  }
+
   function completionNameSegments(name: string): string[] {
     return name
       .replace(/([a-z0-9])([A-Z])/g, "$1 $2")
@@ -2432,53 +2513,36 @@ export const useConnectionStore = defineStore("connection", () => {
       await ensureConnected(connectionId);
 
       if (isSchemaAwareDatabase(connectionId)) {
-        const schemas = schema ? [schema] : await listCompletionSchemas(connectionId, database);
         if (normalizedFilter || limit) {
-          const batchSize = 5;
-          const results: SqlCompletionTable[] = [];
-          const maxResults = limit ?? Infinity;
-          for (let i = 0; i < schemas.length && results.length < maxResults; i += batchSize) {
-            const batch = schemas.slice(i, i + batchSize);
-            const batchResults = await Promise.all(
-              batch.map(async (s) => {
-                try {
-                  const tables = await api.listTables(connectionId, database, s, normalizedFilter, limit);
-                  return tables.map((table) => ({
-                    name: table.name,
-                    schema: s,
-                    type: table.table_type === "VIEW" || table.table_type === "MATERIALIZED_VIEW" ? ("view" as const) : ("table" as const),
-                  })) as SqlCompletionTable[];
-                } catch {
-                  return [] as SqlCompletionTable[];
-                }
-              }),
-            );
-            for (const group of batchResults) {
-              results.push(...group);
-              indexCompletionTables(connectionId, database, undefined, group);
+          let results: SqlCompletionTable[] = [];
+          try {
+            results = await listCompletionAssistantTables(connectionId, database, normalizedFilter, limit, schema);
+          } catch {
+            if (schema) {
+              const tables = await api.listTables(connectionId, database, schema, normalizedFilter, limit);
+              results = tables.map((table) => ({
+                name: table.name,
+                schema,
+                type: table.table_type === "VIEW" || table.table_type === "MATERIALIZED_VIEW" ? ("view" as const) : ("table" as const),
+              }));
+            } else {
+              results = lookupLocalCompletionTables(connectionId, database, normalizedFilter, limit);
             }
           }
           if (results.length === 0 && relaxedFilter) {
-            for (let i = 0; i < schemas.length && results.length < maxResults; i += batchSize) {
-              const batch = schemas.slice(i, i + batchSize);
-              const batchResults = await Promise.all(
-                batch.map(async (s) => {
-                  try {
-                    const tables = await api.listTables(connectionId, database, s, relaxedFilter, expandedCompletionLimit(limit));
-                    return tables.map((table) => ({
-                      name: table.name,
-                      schema: s,
-                      type: table.table_type === "VIEW" || table.table_type === "MATERIALIZED_VIEW" ? ("view" as const) : ("table" as const),
-                    })) as SqlCompletionTable[];
-                  } catch {
-                    return [] as SqlCompletionTable[];
-                  }
-                }),
-              );
-              for (const group of batchResults) {
-                results.push(...group);
-                indexCompletionTables(connectionId, database, undefined, group);
+            if (schema) {
+              try {
+                const tables = await api.listTables(connectionId, database, schema, relaxedFilter, expandedCompletionLimit(limit));
+                results = tables.map((table) => ({
+                  name: table.name,
+                  schema,
+                  type: table.table_type === "VIEW" || table.table_type === "MATERIALIZED_VIEW" ? ("view" as const) : ("table" as const),
+                }));
+              } catch {
+                results = [];
               }
+            } else {
+              results = lookupLocalCompletionTables(connectionId, database, relaxedFilter, expandedCompletionLimit(limit));
             }
           }
           const limitedTables = limit ? dedupeCompletionTables(results).slice(0, limit) : results;
@@ -2488,21 +2552,16 @@ export const useConnectionStore = defineStore("connection", () => {
           return completionTablesCache.value[cacheKey];
         }
 
-        const tableGroups = await Promise.all(
-          schemas.map(async (schema) => {
-            try {
-              const tables = await api.listTables(connectionId, database, schema);
-              return tables.map((table) => ({
-                name: table.name,
-                schema,
-                type: table.table_type === "VIEW" || table.table_type === "MATERIALIZED_VIEW" ? ("view" as const) : ("table" as const),
-              }));
-            } catch {
-              return [];
-            }
-          }),
-        );
-        completionTablesCache.value[cacheKey] = tableGroups.flat();
+        if (schema) {
+          const tables = await api.listTables(connectionId, database, schema);
+          completionTablesCache.value[cacheKey] = tables.map((table) => ({
+            name: table.name,
+            schema,
+            type: table.table_type === "VIEW" || table.table_type === "MATERIALIZED_VIEW" ? ("view" as const) : ("table" as const),
+          }));
+        } else {
+          completionTablesCache.value[cacheKey] = lookupLocalCompletionTables(connectionId, database, normalizedFilter, limit);
+        }
         indexCompletionTables(connectionId, database, undefined, completionTablesCache.value[cacheKey]);
         evictOldestCacheEntries(completionTablesCache.value, COMPLETION_CACHE_MAX);
         return completionTablesCache.value[cacheKey];
@@ -2633,6 +2692,27 @@ export const useConnectionStore = defineStore("connection", () => {
     if (!completionColumnsCache.value[cacheKey]) {
       await withCompletionInFlight(`${cacheKey}:columns`, async () => {
         await ensureConnected(connectionId);
+        try {
+          const assistantColumns = await listCompletionAssistantColumns(connectionId, database, table, schema);
+          if (assistantColumns.length > 0) {
+            completionColumnsCache.value[cacheKey] = assistantColumns.map((column) => ({
+              name: column.name,
+              data_type: column.dataType ?? "",
+              is_nullable: column.isNullable ?? true,
+              column_default: null,
+              is_primary_key: false,
+              extra: null,
+              comment: column.comment ?? null,
+              numeric_precision: null,
+              numeric_scale: null,
+              character_maximum_length: null,
+            }));
+            evictOldestCacheEntries(completionColumnsCache.value, COMPLETION_CACHE_MAX);
+            return;
+          }
+        } catch {
+          // Fall back to the existing metadata path below.
+        }
         const querySchema = metadataQuerySchema(connectionId, database, schema);
         completionColumnsCache.value[cacheKey] = await api.getColumns(connectionId, database, querySchema, table);
         evictOldestCacheEntries(completionColumnsCache.value, COMPLETION_CACHE_MAX);

apps/desktop/src/types/database.ts

@@ -65,6 +65,45 @@ export interface SqlSnippet {
   body: string;
 }
 
+export type CompletionAssistantObjectKind = "database" | "schema" | "table" | "view" | "routine" | "procedure" | "function" | "column";
+
+export type CompletionAssistantCandidateKind = "database" | "schema" | "table" | "view" | "procedure" | "function" | "column" | "object";
+
+export type CompletionAssistantMatchMode = "prefix" | "contains";
+
+export interface CompletionAssistantRequest {
+  connection_id: string;
+  database: string;
+  schema?: string | null;
+  object_kinds?: CompletionAssistantObjectKind[];
+  mask?: string;
+  case_sensitive?: boolean;
+  global_search?: boolean;
+  max_results?: number | null;
+  search_in_comments?: boolean;
+  search_in_definitions?: boolean;
+  parent_schema?: string | null;
+  parent_name?: string | null;
+  match_mode?: CompletionAssistantMatchMode | null;
+}
+
+export interface CompletionAssistantCandidate {
+  name: string;
+  kind: CompletionAssistantCandidateKind;
+  database?: string | null;
+  schema?: string | null;
+  parent_schema?: string | null;
+  parent_name?: string | null;
+  comment?: string | null;
+  data_type?: string | null;
+}
+
+export interface CompletionAssistantResponse {
+  candidates: CompletionAssistantCandidate[];
+  incomplete: boolean;
+  fallback_used: boolean;
+}
+
 export interface ConnectionConfig {
   id: string;
   name: string;