From 59b6de2cfb3f4062ac232233247742b2c5e0a0a9 Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Fri, 2 Oct 2026 17:58:55 +0000 Subject: [PATCH] fix(snowflake): resolve database and schema completion namespaces Closes #10953 --- .../useQueryEditorCompletion.spec.ts | 177 +++++++++++++++++- .../editor/useQueryEditorCompletion.ts | 3 + .../database/databaseFeatureSupport.spec.ts | 8 +- .../__tests__/sql/semantic/dialect.spec.ts | 11 ++ .../sql/sqlCompletionLookupTarget.spec.ts | 28 +++ .../lib/database/databaseFeatureSupport.ts | 2 +- apps/desktop/src/lib/sql/semantic/dialect.ts | 10 + apps/desktop/src/lib/sql/sqlCompletion.ts | 20 +- .../src/lib/sql/sqlCompletionLookupTarget.ts | 3 +- .../connectionStore.snowflake.spec.ts | 133 +++++++++++++ .../assets/database-drivers.manifest.json | 4 +- plugins/connection-types/snowflake.yaml | 2 + 12 files changed, 378 insertions(+), 23 deletions(-) create mode 100644 apps/desktop/src/stores/__tests__/connectionStore.snowflake.spec.ts diff --git a/apps/desktop/src/components/editor/__tests__/useQueryEditorCompletion.spec.ts b/apps/desktop/src/components/editor/__tests__/useQueryEditorCompletion.spec.ts index 259c96bb3e..bc5b8d08f3 100644 --- a/apps/desktop/src/components/editor/__tests__/useQueryEditorCompletion.spec.ts +++ b/apps/desktop/src/components/editor/__tests__/useQueryEditorCompletion.spec.ts @@ -9,7 +9,7 @@ import { useQueryEditorCompletion } from "../useQueryEditorCompletion"; import { useQueryEditorCompletionMetadata } from "../useQueryEditorCompletionMetadata"; import type { QueryEditorProps } from "../queryEditorTypes"; import type { RedisCommandDocumentation } from "@/lib/redis/redisCommandDocs"; -import type { SqlCompletionTable } from "@/lib/sql/sqlCompletion"; +import type { SqlCompletionColumn, SqlCompletionTable } from "@/lib/sql/sqlCompletion"; import { analyzeSqlCompletion, type SqlCompletionAnalysisResult } from "@/lib/sql/sqlCompletionAnalysis"; vi.mock("@/stores/connectionStore", () => ({ COMPLETION_METADATA_CONCURRENCY: 4 })); @@ -32,7 +32,7 @@ afterEach(() => { vi.restoreAllMocks(); }); -function createHarness(overrides: Partial = {}, configureMetadata?: (metadata: ReturnType) => void) { +function createHarness(overrides: Partial = {}, configureMetadata?: (metadata: ReturnType) => void, semanticCompletionEnabled = false) { const props = reactive({ modelValue: "SEL", databaseType: "mysql", dialect: "mysql", connectionId: "connection", database: "demo", ...overrides }); const parent = document.createElement("div"); document.body.append(parent); @@ -44,14 +44,14 @@ function createHarness(overrides: Partial = {}, configureMetad lookupLocalCompletionColumns: vi.fn(() => []), lookupLocalCompletionColumnsByPrefix: vi.fn(() => []), lookupLocalCompletionObjects: vi.fn(() => []), - lookupLocalCompletionDatabases: vi.fn(() => []), - lookupLocalCompletionSchemas: vi.fn(() => []), + lookupLocalCompletionDatabases: vi.fn((): string[] => []), + lookupLocalCompletionSchemas: vi.fn((): string[] => []), lookupLocalCompletionForeignKeys: vi.fn(() => []), listCompletionTables: vi.fn(async (): Promise => []), listCompletionObjects: vi.fn(async () => []), - listCompletionColumns: vi.fn(async () => []), - listCompletionSchemas: vi.fn(async () => []), - listCompletionDatabases: vi.fn(async () => []), + listCompletionColumns: vi.fn(async (): Promise => []), + listCompletionSchemas: vi.fn(async (): Promise => []), + listCompletionDatabases: vi.fn(async (): Promise => []), refreshCompletionTables: vi.fn(async () => []), refreshCompletionColumns: vi.fn(async () => []), refreshCompletionSchemas: vi.fn(async () => []), @@ -66,7 +66,7 @@ function createHarness(overrides: Partial = {}, configureMetad }; const connectionStore = store as unknown as Options["connectionStore"]; const settings = reactive({ editorSettings: { completionTriggerMode: "positional", snippets: [], sqlFormatter: { keywordCase: "upper", functionCase: "upper" }, autoAliasTables: false, tableCompletionSchemaQualification: "collision", generateSqlQuoteIdentifiers: false } }); - const metadata = useQueryEditorCompletionMetadata({ props, view, connectionStore, sqlBehaviorDialect: () => props.dialect, remoteLatencyBudgetMs: 40, maxCompletionTables: 100, onDemandMinPrefix: 2, semanticCompletionEnabled: false }); + const metadata = useQueryEditorCompletionMetadata({ props, view, connectionStore, sqlBehaviorDialect: () => props.dialect, remoteLatencyBudgetMs: 40, maxCompletionTables: 100, onDemandMinPrefix: 2, semanticCompletionEnabled }); configureMetadata?.(metadata); const startCompletion = vi.fn(() => true); const runtime: Options["runtime"] = { codeMirrorStartCompletion: startCompletion, codeMirrorInsertCompletionText: insertCompletionText, codeMirrorSnippetCompletion: snippetCompletion, codeMirrorCompletionStatus: () => null, imeCompositionActive: false }; @@ -92,7 +92,7 @@ function createHarness(overrides: Partial = {}, configureMetad triggerDeferDelayMs: 25, maxCompletionTables: 100, onDemandTableLimit: 100, - semanticCompletionEnabled: false, + semanticCompletionEnabled, }); cleanups.push(() => { completion.clearDeferredCompletionTrigger(); @@ -104,6 +104,165 @@ function createHarness(overrides: Partial = {}, configureMetad return { props, currentView, view, store, settings, runtime, completion, provide, startCompletion }; } +describe.each([false, true])("Snowflake namespace completion (semantic=%s)", (semanticCompletionEnabled) => { + beforeEach(() => vi.useRealTimers()); + + function snowflakeHarness(sql: string, schema: string | undefined = "PUBLIC") { + const harness = createHarness({ databaseType: "snowflake", dialect: "sql", database: "DEFAULT_DB", schema, modelValue: sql }, undefined, semanticCompletionEnabled); + const catalogs: Record> = { + DEFAULT_DB: { PUBLIC: ["DEFAULT_TABLE"] }, + OTHER_DB: { PUBLIC: ["OTHER_TABLE"] }, + ADVENTUREWORKS2025: { PUBLIC: ["DIMCUSTOMER"] }, + Other_Db: { PUBLIC: ["QUOTED_TABLE"] }, + "My.DB": { "Sales Area": ["Dim.Customer"], 'A"B': ["ESCAPED_TABLE"] }, + }; + const databases = Object.keys(catalogs); + harness.store.listCompletionDatabases.mockResolvedValue(databases); + harness.store.listCompletionSchemas.mockImplementation(async (...args: unknown[]) => Object.keys(catalogs[String(args[1])] ?? {})); + harness.store.listCompletionTables.mockImplementation(async (...args: unknown[]) => (catalogs[String(args[1])]?.[String(args[4])] ?? []).map((name) => ({ name, schema: String(args[4]) }))); + return { ...harness, catalogs }; + } + + it.each(["DEFAULT_DB", "OTHER_DB", "ADVENTUREWORKS2025", "Adventureworks2025"])("lists schemas for explicit database %s with no selected schema", async (database) => { + const { provide, store, props } = snowflakeHarness(`SELECT * FROM ${database}.`); + props.schema = undefined; + const result = await provide(); + expect(result?.options.map((option) => option.label)).toContain("PUBLIC"); + expect(store.listCompletionSchemas).toHaveBeenCalledWith("connection", database.toUpperCase()); + expect(result?.options.map((option) => option.label)).toEqual(["PUBLIC"]); + }); + + it.each([ + ['"Other_Db"', "Other_Db", ["PUBLIC"]], + ['"My.DB"', "My.DB", ['A"B', "Sales Area"]], + ] as const)("keeps the exact quoted database %s when listing schemas", async (qualifier, database, schemas) => { + const { provide, store } = snowflakeHarness(`SELECT * FROM ${qualifier}.`); + expect((await provide())?.options.map((option) => option.label).sort()).toEqual(schemas); + expect(store.listCompletionSchemas).toHaveBeenCalledWith("connection", database); + }); + + it.each([ + ["DEFAULT_DB.PUBLIC", "DEFAULT_DB", "PUBLIC", "DEFAULT_TABLE"], + ["other_db.public", "OTHER_DB", "PUBLIC", "OTHER_TABLE"], + ["Adventureworks2025.Public", "ADVENTUREWORKS2025", "PUBLIC", "DIMCUSTOMER"], + ['"Other_Db".PUBLIC', "Other_Db", "PUBLIC", "QUOTED_TABLE"], + ['"My.DB"."Sales Area"', "My.DB", "Sales Area", "Dim.Customer"], + ['"My.DB"."A""B"', "My.DB", 'A"B', "ESCAPED_TABLE"], + ])("loads only tables in %s and preserves the typed prefix", async (qualifier, database, schema, table) => { + const sql = `SELECT * FROM ${qualifier}.`; + const { provide, store, currentView, settings } = snowflakeHarness(sql); + settings.editorSettings.tableCompletionSchemaQualification = "always"; + settings.editorSettings.generateSqlQuoteIdentifiers = true; + const result = await provide(); + expect(store.listCompletionTables.mock.calls.map((args) => (args as unknown[]).slice(0, 5))).toContainEqual(["connection", database, "", 100, schema]); + expect(result?.options.map((option) => option.label)).toEqual([table]); + const option = result!.options[0]!; + if (typeof option.apply === "function") option.apply(currentView, option, result!.from, currentView.state.doc.length); + else currentView.dispatch(insertCompletionText(currentView.state, option.apply ?? option.label, result!.from, currentView.state.doc.length)); + expect(currentView.state.doc.toString()).toBe(`${sql}${table.includes(".") ? `"${table}"` : table}`); + }); + + it.each(["collision", "always", "never"])("preserves %s schema qualification for default tables", async (setting) => { + const { provide, settings, currentView } = snowflakeHarness("SELECT * FROM "); + settings.editorSettings.tableCompletionSchemaQualification = setting; + const result = (await provide())!; + const table = result.options.find((option) => option.label === "DEFAULT_TABLE")!; + if (typeof table.apply === "function") table.apply(currentView, table, result.from, currentView.state.doc.length); + else currentView.dispatch(insertCompletionText(currentView.state, table.apply ?? table.label, result.from, currentView.state.doc.length)); + expect(currentView.state.doc.toString()).toBe(`SELECT * FROM ${setting === "always" ? "PUBLIC.DEFAULT_TABLE" : "DEFAULT_TABLE"}`); + }); + + it("prefers a known schema over a same-name database after uppercase folding", async () => { + const { catalogs, provide, store } = snowflakeHarness("SELECT * FROM other_db."); + catalogs.DEFAULT_DB!.OTHER_DB = ["SCHEMA_TABLE"]; + expect((await provide())?.options.map((option) => option.label)).toEqual(["SCHEMA_TABLE"]); + expect(store.listCompletionSchemas).not.toHaveBeenCalledWith("connection", "OTHER_DB"); + }); + + it.each(['"other_db".', '"other_db".PUBLIC.', "MISSING_DB.", '"OTHER_DB"."public".'])("does not substitute case-distinct or unknown namespaces for %s", async (qualifier) => { + const { provide } = snowflakeHarness(`SELECT * FROM ${qualifier}`); + expect((await provide())?.options ?? []).toEqual([]); + }); + + it.each(["empty", "denied"])("does not invent tables when explicit schema metadata is %s", async (failure) => { + const { provide, store } = snowflakeHarness("SELECT * FROM OTHER_DB.PUBLIC."); + if (failure === "empty") store.listCompletionTables.mockResolvedValue([]); + else store.listCompletionTables.mockRejectedValue(new Error("permission denied")); + expect((await provide())?.options ?? []).toEqual([]); + }); + + it("offers accessible database names alongside the default schema and tables", async () => { + const { provide } = snowflakeHarness("SELECT * FROM "); + expect((await provide())?.options.map((option) => option.label)).toEqual(expect.arrayContaining(["DEFAULT_DB", "OTHER_DB", "PUBLIC", "DEFAULT_TABLE"])); + }); + + it("keeps cached table metadata separate when moving between databases", async () => { + const { provide, currentView, props } = snowflakeHarness("SELECT * FROM OTHER_DB.PUBLIC."); + expect((await provide())?.options.map((option) => option.label)).toEqual(["OTHER_TABLE"]); + const sql = "SELECT * FROM DEFAULT_DB.PUBLIC."; + props.modelValue = sql; + currentView.dispatch({ changes: { from: 0, to: currentView.state.doc.length, insert: sql }, selection: { anchor: sql.length } }); + expect((await provide())?.options.map((option) => option.label)).toEqual(["DEFAULT_TABLE"]); + const quotedSql = 'SELECT * FROM "Other_Db".PUBLIC.'; + props.modelValue = quotedSql; + currentView.dispatch({ changes: { from: 0, to: currentView.state.doc.length, insert: quotedSql }, selection: { anchor: quotedSql.length } }); + expect((await provide())?.options.map((option) => option.label)).toEqual(["QUOTED_TABLE"]); + }); + + it.each([false, true])("preserves the quote-identifiers preference (%s) for mixed-case table insertion", async (quoteIdentifiers) => { + const { catalogs, provide, currentView, settings } = snowflakeHarness("SELECT * FROM "); + catalogs.DEFAULT_DB!.PUBLIC = ["DimCustomer"]; + settings.editorSettings.generateSqlQuoteIdentifiers = quoteIdentifiers; + settings.editorSettings.tableCompletionSchemaQualification = "always"; + const result = (await provide())!; + const table = result.options.find((option) => option.label === "DimCustomer")!; + if (typeof table.apply === "function") table.apply(currentView, table, result.from, currentView.state.doc.length); + else currentView.dispatch(insertCompletionText(currentView.state, table.apply ?? table.label, result.from, currentView.state.doc.length)); + expect(currentView.state.doc.toString()).toBe(`SELECT * FROM PUBLIC.${quoteIdentifiers ? '"DimCustomer"' : "DimCustomer"}`); + }); + + it.each(["collision", "always", "never"])("preserves %s behavior when table names collide across schemas", async (setting) => { + const { provide, store, currentView, settings } = snowflakeHarness("SELECT * FROM "); + store.listCompletionTables.mockResolvedValue([ + { name: "SHARED", schema: "PUBLIC" }, + { name: "SHARED", schema: "SALES" }, + ]); + settings.editorSettings.tableCompletionSchemaQualification = setting; + const result = (await provide())!; + const table = result.options.find((option) => option.label === "SHARED")!; + if (typeof table.apply === "function") table.apply(currentView, table, result.from, currentView.state.doc.length); + else currentView.dispatch(insertCompletionText(currentView.state, table.apply ?? table.label, result.from, currentView.state.doc.length)); + expect(currentView.state.doc.toString()).toBe(`SELECT * FROM ${setting === "never" ? "SHARED" : "PUBLIC.SHARED"}`); + }); + + it.each([ + ["other_db.public.other_table", "OTHER_DB", "PUBLIC", "OTHER_TABLE"], + ['"Other_Db".PUBLIC.QUOTED_TABLE', "Other_Db", "PUBLIC", "QUOTED_TABLE"], + ['"My.DB"."Sales Area"."Dim.Customer"', "My.DB", "Sales Area", "Dim.Customer"], + ])("loads columns from the explicit reference %s", async (reference, database, schema, table) => { + const { provide, store } = snowflakeHarness(`SELECT * FROM ${reference} t WHERE t.`); + store.listCompletionColumns.mockImplementation(async (...args: unknown[]) => (args[1] === database && args[2] === table && args[3] === schema ? [{ name: "ID", table, schema, dataType: "NUMBER" }] : [])); + const result = await provide(); + expect(store.listCompletionColumns.mock.calls.map((args) => (args as unknown[]).slice(0, 4))).toContainEqual(["connection", database, table, schema]); + expect(result?.options.map((option) => option.label)).toContain("ID"); + }); + + it.each(["empty", "denied"])("does not invent schemas when metadata is %s", async (failure) => { + const { provide, store } = snowflakeHarness("SELECT * FROM OTHER_DB."); + if (failure === "empty") store.listCompletionSchemas.mockResolvedValue([]); + else store.listCompletionSchemas.mockRejectedValue(new Error("permission denied")); + expect((await provide())?.options ?? []).toEqual([]); + }); + + it("uses locally cached database and schema metadata", async () => { + const { provide, store } = snowflakeHarness("SELECT * FROM Other_Db."); + store.lookupLocalCompletionDatabases.mockReturnValue(["OTHER_DB"]); + store.lookupLocalCompletionSchemas.mockImplementation((...args: unknown[]) => (args[1] === "OTHER_DB" ? ["PUBLIC"] : [])); + expect((await provide())?.options.map((option) => option.label)).toEqual(["PUBLIC"]); + expect(store.refreshCompletionSchemas).toHaveBeenCalledWith("connection", "OTHER_DB"); + }); +}); + describe("QueryEditor completion provider ownership", () => { it.each(["cursor", "document", "view", "composition", "mode"])("rejects asynchronous analysis after a %s change", async (change) => { let resolve!: (result: SqlCompletionAnalysisResult) => void; diff --git a/apps/desktop/src/components/editor/useQueryEditorCompletion.ts b/apps/desktop/src/components/editor/useQueryEditorCompletion.ts index e5e1ea414b..eeaf0f4923 100644 --- a/apps/desktop/src/components/editor/useQueryEditorCompletion.ts +++ b/apps/desktop/src/components/editor/useQueryEditorCompletion.ts @@ -1071,6 +1071,7 @@ export function useQueryEditorCompletion(options: QueryEditorCompletionOptions) }); const shouldLoadTables = !schemaLookupDatabase && (completionContext.suggestTables || (!!completionContext.qualifier && !isReferencedTableQualifier(completionContext))); const tableLookupTarget = resolveSqlCompletionTableLookupTarget({ + databaseType: props.databaseType, currentDatabase: scope.database, currentSchema: scope.schema, supportsDatabaseQualifier: supportsDatabaseQualifierCompletion(), @@ -1233,6 +1234,7 @@ export function useQueryEditorCompletion(options: QueryEditorCompletionOptions) knownSchemas: currentDatabaseSchemaNames, }); const tableLookupTarget = resolveSqlCompletionTableLookupTarget({ + databaseType: props.databaseType, currentDatabase: database, currentSchema: scope.schema, supportsDatabaseQualifier: supportsDatabaseQualifierCompletion(), @@ -1443,6 +1445,7 @@ export function useQueryEditorCompletion(options: QueryEditorCompletionOptions) }); const shouldLoadTables = !schemaLookupDatabase && (completionContext.suggestTables || (!!completionContext.qualifier && !isReferencedTableQualifier(completionContext))); const tableLookupTarget = resolveSqlCompletionTableLookupTarget({ + databaseType: props.databaseType, currentDatabase: scope.database, currentSchema: scope.schema, supportsDatabaseQualifier: supportsDatabaseQualifierCompletion(), diff --git a/apps/desktop/src/lib/__tests__/database/databaseFeatureSupport.spec.ts b/apps/desktop/src/lib/__tests__/database/databaseFeatureSupport.spec.ts index b07e1dc4e8..a0aa2e66ba 100644 --- a/apps/desktop/src/lib/__tests__/database/databaseFeatureSupport.spec.ts +++ b/apps/desktop/src/lib/__tests__/database/databaseFeatureSupport.spec.ts @@ -145,19 +145,19 @@ describe("zookeeper query capabilities", () => { }); describe("database and schema qualifiers", () => { - it.each(["sqlserver", "trino", "prestosql"] as const)("supports three-part object names for %s", (databaseType) => { + it.each(["sqlserver", "trino", "prestosql", "snowflake"] as const)("supports three-part object names for %s", (databaseType) => { expect(supportsDatabaseSchemaQualifier(databaseType)).toBe(true); }); - it.each(["mysql", "postgres", "oracle", "snowflake"] as const)("does not widen unverified three-part completion for %s", (databaseType) => { + it.each(["mysql", "postgres", "oracle"] as const)("does not widen unverified three-part completion for %s", (databaseType) => { expect(supportsDatabaseSchemaQualifier(databaseType)).toBe(false); }); - it.each(["mysql", "sqlite", "sqlserver"] as const)("suggests database names for %s", (databaseType) => { + it.each(["mysql", "sqlite", "sqlserver", "snowflake"] as const)("suggests database names for %s", (databaseType) => { expect(supportsDatabaseNameCompletion(databaseType)).toBe(true); }); - it.each(["postgres", "oracle", "snowflake", "trino", "prestosql"] as const)("does not add database name completion for %s", (databaseType) => { + it.each(["postgres", "oracle", "trino", "prestosql"] as const)("does not add database name completion for %s", (databaseType) => { expect(supportsDatabaseNameCompletion(databaseType)).toBe(false); }); }); diff --git a/apps/desktop/src/lib/__tests__/sql/semantic/dialect.spec.ts b/apps/desktop/src/lib/__tests__/sql/semantic/dialect.spec.ts index 761c57edad..9d87b81472 100644 --- a/apps/desktop/src/lib/__tests__/sql/semantic/dialect.spec.ts +++ b/apps/desktop/src/lib/__tests__/sql/semantic/dialect.spec.ts @@ -1,6 +1,17 @@ import { describe, expect, it } from "vitest"; import { sqlReferenceAnalysisDialectFor, sqlSemanticDialectFor } from "@/lib/sql/semantic/dialect"; +describe("Snowflake semantic dialect", () => { + it.each([undefined, "mysql", "postgres"] as const)("preserves Snowflake identifiers with fallback %s", (fallback) => { + const dialect = sqlSemanticDialectFor({ databaseType: "snowflake", dialect: fallback }); + expect(dialect.id).toBe("snowflake"); + expect(dialect.normalizeIdentifier("Adventureworks2025")).toBe("ADVENTUREWORKS2025"); + expect(dialect.normalizeIdentifier("My.DB", true)).toBe("My.DB"); + expect(dialect.quoteIdentifier('Sales "Area"')).toBe('"Sales ""Area"""'); + expect(dialect.supportsAsForTableAlias).toBe(true); + }); +}); + describe("ClickHouse semantic dialect", () => { it("takes precedence over the legacy MySQL behavior dialect", () => { const dialect = sqlSemanticDialectFor({ diff --git a/apps/desktop/src/lib/__tests__/sql/sqlCompletionLookupTarget.spec.ts b/apps/desktop/src/lib/__tests__/sql/sqlCompletionLookupTarget.spec.ts index 4c800a424b..e4346b7992 100644 --- a/apps/desktop/src/lib/__tests__/sql/sqlCompletionLookupTarget.spec.ts +++ b/apps/desktop/src/lib/__tests__/sql/sqlCompletionLookupTarget.spec.ts @@ -17,6 +17,34 @@ import { } from "@/lib/sql/sqlCompletionLookupTarget"; describe("sqlCompletionLookupTarget", () => { + it.each([ + ["Adventureworks2025.Public", "ADVENTUREWORKS2025", "PUBLIC"], + ['"Adventureworks2025".public', "Adventureworks2025", "PUBLIC"], + ['"My.DB"."Sales Area"', "My.DB", "Sales Area"], + ['"My.DB"."A""B"', "My.DB", 'A"B'], + ['"adventureworks2025".PUBLIC', "adventureworks2025", "PUBLIC"], + ])("resolves Snowflake %s using unquoted uppercase and quoted exact names", (qualifier, database, schema) => { + const sql = `SELECT * FROM ${qualifier}.Di`; + const completionContext = getSqlCompletionContext(sql, sql.length, { databaseType: "snowflake" }); + expect(completionContext.qualifierParts).toEqual([database, schema]); + expect( + resolveSqlCompletionTableLookupTarget({ + databaseType: "snowflake", + currentDatabase: "DEFAULT_DB", + currentSchema: "PUBLIC", + supportsDatabaseQualifier: false, + supportsDatabaseSchemaQualifier: true, + knownDatabases: ["Adventureworks2025", "ADVENTUREWORKS2025", "My.DB"], + completionContext, + }), + ).toEqual({ database, schema, filter: "Di", qualifierDatabase: database }); + }); + + it.each(["postgres", "sqlserver", "trino", "mysql"] as const)("preserves %s qualifier casing", (databaseType) => { + const sql = "SELECT * FROM Adventureworks2025.Public."; + expect(getSqlCompletionContext(sql, sql.length, { databaseType }).qualifierParts).toEqual(["Adventureworks2025", "Public"]); + }); + it("treats qualified table completion as a database lookup for MySQL-compatible engines", () => { const target = resolveSqlCompletionTableLookupTarget({ currentDatabase: "default_db", diff --git a/apps/desktop/src/lib/database/databaseFeatureSupport.ts b/apps/desktop/src/lib/database/databaseFeatureSupport.ts index eb3a273c7a..6921292406 100644 --- a/apps/desktop/src/lib/database/databaseFeatureSupport.ts +++ b/apps/desktop/src/lib/database/databaseFeatureSupport.ts @@ -13,7 +13,7 @@ export function supportsDatabaseSchemaQualifier(dbType?: DatabaseType): boolean } export function supportsDatabaseNameCompletion(dbType?: DatabaseType): boolean { - return !!dbType && ((!isSchemaAware(dbType) && !isSingleDatabase(dbType)) || dbType === "sqlserver"); + return !!dbType && ((!isSchemaAware(dbType) && !isSingleDatabase(dbType)) || dbType === "sqlserver" || dbType === "snowflake"); } /** diff --git a/apps/desktop/src/lib/sql/semantic/dialect.ts b/apps/desktop/src/lib/sql/semantic/dialect.ts index a1fe2cd01d..a33a4f0b62 100644 --- a/apps/desktop/src/lib/sql/semantic/dialect.ts +++ b/apps/desktop/src/lib/sql/semantic/dialect.ts @@ -64,6 +64,15 @@ export const SQL_SEMANTIC_DIALECTS: Record = quoteIdentifier: (identifier) => quoteWith(identifier, '"'), qualifierRole: roleForGenericQualifier, }, + snowflake: { + id: "snowflake", + identifierQuotes: [{ open: '"', close: '"' }], + supportsAsForTableAlias: true, + projectionAliasVisibility: defaultProjectionAliasVisibility, + normalizeIdentifier: upperUnquoted, + quoteIdentifier: (identifier) => quoteWith(identifier, '"'), + qualifierRole: roleForGenericQualifier, + }, postgres: { id: "postgres", identifierQuotes: [{ open: '"', close: '"' }], @@ -190,6 +199,7 @@ export function sqlReferenceAnalysisDialectFor(options: { databaseType?: Databas } export function sqlSemanticDialectFor(options: { databaseType?: DatabaseType; dialect?: "mysql" | "postgres" | "sqlserver" | "clickhouse" | "doris" | "soql" }): SqlSemanticDialectAdapter { + if (options.databaseType === "snowflake") return SQL_SEMANTIC_DIALECTS.snowflake; if (options.databaseType === "salesforce") return SQL_SEMANTIC_DIALECTS.soql; if (options.databaseType === "clickhouse") return SQL_SEMANTIC_DIALECTS.clickhouse; // Doris/StarRocks connections ride the editor's MySQL fallback dialect (codeMirrorSqlDialect maps diff --git a/apps/desktop/src/lib/sql/sqlCompletion.ts b/apps/desktop/src/lib/sql/sqlCompletion.ts index cd3a1e2496..20408067ec 100644 --- a/apps/desktop/src/lib/sql/sqlCompletion.ts +++ b/apps/desktop/src/lib/sql/sqlCompletion.ts @@ -1471,7 +1471,7 @@ type SqlCompletionApplyDialect = "mysql" | "postgres" | "sqlserver" | "oracle" | const MYSQL_LIKE_IDENTIFIER_DATABASES = new Set(["mysql", "clickhouse", "hive", "argo", "transwarp", "kyuubi", "impala", "spark", "databend", "tdengine", "access", "doris", "starrocks"]); const POSTGRES_LIKE_IDENTIFIER_DATABASES = new Set(["postgres", "redshift", "gaussdb", "kingbase", "highgo", "uxdb", "vastbase", "kwdb", "opengauss"]); export const ORACLE_COMPAT_IDENTIFIER_DATABASES = new Set(["oracle", "oceanbase-oracle", "yashandb", "oscar", "xugu"]); -const UPPER_FOLDING_IDENTIFIER_DATABASES = new Set(["dameng", "db2"]); +const UPPER_FOLDING_IDENTIFIER_DATABASES = new Set(["dameng", "db2", "snowflake"]); function sqlCompletionApplyDialect(databaseType: DatabaseType | undefined, fallback: "mysql" | "postgres" | "sqlserver" | undefined): SqlCompletionApplyDialect | undefined { if (!databaseType) return fallback; @@ -2850,7 +2850,7 @@ function parseTrailingIdentifierContext(input: string, databaseType?: DatabaseTy index -= 1; continue; } - parts.unshift(unquoteIdentifier(parsed.raw)); + parts.unshift(parsed.raw); index = parsed.start; if (index <= 0 || tail[index - 1] !== ".") break; index -= 1; @@ -2860,8 +2860,8 @@ function parseTrailingIdentifierContext(input: string, databaseType?: DatabaseTy const start = index; if (parts.length >= 2 || endsWithDot) { - const qualifierParts = endsWithDot ? parts : parts.slice(0, -1); - const prefixPart = endsWithDot ? "" : (parts[parts.length - 1] ?? ""); + const qualifierParts = (endsWithDot ? parts : parts.slice(0, -1)).map((part) => completionLookupIdentifier(part, databaseType)); + const prefixPart = endsWithDot ? "" : unquoteIdentifier(parts[parts.length - 1] ?? ""); const qualifierValue = qualifierParts.join("."); return { start, @@ -2873,7 +2873,7 @@ function parseTrailingIdentifierContext(input: string, databaseType?: DatabaseTy return { start, - prefix: parts[0] ?? "", + prefix: unquoteIdentifier(parts[0] ?? ""), }; } @@ -3625,8 +3625,8 @@ function extractReferencedTables(sql: string, databaseType?: DatabaseType): SqlC } const rawParts = splitQualifiedNameRawParts(rawName); const omittedSqlServerSchema = databaseType === "sqlserver" && rawParts.length >= 3 && rawParts[rawParts.length - 2] === ""; - const unquotedRawParts = rawParts.map((part) => unquoteIdentifier(part)); - const parts = rawParts.map((part) => unquoteIdentifier(part)).filter(Boolean); + const unquotedRawParts = rawParts.map((part) => completionLookupIdentifier(part, databaseType)); + const parts = unquotedRawParts.filter(Boolean); const name = unquotedRawParts[unquotedRawParts.length - 1]; if (!name) continue; const table: SqlCompletionReferencedTable = { @@ -3994,6 +3994,12 @@ function unquoteIdentifier(value: string): string { return value; } +function completionLookupIdentifier(value: string, databaseType?: DatabaseType): string { + const name = unquoteIdentifier(value); + if (databaseType !== "snowflake") return name; + return value.startsWith('"') ? name.replaceAll('""', '"') : name.toUpperCase(); +} + export function quoteSqlIdentifier(identifier: string, dialect?: SqlCompletionApplyDialect): string { if (dialect === "oracle") { // Oracle folds bare identifiers to uppercase, so only all-uppercase names diff --git a/apps/desktop/src/lib/sql/sqlCompletionLookupTarget.ts b/apps/desktop/src/lib/sql/sqlCompletionLookupTarget.ts index 7f2501aeae..68011c1a1c 100644 --- a/apps/desktop/src/lib/sql/sqlCompletionLookupTarget.ts +++ b/apps/desktop/src/lib/sql/sqlCompletionLookupTarget.ts @@ -274,6 +274,7 @@ export function resolveSqlCompletionSchemaLookupDatabase(options: { } export function resolveSqlCompletionTableLookupTarget(options: { + databaseType?: DatabaseType; currentDatabase: string; currentSchema?: string; supportsDatabaseQualifier: boolean; @@ -287,7 +288,7 @@ export function resolveSqlCompletionTableLookupTarget(options: { if (options.supportsDatabaseSchemaQualifier && completionContext.suggestTables && !completionContext.insertTable && qualifierParts.length >= 2) { const databaseQualifier = qualifierParts[qualifierParts.length - 2]!; const schema = qualifierParts[qualifierParts.length - 1]!; - const database = findExactName(options.knownDatabases, databaseQualifier) ?? databaseQualifier; + const database = (options.databaseType === "snowflake" ? findCaseSensitiveName(options.knownDatabases, databaseQualifier) : findExactName(options.knownDatabases, databaseQualifier)) ?? databaseQualifier; return { database, schema, diff --git a/apps/desktop/src/stores/__tests__/connectionStore.snowflake.spec.ts b/apps/desktop/src/stores/__tests__/connectionStore.snowflake.spec.ts new file mode 100644 index 0000000000..9ecadc6070 --- /dev/null +++ b/apps/desktop/src/stores/__tests__/connectionStore.snowflake.spec.ts @@ -0,0 +1,133 @@ +import { createPinia, setActivePinia } from "pinia"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import type { ConnectionConfig, TreeNode } from "@/types/database"; + +beforeEach(() => { + vi.resetModules(); + const storage = new Map(); + vi.stubGlobal("localStorage", { + getItem: (key: string) => storage.get(key) ?? null, + setItem: (key: string, value: string) => storage.set(key, value), + removeItem: (key: string) => storage.delete(key), + }); + setActivePinia(createPinia()); +}); + +afterEach(() => vi.unstubAllGlobals()); + +async function createHarness(mode: "simple" | "grouped" = "simple", disk = new Map()) { + const listSchemaInfos = vi.fn(async (_id: string, _database: string) => [ + { name: "PUBLIC", comment: null }, + { name: "Sales Area", comment: null }, + ]); + const listTables = vi.fn(async (_id: string, database: string, schema: string) => [{ name: "SHARED_TABLE", table_type: "TABLE", schema, comment: database }]); + const getColumns = vi.fn().mockResolvedValue([{ name: "ID", data_type: "NUMBER", is_nullable: false }]); + const loadSchemaCache = vi.fn(async (key: string) => disk.get(key) ?? null); + vi.doMock("@/lib/backend/tauriRuntime", () => ({ isTauriRuntime: () => false })); + vi.doMock("@/lib/backend/api", () => ({ + checkConnectionHealth: vi.fn().mockResolvedValue(undefined), + listInstalledAgents: vi.fn().mockResolvedValue([]), + deleteSchemaCachePrefix: vi.fn(async (prefix: string) => { + for (const key of disk.keys()) if (key.startsWith(prefix)) disk.delete(key); + }), + listSchemaInfos, + listTables, + getColumns, + listObjects: vi.fn().mockResolvedValue([]), + loadSchemaCache, + saveSchemaCache: vi.fn(async (key: string, value: string) => { + disk.set(key, value); + }), + saveConnections: vi.fn().mockResolvedValue(undefined), + saveSidebarLayout: vi.fn().mockResolvedValue(undefined), + })); + const { useConnectionStore } = await import("@/stores/connectionStore"); + const { useSettingsStore } = await import("@/stores/settingsStore"); + useSettingsStore().editorSettings.sidebarObjectDisplay = mode; + const store = useConnectionStore(); + const connection = { id: "snowflake-test", name: "Snowflake", db_type: "snowflake", database: "DEFAULT_DB", host: "fixture.invalid", port: 443, username: "", password: "" } as ConnectionConfig; + const databases: TreeNode[] = ["DEFAULT_DB", "OTHER_DB"].map((database) => ({ id: `${connection.id}:${database}`, label: database, type: "database", connectionId: connection.id, database, children: [] })); + store.connections = [connection]; + store.connectedIds.add(connection.id); + store.treeNodes = [{ id: connection.id, label: connection.name, type: "connection", connectionId: connection.id, children: databases }]; + return { store, connection, databases, listSchemaInfos, listTables, getColumns, loadSchemaCache, disk }; +} + +describe("Snowflake sidebar namespace metadata", () => { + it.each(["simple", "grouped"] as const)("loads database/schema/table/column scopes in %s mode", async (mode) => { + const { store, connection, databases, listSchemaInfos, listTables, getColumns } = await createHarness(mode); + for (const database of databases) { + await store.loadTreeNodeChildren(database); + expect(listSchemaInfos).toHaveBeenCalledWith(connection.id, database.database); + const schemas = database.children!.filter((node) => node.type === "schema"); + expect(schemas.map((node) => node.schema)).toEqual(["PUBLIC", "Sales Area"]); + for (const schema of schemas) { + await store.loadTreeNodeChildren(schema); + const parent = mode === "simple" ? schema : schema.children!.find((node) => node.type === "group-tables")!; + if (mode === "grouped") await store.loadTreeNodeChildren(parent); + const table = parent.children!.find((node) => node.type === "table")!; + expect(table).toMatchObject({ label: "SHARED_TABLE", database: database.database, schema: schema.schema, comment: database.database }); + await store.loadTreeNodeChildren(table); + const columns = table.children!.find((node) => node.type === "group-columns")!; + await store.loadTreeNodeChildren(columns); + expect(getColumns).toHaveBeenCalledWith(connection.id, database.database, schema.schema, "SHARED_TABLE", undefined); + expect(columns.children?.map((node) => node.label)).toContain("ID (NUMBER)"); + } + } + expect(listTables.mock.calls.every((args) => ["PUBLIC", "Sales Area"].includes(args[2]))).toBe(true); + }); + + it("reuses cached schema metadata and refreshes changed schemas explicitly", async () => { + const { store, databases, listSchemaInfos } = await createHarness(); + const database = databases[1]!; + await store.loadTreeNodeChildren(database); + await store.loadTreeNodeChildren(database); + expect(listSchemaInfos).toHaveBeenCalledOnce(); + listSchemaInfos.mockResolvedValue([{ name: "NEW_SCHEMA", comment: null }]); + await store.loadTreeNodeChildren(database, { force: true }); + expect(listSchemaInfos).toHaveBeenCalledTimes(2); + expect(database.children?.filter((node) => node.type === "schema").map((node) => node.schema)).toEqual(["NEW_SCHEMA"]); + }); + + it("restores persisted schema and table metadata after store reload", async () => { + const first = await createHarness(); + await first.store.loadTreeNodeChildren(first.databases[1]!); + const schema = first.databases[1]!.children!.find((node) => node.schema === "PUBLIC")!; + await first.store.loadTreeNodeChildren(schema); + expect(first.disk.size).toBeGreaterThan(0); + vi.resetModules(); + setActivePinia(createPinia()); + const restored = await createHarness("simple", first.disk); + let finishRefresh!: (tables: Awaited>) => void; + restored.listTables.mockReturnValue( + new Promise((resolve) => { + finishRefresh = resolve; + }), + ); + await restored.store.loadTreeNodeChildren(restored.databases[1]!); + const restoredSchema = restored.databases[1]!.children!.find((node) => node.schema === "PUBLIC")!; + await restored.store.loadTreeNodeChildren(restoredSchema); + expect(restored.listSchemaInfos).not.toHaveBeenCalled(); + expect(restoredSchema.children).toEqual(expect.arrayContaining([expect.objectContaining({ label: "SHARED_TABLE", database: "OTHER_DB", schema: "PUBLIC" })])); + finishRefresh([{ name: "REFRESHED_TABLE", table_type: "TABLE", schema: "PUBLIC", comment: "OTHER_DB" }]); + await vi.waitFor(() => expect(restoredSchema.children).toEqual(expect.arrayContaining([expect.objectContaining({ label: "REFRESHED_TABLE", database: "OTHER_DB", schema: "PUBLIC" })]))); + }); + + it.each(["empty", "denied"])("preserves %s schema results without a flat table fallback", async (state) => { + const { store, databases, listSchemaInfos, listTables } = await createHarness(); + const database = databases[1]!; + if (state === "empty") { + listSchemaInfos.mockResolvedValue([]); + await store.loadTreeNodeChildren(database); + } else { + listSchemaInfos.mockRejectedValueOnce(new Error("permission denied")); + await expect(store.loadTreeNodeChildren(database)).rejects.toThrow("permission denied"); + } + expect(database.children?.filter((node) => node.type === "schema" || node.type === "table")).toEqual([]); + expect(listTables).not.toHaveBeenCalled(); + if (state === "denied") { + await store.loadTreeNodeChildren(database, { force: true }); + expect(database.children?.filter((node) => node.type === "schema")).toHaveLength(2); + } + }); +}); diff --git a/crates/dbx-core/assets/database-drivers.manifest.json b/crates/dbx-core/assets/database-drivers.manifest.json index 0e6dbb7151..d4ba91a8b8 100644 --- a/crates/dbx-core/assets/database-drivers.manifest.json +++ b/crates/dbx-core/assets/database-drivers.manifest.json @@ -1750,7 +1750,9 @@ "skipTcpProbe": true, "defaultPort": 443, "traits": { - "schemaAware": true + "schemaAware": true, + "databaseSchemaQualified": true, + "treeSchema": true }, "supportLevel": "operate", "capabilities": { diff --git a/plugins/connection-types/snowflake.yaml b/plugins/connection-types/snowflake.yaml index bb33f97967..58354b3f37 100644 --- a/plugins/connection-types/snowflake.yaml +++ b/plugins/connection-types/snowflake.yaml @@ -14,6 +14,8 @@ skipTcpProbe: true defaultPort: 443 traits: schemaAware: true + databaseSchemaQualified: true + treeSchema: true supportLevel: operate capabilities: queryExecution: true