From d1cce3ed1ec7b21075d0062abb981ba801dd0a3a Mon Sep 17 00:00:00 2001 From: zipg Date: Mon, 13 Jul 2026 18:51:32 +0800 Subject: [PATCH] fix(kingbase): quote MySQL-compatible identifiers --- .../java/com/dbx/agent/AbstractJdbcAgent.java | 10 ++++ .../java/com/dbx/agent/AgentProtocol.java | 2 + .../java/com/dbx/agent/DatabaseAgent.java | 4 ++ .../java/com/dbx/agent/JsonRpcServer.java | 3 + .../java/com/dbx/agent/PostgresLikeAgent.java | 5 ++ .../src/main/resources/agent-protocol-v1.json | 1 + .../src/main/resources/agent-protocol-v2.json | 1 + .../com/dbx/agent/kingbase/KingbaseAgent.java | 14 ++++- .../dbx/agent/kingbase/KingbaseAgentTest.java | 36 ++++++++++++ apps/desktop/src/components/grid/DataGrid.vue | 7 ++- .../src/components/objects/ObjectBrowser.vue | 2 + .../src/components/sidebar/TreeItem.vue | 2 + .../src/composables/useDataGridActions.ts | 5 +- .../src/composables/useNavigationTargets.ts | 5 ++ .../__tests__/table/tableSelectSql.spec.ts | 15 ++++- apps/desktop/src/lib/backend/api.ts | 1 + apps/desktop/src/lib/backend/http.ts | 5 ++ apps/desktop/src/lib/backend/tauri.ts | 5 ++ apps/desktop/src/lib/dataGrid/dataGridSql.ts | 1 + apps/desktop/src/lib/export/databaseExport.ts | 2 + apps/desktop/src/lib/table/tableDataExport.ts | 2 + apps/desktop/src/lib/table/tableSelectSql.ts | 9 +++ .../queryStore.hiddenPrimaryKey.spec.ts | 1 + .../queryStore.tableDataRefresh.spec.ts | 1 + apps/desktop/src/stores/connectionStore.ts | 28 ++++++++++ apps/desktop/src/stores/queryStore.ts | 14 +++-- crates/dbx-core/assets/agent-protocol-v1.json | 1 + crates/dbx-core/assets/agent-protocol-v2.json | 1 + crates/dbx-core/src/connection.rs | 29 ++++++++++ crates/dbx-core/src/data_grid_sql.rs | 37 ++++++++++-- crates/dbx-core/src/db/agent_driver.rs | 15 ++++- crates/dbx-core/src/sql_dialect.rs | 1 + .../dbx-core/src/sql_dialect/table_select.rs | 56 +++++++++++++++++-- crates/dbx-core/src/sql_dialect/tests.rs | 43 ++++++++++++++ crates/dbx-core/src/sql_dialect/types.rs | 2 + crates/dbx-web/src/main.rs | 1 + crates/dbx-web/src/routes/connection.rs | 19 +++++++ src-tauri/src/commands/connection.rs | 9 +++ src-tauri/src/lib.rs | 1 + 39 files changed, 373 insertions(+), 23 deletions(-) diff --git a/agents/common/src/main/java/com/dbx/agent/AbstractJdbcAgent.java b/agents/common/src/main/java/com/dbx/agent/AbstractJdbcAgent.java index bbb7dcb21b..6a0a0eed69 100644 --- a/agents/common/src/main/java/com/dbx/agent/AbstractJdbcAgent.java +++ b/agents/common/src/main/java/com/dbx/agent/AbstractJdbcAgent.java @@ -20,6 +20,16 @@ public abstract class AbstractJdbcAgent extends BaseDatabaseAgent { return connection; } + @Override + public String getIdentifierQuote() { + try { + String quote = requireConnected().getMetaData().getIdentifierQuoteString(); + return quote == null || quote.trim().isEmpty() ? "" : quote.trim(); + } catch (Exception ignored) { + return ""; + } + } + @Override public final void connect(ConnectParams params) { uncheckedVoid(() -> { diff --git a/agents/common/src/main/java/com/dbx/agent/AgentProtocol.java b/agents/common/src/main/java/com/dbx/agent/AgentProtocol.java index 17ef879e1a..6f54536bf7 100644 --- a/agents/common/src/main/java/com/dbx/agent/AgentProtocol.java +++ b/agents/common/src/main/java/com/dbx/agent/AgentProtocol.java @@ -16,6 +16,7 @@ public final class AgentProtocol { public static final String METHOD_CANCEL_SESSION = "cancel_session"; public static final String METHOD_TEST_CONNECTION = "test_connection"; public static final String METHOD_VALIDATE_CONNECTION = "validate_connection"; + public static final String METHOD_CONNECTION_INFO = "connection_info"; public static final String METHOD_LIST_DATABASES = "list_databases"; public static final String METHOD_LIST_SCHEMAS = "list_schemas"; public static final String METHOD_LIST_TABLES = "list_tables"; @@ -100,6 +101,7 @@ public final class AgentProtocol { METHOD_CONNECT, METHOD_TEST_CONNECTION, METHOD_VALIDATE_CONNECTION, + METHOD_CONNECTION_INFO, METHOD_LIST_DATABASES, METHOD_LIST_SCHEMAS, METHOD_LIST_TABLES, diff --git a/agents/common/src/main/java/com/dbx/agent/DatabaseAgent.java b/agents/common/src/main/java/com/dbx/agent/DatabaseAgent.java index cc29afa5ca..0c8ca64cba 100644 --- a/agents/common/src/main/java/com/dbx/agent/DatabaseAgent.java +++ b/agents/common/src/main/java/com/dbx/agent/DatabaseAgent.java @@ -13,6 +13,10 @@ public interface DatabaseAgent { boolean testConnection(ConnectParams params); + default String getIdentifierQuote() { + return ""; + } + List listDatabases(); List listSchemas(); diff --git a/agents/common/src/main/java/com/dbx/agent/JsonRpcServer.java b/agents/common/src/main/java/com/dbx/agent/JsonRpcServer.java index 3a59784f2f..d93de45474 100644 --- a/agents/common/src/main/java/com/dbx/agent/JsonRpcServer.java +++ b/agents/common/src/main/java/com/dbx/agent/JsonRpcServer.java @@ -125,6 +125,9 @@ public final class JsonRpcServer { return Collections.singletonMap("ok", true); } ensureLiveConnection(method); + if (AgentProtocol.METHOD_CONNECTION_INFO.equals(method)) { + return Collections.singletonMap("identifierQuote", agent.getIdentifierQuote()); + } if (AgentProtocol.METHOD_LIST_DATABASES.equals(method)) { return agent.listDatabases(); } diff --git a/agents/common/src/main/java/com/dbx/agent/PostgresLikeAgent.java b/agents/common/src/main/java/com/dbx/agent/PostgresLikeAgent.java index df591b45cf..ecba927181 100644 --- a/agents/common/src/main/java/com/dbx/agent/PostgresLikeAgent.java +++ b/agents/common/src/main/java/com/dbx/agent/PostgresLikeAgent.java @@ -28,6 +28,11 @@ public abstract class PostgresLikeAgent extends AbstractJdbcAgent { return mysqlCompatMode; } + @Override + public String getIdentifierQuote() { + return mysqlCompatMode ? "`" : super.getIdentifierQuote(); + } + @Override public QueryResult executeQuery(String sql, String schema, ExecuteQueryOptions options) { return JdbcExecutor.current().execute( diff --git a/agents/common/src/main/resources/agent-protocol-v1.json b/agents/common/src/main/resources/agent-protocol-v1.json index 29d9ca4359..db51f9650f 100644 --- a/agents/common/src/main/resources/agent-protocol-v1.json +++ b/agents/common/src/main/resources/agent-protocol-v1.json @@ -39,6 +39,7 @@ "connect", "test_connection", "validate_connection", + "connection_info", "list_databases", "list_schemas", "list_tables", diff --git a/agents/common/src/main/resources/agent-protocol-v2.json b/agents/common/src/main/resources/agent-protocol-v2.json index 33d2400f72..08900d5d76 100644 --- a/agents/common/src/main/resources/agent-protocol-v2.json +++ b/agents/common/src/main/resources/agent-protocol-v2.json @@ -46,6 +46,7 @@ "cancel_session", "test_connection", "validate_connection", + "connection_info", "list_databases", "list_schemas", "list_tables", diff --git a/agents/drivers/kingbase/src/main/java/com/dbx/agent/kingbase/KingbaseAgent.java b/agents/drivers/kingbase/src/main/java/com/dbx/agent/kingbase/KingbaseAgent.java index 09beafb27a..ee26db8bed 100644 --- a/agents/drivers/kingbase/src/main/java/com/dbx/agent/kingbase/KingbaseAgent.java +++ b/agents/drivers/kingbase/src/main/java/com/dbx/agent/kingbase/KingbaseAgent.java @@ -58,12 +58,24 @@ public final class KingbaseAgent extends PostgresLikeAgent { @Override protected void afterConnect(ConnectParams params, Connection connection) { postgresCatalogMode = false; + setMysqlCompatMode(params.isMysql_compat_mode()); if (params.isMysql_compat_mode()) { - setMysqlCompatMode(true); return; } postgresCatalogMode = !catalogExists(connection, "sys_catalog.sys_namespace") && catalogExists(connection, "pg_catalog.pg_namespace"); + if (!postgresCatalogMode && mysqlSqlModeExists(connection)) { + setMysqlCompatMode(true); + } + } + + private static boolean mysqlSqlModeExists(Connection connection) { + try (Statement stmt = connection.createStatement(); + ResultSet rs = stmt.executeQuery("SELECT 1 FROM sys_settings WHERE LOWER(name) = 'sql_mode'")) { + return rs.next(); + } catch (Exception ignored) { + return false; + } } private static boolean catalogExists(Connection connection, String catalog) { diff --git a/agents/drivers/kingbase/src/test/java/com/dbx/agent/kingbase/KingbaseAgentTest.java b/agents/drivers/kingbase/src/test/java/com/dbx/agent/kingbase/KingbaseAgentTest.java index 0e7a55be82..9513b0523f 100644 --- a/agents/drivers/kingbase/src/test/java/com/dbx/agent/kingbase/KingbaseAgentTest.java +++ b/agents/drivers/kingbase/src/test/java/com/dbx/agent/kingbase/KingbaseAgentTest.java @@ -137,6 +137,22 @@ class KingbaseAgentTest extends JdbcFakeExecutionBehaviorTest { Assertions.assertEquals("SET search_path TO \"app\"", agent.setSchemaSQL("app")); } + @Test + void detectsMysqlCompatModeFromServerSqlMode() throws Exception { + List sql = new ArrayList<>(); + KingbaseAgent agent = new KingbaseAgent(); + Connection connection = mysqlCompatConnection(sql); + + Method afterConnect = KingbaseAgent.class.getDeclaredMethod("afterConnect", ConnectParams.class, Connection.class); + afterConnect.setAccessible(true); + afterConnect.invoke(agent, new ConnectParams(), connection); + + Assertions.assertTrue(agent.isMysqlCompatMode()); + Assertions.assertEquals("`", agent.getIdentifierQuote()); + Assertions.assertEquals("SELECT 1 FROM sys_catalog.sys_namespace WHERE 1 = 0", sql.get(0)); + Assertions.assertEquals("SELECT 1 FROM sys_settings WHERE LOWER(name) = 'sql_mode'", sql.get(1)); + } + @Test void mysqlCompatListTablesUsesInformationSchema() { List sql = new ArrayList<>(); @@ -512,6 +528,26 @@ class KingbaseAgentTest extends JdbcFakeExecutionBehaviorTest { }); } + private static Connection mysqlCompatConnection(List sql) { + return proxy(Connection.class, (method, args) -> { + if ("createStatement".equals(method.getName())) { + return proxy(Statement.class, (statementMethod, statementArgs) -> { + if ("executeQuery".equals(statementMethod.getName())) { + String query = String.valueOf(statementArgs[0]); + sql.add(query); + if (query.contains("sys_settings")) { + return resultSet(new String[]{"probe"}, new Object[][]{{1}}); + } + return resultSet(new String[]{"probe"}, new Object[][]{}); + } + return defaultValue(statementMethod.getReturnType()); + }); + } + if ("isClosed".equals(method.getName())) return false; + return defaultValue(method.getReturnType()); + }); + } + private static ResultSet resultSet(String[] columns, Object[][] rows) { int[] index = {-1}; return proxy(ResultSet.class, (method, args) -> { diff --git a/apps/desktop/src/components/grid/DataGrid.vue b/apps/desktop/src/components/grid/DataGrid.vue index d71f421266..0487a0805d 100644 --- a/apps/desktop/src/components/grid/DataGrid.vue +++ b/apps/desktop/src/components/grid/DataGrid.vue @@ -89,7 +89,7 @@ import { formatElapsedSeconds } from "@/lib/common/elapsedTime"; import { dataGridCellDisplayText, dataGridCellEditorText } from "@/lib/dataGrid/dataGridCellCoercion"; import { createColumnDrafts } from "@/lib/table/tableStructureEditorState"; import type { BuildSingleColumnAlterSqlOptions } from "@/lib/table/tableStructureEditorSql"; -import { buildTableSelectSql, quoteTableIdentifier } from "@/lib/table/tableSelectSql"; +import { buildTableSelectSql, quoteTableDataIdentifier } from "@/lib/table/tableSelectSql"; import { uuid } from "@/lib/common/utils"; import { compactHeaderColumnType, resolveHeaderColumnType } from "@/lib/dataGrid/dataGridColumnType"; import { @@ -3958,6 +3958,7 @@ async function buildCurrentCountTarget(): Promise<{ sql: string; schema?: string if (props.tableMeta) { const sql = await buildDataGridCountSql({ databaseType: props.databaseType, + identifierQuote: connectionStore.connectionIdentifierQuote?.(props.connectionId), catalog: props.tableMeta.catalog, schema: props.tableMeta.schema, tableName: props.tableMeta.tableName, @@ -5612,6 +5613,7 @@ async function applyOrderBySearch() { if (!tableMeta) return; const sql = await buildTableSelectSql({ databaseType: resolvedDatabaseType.value, + identifierQuote: connectionStore.connectionIdentifierQuote?.(props.connectionId), catalog: tableMeta.catalog, schema: tableMeta.schema, tableName: tableMeta.tableName, @@ -5644,6 +5646,7 @@ async function applyWhereFilter() { if (!tableMeta) return; const sql = await buildTableSelectSql({ databaseType: resolvedDatabaseType.value, + identifierQuote: connectionStore.connectionIdentifierQuote?.(props.connectionId), catalog: tableMeta.catalog, schema: tableMeta.schema, tableName: tableMeta.tableName, @@ -5753,7 +5756,7 @@ function columnTypeCacheSignature(): string { watch(() => [props.result.columns.join("\u0000"), columnFormatterSignatures.value.join("\u0000"), columnTypeCacheSignature()], clearCellFormatCache); function quoteIdent(name: string): string { - return quoteTableIdentifier(props.databaseType, name); + return quoteTableDataIdentifier(props.databaseType, name, connectionStore.connectionIdentifierQuote?.(props.connectionId)); } function queryColumnRef(name: string): string { diff --git a/apps/desktop/src/components/objects/ObjectBrowser.vue b/apps/desktop/src/components/objects/ObjectBrowser.vue index 1d5d1e0d54..c7f18ecbe4 100644 --- a/apps/desktop/src/components/objects/ObjectBrowser.vue +++ b/apps/desktop/src/components/objects/ObjectBrowser.vue @@ -1048,6 +1048,7 @@ async function openNewQuery(row: ObjectBrowserRow) { tabId, await buildTableSelectSql({ databaseType: effectiveDatabaseType.value, + identifierQuote: connectionStore.connectionIdentifierQuote?.(props.connection.id), schema: row.schema || selectedSchema.value, tableName: row.name, limit: 100, @@ -1585,6 +1586,7 @@ async function exportDataLegacy(row: ObjectBrowserRow, format: "json" | "sql") { const queryColumns = props.connection.db_type === "neo4j" ? (tableColumns ?? (await api.getColumns(props.connection.id, props.database, schema || props.database, row.name))).map((column) => column.name) : undefined; const result = await fetchTableDataForExport({ databaseType: effectiveDatabaseType.value, + identifierQuote: connectionStore.connectionIdentifierQuote?.(props.connection.id), schema, tableName: row.name, columns: queryColumns, diff --git a/apps/desktop/src/components/sidebar/TreeItem.vue b/apps/desktop/src/components/sidebar/TreeItem.vue index a93bf9ef81..62d8cecac7 100644 --- a/apps/desktop/src/components/sidebar/TreeItem.vue +++ b/apps/desktop/src/components/sidebar/TreeItem.vue @@ -1394,6 +1394,7 @@ async function openData() { const includeRowId = usesSyntheticRowIdKey(effectiveDbType, primaryKeys, tableType); const sql = await buildTableSelectSql({ databaseType: effectiveDbType, + identifierQuote: connectionStore.connectionIdentifierQuote?.(node.connectionId), schema: tableSchema, tableName: node.label, tableType, @@ -3768,6 +3769,7 @@ async function exportDataLegacy(format: "csv" | "json" | "sql") { const effectiveDbType = effectiveDatabaseTypeForConnection(config); const result = await fetchTableDataForExport({ databaseType: effectiveDbType, + identifierQuote: connectionStore.connectionIdentifierQuote?.(connectionId), schema: node.schema, tableName: node.label, tableType: node.tableType, diff --git a/apps/desktop/src/composables/useDataGridActions.ts b/apps/desktop/src/composables/useDataGridActions.ts index ae21c2f9eb..fd97ba53c2 100644 --- a/apps/desktop/src/composables/useDataGridActions.ts +++ b/apps/desktop/src/composables/useDataGridActions.ts @@ -3,7 +3,7 @@ import { useI18n } from "vue-i18n"; import { useConnectionStore } from "@/stores/connectionStore"; import { useQueryStore } from "@/stores/queryStore"; import { useSettingsStore } from "@/stores/settingsStore"; -import { buildTableSelectSql, quoteTableIdentifier } from "@/lib/table/tableSelectSql"; +import { buildTableSelectSql, quoteTableDataIdentifier } from "@/lib/table/tableSelectSql"; import { editableRowIdentifierColumns, usesSyntheticRowIdKey } from "@/lib/table/tableEditing"; import { tableMetaForDataTab } from "@/lib/table/tableDataTabMeta"; import * as api from "@/lib/backend/api"; @@ -40,7 +40,7 @@ export function useDataGridActions(activeTab: ComputedRef) function quoteIdent(tab: QueryTab, name: string): string { const config = connectionStore.getConfig(tab.connectionId); - return quoteTableIdentifier(effectiveDatabaseTypeForConnection(config), name); + return quoteTableDataIdentifier(effectiveDatabaseTypeForConnection(config), name, connectionStore.connectionIdentifierQuote?.(tab.connectionId)); } function buildTableSql(tab: QueryTab, options: { orderBy?: string; limit?: number; offset?: number; whereInput?: string } = {}): Promise { @@ -51,6 +51,7 @@ export function useDataGridActions(activeTab: ComputedRef) const useRowId = usesSyntheticRowIdKey(effectiveDbType, primaryKeys, tableMeta?.tableType); return buildTableSelectSql({ databaseType: effectiveDbType, + identifierQuote: connectionStore.connectionIdentifierQuote?.(tab.connectionId), schema: tableMeta?.schema, tableName: tableMeta?.tableName ?? "", tableType: tableMeta?.tableType, diff --git a/apps/desktop/src/composables/useNavigationTargets.ts b/apps/desktop/src/composables/useNavigationTargets.ts index 5aa513a645..d3f270b95e 100644 --- a/apps/desktop/src/composables/useNavigationTargets.ts +++ b/apps/desktop/src/composables/useNavigationTargets.ts @@ -56,6 +56,7 @@ async function openTableTarget(target: NavigationTarget, options: { tableInfoTab await connectionStore.ensureConnected(target.connectionId); if (!config) throw new Error("Connection config not found"); const effectiveDbType = effectiveDatabaseTypeForConnection(config); + const identifierQuote = connectionStore.connectionIdentifierQuote?.(target.connectionId); const querySchema = metadataSchemaForConnection(config, target.database, target.schema); const targetTableType = target.tableType ?? "TABLE"; if (config.db_type === "neo4j") { @@ -63,6 +64,7 @@ async function openTableTarget(target: NavigationTarget, options: { tableInfoTab const primaryKeys = editableRowIdentifierColumns(effectiveDbType, columns, undefined, targetTableType); const sql = await buildTableSelectSql({ databaseType: effectiveDbType, + identifierQuote, schema: target.schema, catalog: target.catalog, tableName: target.tableName, @@ -86,6 +88,7 @@ async function openTableTarget(target: NavigationTarget, options: { tableInfoTab } const sql = await buildTableSelectSql({ databaseType: effectiveDbType, + identifierQuote, schema: target.schema, catalog: target.catalog, tableName: target.tableName, @@ -113,6 +116,7 @@ async function openTableTarget(target: NavigationTarget, options: { tableInfoTab if (fellBackToLimitZero) { const emptySql = await buildTableSelectSql({ databaseType: effectiveDbType, + identifierQuote, schema: target.schema, catalog: target.catalog, tableName: target.tableName, @@ -139,6 +143,7 @@ async function openTableTarget(target: NavigationTarget, options: { tableInfoTab if (!fellBackToLimitZero && (useRowId || config.db_type === "tdengine")) { const newSql = await buildTableSelectSql({ databaseType: effectiveDbType, + identifierQuote, schema: target.schema, catalog: target.catalog, tableName: target.tableName, diff --git a/apps/desktop/src/lib/__tests__/table/tableSelectSql.spec.ts b/apps/desktop/src/lib/__tests__/table/tableSelectSql.spec.ts index e3b82f1640..f04eced004 100644 --- a/apps/desktop/src/lib/__tests__/table/tableSelectSql.spec.ts +++ b/apps/desktop/src/lib/__tests__/table/tableSelectSql.spec.ts @@ -1,5 +1,5 @@ import { describe, expect, it } from "vitest"; -import { qualifiedTableName, quoteTableIdentifier } from "@/lib/table/tableSelectSql"; +import { qualifiedTableName, quoteTableDataIdentifier, quoteTableIdentifier } from "@/lib/table/tableSelectSql"; describe("qualifiedTableName — Doris/StarRocks multi-catalog", () => { it("prefixes external catalog for Doris (no schema)", () => { @@ -36,4 +36,17 @@ describe("quoteTableIdentifier", () => { it("bracket-quotes sqlserver identifiers", () => { expect(quoteTableIdentifier("sqlserver", "orders")).toBe("[orders]"); }); + + it("uses the connection-reported quote for Kingbase table-data identifiers", () => { + expect(quoteTableDataIdentifier("kingbase", "order", "`")).toBe("`order`"); + expect(quoteTableDataIdentifier("kingbase", "MixedCase", '"')).toBe('"MixedCase"'); + expect(quoteTableDataIdentifier("kingbase", "order detail", "`")).toBe("`order detail`"); + }); + + it("escapes Kingbase identifiers without maintaining a reserved-word list", () => { + expect(quoteTableDataIdentifier("kingbase", "ANALYZE", "`")).toBe("`ANALYZE`"); + expect(quoteTableDataIdentifier("kingbase", "AUTHORIZATION", '"')).toBe('"AUTHORIZATION"'); + expect(quoteTableDataIdentifier("kingbase", "COLLATE", "`")).toBe("`COLLATE`"); + expect(quoteTableDataIdentifier("kingbase", "a`b", "`")).toBe("`a``b`"); + }); }); diff --git a/apps/desktop/src/lib/backend/api.ts b/apps/desktop/src/lib/backend/api.ts index de8735d9a4..55334841f3 100644 --- a/apps/desktop/src/lib/backend/api.ts +++ b/apps/desktop/src/lib/backend/api.ts @@ -55,6 +55,7 @@ export const connectDb = forward("connectDb"); export const connectionFinalProxyPort = forward("connectionFinalProxyPort"); export const disconnectDb = forward("disconnectDb"); export const checkConnectionHealth = forward("checkConnectionHealth"); +export const connectionIdentifierQuote = forward("connectionIdentifierQuote"); export const closeDatabaseConnection = forward("closeDatabaseConnection"); export const refreshConnections = forward("refreshConnections"); export const saveConnections = forward("saveConnections"); diff --git a/apps/desktop/src/lib/backend/http.ts b/apps/desktop/src/lib/backend/http.ts index b62bfd1837..6f9aeed05a 100644 --- a/apps/desktop/src/lib/backend/http.ts +++ b/apps/desktop/src/lib/backend/http.ts @@ -225,6 +225,11 @@ export async function checkConnectionHealth(connectionId: string): Promise return post("/api/connection/check-health", { connectionId }); } +export async function connectionIdentifierQuote(connectionId: string, database?: string): Promise { + const quote = await post("/api/connection/identifier-quote", { connectionId, database }); + return quote ?? undefined; +} + export async function closeDatabaseConnection(connectionId: string, database: string): Promise { return post("/api/connection/close-database", { connectionId, database }); } diff --git a/apps/desktop/src/lib/backend/tauri.ts b/apps/desktop/src/lib/backend/tauri.ts index 7a1d2827aa..edef4c8993 100644 --- a/apps/desktop/src/lib/backend/tauri.ts +++ b/apps/desktop/src/lib/backend/tauri.ts @@ -652,6 +652,11 @@ export async function checkConnectionHealth(connectionId: string): Promise return invoke("check_connection_health", { connectionId }); } +export async function connectionIdentifierQuote(connectionId: string, database?: string): Promise { + const quote = await invoke("connection_identifier_quote", { connectionId, database }); + return quote ?? undefined; +} + export async function closeDatabaseConnection(connectionId: string, database: string): Promise { return invoke("close_database_connection", { connectionId, database }); } diff --git a/apps/desktop/src/lib/dataGrid/dataGridSql.ts b/apps/desktop/src/lib/dataGrid/dataGridSql.ts index 8c3108b02f..6d7782a24e 100644 --- a/apps/desktop/src/lib/dataGrid/dataGridSql.ts +++ b/apps/desktop/src/lib/dataGrid/dataGridSql.ts @@ -93,6 +93,7 @@ export interface DataGridColumnDistinctValuesSqlOptions { export interface DataGridCountSqlOptions { databaseType?: DatabaseType; + identifierQuote?: string; catalog?: string; schema?: string; tableName: string; diff --git a/apps/desktop/src/lib/export/databaseExport.ts b/apps/desktop/src/lib/export/databaseExport.ts index 5ab8183ed4..1afca53932 100644 --- a/apps/desktop/src/lib/export/databaseExport.ts +++ b/apps/desktop/src/lib/export/databaseExport.ts @@ -46,6 +46,7 @@ export interface BuildExportInsertStatementsOptions { export interface BuildExportPageSqlOptions { databaseType?: DatabaseType; + identifierQuote?: string; schema?: string; tableName: string; limit?: number; @@ -72,6 +73,7 @@ export function buildInsertStatements(options: BuildExportInsertStatementsOption export async function buildExportPageSql(options: BuildExportPageSqlOptions): Promise { return buildTableSelectSql({ databaseType: options.databaseType, + identifierQuote: options.identifierQuote, schema: options.schema, tableName: options.tableName, limit: options.limit ?? DATABASE_EXPORT_PAGE_SIZE, diff --git a/apps/desktop/src/lib/table/tableDataExport.ts b/apps/desktop/src/lib/table/tableDataExport.ts index 06619062df..9191d71270 100644 --- a/apps/desktop/src/lib/table/tableDataExport.ts +++ b/apps/desktop/src/lib/table/tableDataExport.ts @@ -5,6 +5,7 @@ export const TABLE_DATA_EXPORT_PAGE_SIZE = 10_000; export interface FetchTableDataForExportOptions { databaseType?: DatabaseType; + identifierQuote?: string; schema?: string; tableName: string; tableType?: string; @@ -24,6 +25,7 @@ export async function fetchTableDataForExport(options: FetchTableDataForExportOp while (true) { const sql = await (options.buildPageSql ?? buildTableSelectSql)({ databaseType: options.databaseType, + identifierQuote: options.identifierQuote, schema: options.schema, tableName: options.tableName, tableType: options.tableType, diff --git a/apps/desktop/src/lib/table/tableSelectSql.ts b/apps/desktop/src/lib/table/tableSelectSql.ts index e56238bcb5..dc557ec3e8 100644 --- a/apps/desktop/src/lib/table/tableSelectSql.ts +++ b/apps/desktop/src/lib/table/tableSelectSql.ts @@ -5,6 +5,7 @@ import { parseSqlServerLinkedSchema, sqlServerLinkedTableName } from "@/lib/data export interface BuildTableSelectSqlOptions { databaseType?: DatabaseType; + identifierQuote?: string; schema?: string; tableName: string; tableType?: string; @@ -32,6 +33,14 @@ export function quoteTableIdentifier(databaseType: DatabaseType | undefined, nam return `"${name.replace(/"/g, '""')}"`; } +export function quoteTableDataIdentifier(databaseType: DatabaseType | undefined, name: string, identifierQuote?: string): string { + if (databaseType === "kingbase" && identifierQuote != null) { + if (!identifierQuote) return name; + return `${identifierQuote}${name.replaceAll(identifierQuote, identifierQuote + identifierQuote)}${identifierQuote}`; + } + return quoteTableIdentifier(databaseType, name); +} + function quoteCypherIdentifier(name: string): string { return `\`${name.replace(/`/g, "``")}\``; } diff --git a/apps/desktop/src/stores/__tests__/queryStore.hiddenPrimaryKey.spec.ts b/apps/desktop/src/stores/__tests__/queryStore.hiddenPrimaryKey.spec.ts index ec2b40e6e9..4e99c4011f 100644 --- a/apps/desktop/src/stores/__tests__/queryStore.hiddenPrimaryKey.spec.ts +++ b/apps/desktop/src/stores/__tests__/queryStore.hiddenPrimaryKey.spec.ts @@ -462,6 +462,7 @@ describe("queryStore hidden primary key editing", () => { expect(buildDataGridCountSql).toHaveBeenCalledWith({ databaseType: "mysql", + identifierQuote: undefined, catalog: undefined, schema: "public", tableName: "users", diff --git a/apps/desktop/src/stores/__tests__/queryStore.tableDataRefresh.spec.ts b/apps/desktop/src/stores/__tests__/queryStore.tableDataRefresh.spec.ts index 4e7c07dfb3..816780db46 100644 --- a/apps/desktop/src/stores/__tests__/queryStore.tableDataRefresh.spec.ts +++ b/apps/desktop/src/stores/__tests__/queryStore.tableDataRefresh.spec.ts @@ -105,6 +105,7 @@ describe("queryStore table data refresh", () => { expect(refreshed).toBe(1); expect(mocks.buildTableSelectSql).toHaveBeenCalledWith({ databaseType: "postgres", + identifierQuote: undefined, schema: "public", tableName: "users", tableType: "TABLE", diff --git a/apps/desktop/src/stores/connectionStore.ts b/apps/desktop/src/stores/connectionStore.ts index 6971f15a6e..0c136155f4 100644 --- a/apps/desktop/src/stores/connectionStore.ts +++ b/apps/desktop/src/stores/connectionStore.ts @@ -233,6 +233,7 @@ export const useConnectionStore = defineStore("connection", () => { const treeNodes = ref([]); const pinnedTreeNodeIds = ref>(new Set()); const connectedIds = ref>(new Set()); + const identifierQuotes = ref>({}); const lastConnectionHealthCheckAt = ref>({}); const agentDrivers = ref([]); let agentDriversRefreshPromise: Promise | null = null; @@ -385,6 +386,25 @@ export const useConnectionStore = defineStore("connection", () => { return configById.value.get(connectionId); } + function connectionIdentifierQuote(connectionId?: string): string | undefined { + if (!connectionId) return undefined; + return identifierQuotes.value[connectionId]; + } + + function clearConnectionIdentifierQuote(connectionId: string) { + if (!(connectionId in identifierQuotes.value)) return; + const next = { ...identifierQuotes.value }; + delete next[connectionId]; + identifierQuotes.value = next; + } + + async function refreshConnectionIdentifierQuote(connectionId: string, config: ConnectionConfig) { + clearConnectionIdentifierQuote(connectionId); + if (config.db_type !== "kingbase") return; + const quote = await api.connectionIdentifierQuote(connectionId).catch(() => undefined); + if (quote != null) identifierQuotes.value = { ...identifierQuotes.value, [connectionId]: quote }; + } + function connectionErrorMessage(error: unknown): string { if (error instanceof Error) return error.message; return String(error); @@ -711,6 +731,7 @@ export const useConnectionStore = defineStore("connection", () => { function markConnectionLost(connectionId: string, error: unknown) { connectedIds.value.delete(connectionId); + clearConnectionIdentifierQuote(connectionId); clearConnectionNodeLoading(connectionId); clearConnectionHealthCheck(connectionId); if (activeConnectionId.value === connectionId) activeConnectionId.value = null; @@ -1629,6 +1650,7 @@ export const useConnectionStore = defineStore("connection", () => { for (const id of removedIds) { clearConnectionError(id); connectedIds.value.delete(id); + clearConnectionIdentifierQuote(id); clearConnectionHealthCheck(id); sidebarLayout.value = removeConnectionFromSidebarLayout(sidebarLayout.value, id); } @@ -1660,6 +1682,7 @@ export const useConnectionStore = defineStore("connection", () => { connections.value = nextConnections; rebuildTreeNodes(); connectedIds.value.delete(config.id); + clearConnectionIdentifierQuote(config.id); clearConnectionHealthCheck(config.id); invalidateCompletionCache(config.id); clearLoadedChildrenCache(config.id); @@ -1853,6 +1876,7 @@ export const useConnectionStore = defineStore("connection", () => { await ensureLocalConnectionAttemptActiveAfterConnectResult(config.id, localAttempt, id); activeConnectionId.value = id; connectedIds.value.add(id); + await refreshConnectionIdentifierQuote(id, { ...config, id }); if (id !== config.id) markSuccessfulLocalConnectionAttempt(config.id, localAttempt); markSuccessfulLocalConnectionAttempt(id, localAttempt); markConnectionHealthChecked(id); @@ -1900,6 +1924,7 @@ export const useConnectionStore = defineStore("connection", () => { if (!cancelled) return false; clearConnectionError(connectionId); connectedIds.value.delete(connectionId); + clearConnectionIdentifierQuote(connectionId); clearConnectionHealthCheck(connectionId); if (activeConnectionId.value === connectionId) activeConnectionId.value = null; invalidateCompletionCache(connectionId); @@ -1914,6 +1939,7 @@ export const useConnectionStore = defineStore("connection", () => { const shouldRemoveOneTimeConnection = getConfig(connectionId)?.one_time === true; connectedIds.value.delete(connectionId); + clearConnectionIdentifierQuote(connectionId); forgetSuccessfulLocalConnectionAttempt(connectionId); clearConnectionHealthCheck(connectionId); const node = findNode(treeNodes.value, connectionId); @@ -2014,6 +2040,7 @@ export const useConnectionStore = defineStore("connection", () => { await syncMongoLegacyDriverFallback(connectionId, config); await ensureLocalConnectionAttemptActiveAfterConnectResult(connectionId, localAttempt, id); connectedIds.value.add(connectionId); + await refreshConnectionIdentifierQuote(connectionId, config); markSuccessfulLocalConnectionAttempt(connectionId, localAttempt); markConnectionHealthChecked(connectionId); activeConnectionId.value = connectionId; @@ -5089,6 +5116,7 @@ export const useConnectionStore = defineStore("connection", () => { recordConnectionLostError, sidebarLayout, getConfig, + connectionIdentifierQuote, isTreeNodePinned, toggleTreeNodePin, addConnection, diff --git a/apps/desktop/src/stores/queryStore.ts b/apps/desktop/src/stores/queryStore.ts index 7acc17e3b9..b6b7dcbe65 100644 --- a/apps/desktop/src/stores/queryStore.ts +++ b/apps/desktop/src/stores/queryStore.ts @@ -34,7 +34,7 @@ import { tableMetaForDataTab } from "@/lib/table/tableDataTabMeta"; import { dataTabExecutionDatabase } from "@/lib/table/dataTabExecutionDatabase"; import { tableOpenPageLimit } from "@/lib/table/tableOpenPageLimit"; import { loadTableMetadata } from "@/lib/metadata/tableMetadataCache"; -import { buildTableSelectSql, quoteTableIdentifier } from "@/lib/table/tableSelectSql"; +import { buildTableSelectSql, quoteTableDataIdentifier } from "@/lib/table/tableSelectSql"; import { connectionQueryExecutionSchema, effectiveDatabaseTypeForConnection, metadataSchemaForConnection } from "@/lib/database/jdbcDialect"; import { frontendQueryTimeoutSecsForSql, queryTimeoutSecsForConnection } from "@/lib/sql/queryTimeout"; import { queryResultSourceLabel } from "@/lib/sql/queryResultSource"; @@ -1625,15 +1625,18 @@ export const useQueryStore = defineStore("query", () => { for (const tab of matchingTabs) { const tableMeta = tableMetaForDataTab(tab); if (!tableMeta?.tableName) continue; - const conn = useConnectionStore().getConfig(tab.connectionId); + const connStore = useConnectionStore(); + const conn = connStore.getConfig(tab.connectionId); const effectiveDbType = effectiveDatabaseTypeForConnection(conn); + const identifierQuote = connStore.connectionIdentifierQuote?.(tab.connectionId); const primaryKeys = tab.tableMeta ? tab.tableMeta.primaryKeys : tableMeta.primaryKeys; - const sortOrder = tab.resultSortColumn && tab.resultSortDirection ? `${quoteTableIdentifier(effectiveDbType, tab.resultSortColumn)} ${tab.resultSortDirection.toUpperCase()}` : undefined; + const sortOrder = tab.resultSortColumn && tab.resultSortDirection ? `${quoteTableDataIdentifier(effectiveDbType, tab.resultSortColumn, identifierQuote)} ${tab.resultSortDirection.toUpperCase()}` : undefined; const orderBy = tab.orderByInput?.trim() || sortOrder; const limit = tab.resultPageLimit ?? settingsStore.editorSettings.pageSize ?? tableOpenPageLimit(); const offset = tab.resultPageOffset ?? 0; const sql = await buildTableSelectSql({ databaseType: effectiveDbType, + identifierQuote, schema: tableMeta.schema, tableName: tableMeta.tableName, tableType: tableMeta.tableType, @@ -2967,6 +2970,7 @@ export const useQueryStore = defineStore("query", () => { if (!tableMeta?.tableName) return undefined; return { databaseType: effectiveDbType, + identifierQuote: useConnectionStore().connectionIdentifierQuote?.(current.connectionId), catalog: tableMeta.catalog, schema: tableMeta.schema, tableName: tableMeta.tableName, @@ -3578,8 +3582,9 @@ export const useQueryStore = defineStore("query", () => { const totalRows = typeof tab.resultTotalRowCount === "number" ? tab.resultTotalRowCount : null; const pageLimit = TABLE_DATA_EXPORT_PAGE_SIZE; const effectiveDbType = effectiveDatabaseTypeForConnection(conn); + const identifierQuote = connStore.connectionIdentifierQuote?.(tab.connectionId); const primaryKeys = tab.tableMeta ? tab.tableMeta.primaryKeys : tableMeta.primaryKeys; - const sortOrder = tab.resultSortColumn && tab.resultSortDirection ? `${quoteTableIdentifier(effectiveDbType, tab.resultSortColumn)} ${tab.resultSortDirection.toUpperCase()}` : undefined; + const sortOrder = tab.resultSortColumn && tab.resultSortDirection ? `${quoteTableDataIdentifier(effectiveDbType, tab.resultSortColumn, identifierQuote)} ${tab.resultSortDirection.toUpperCase()}` : undefined; const orderBy = tab.orderByInput?.trim() || sortOrder; const queryTimeoutSecs = queryTimeoutSecsForConnection(conn); const executionDatabase = dataTabExecutionDatabase(conn, tab.database, tableMeta.catalog); @@ -3594,6 +3599,7 @@ export const useQueryStore = defineStore("query", () => { while (true) { const sql = await api.buildTableSelectSql({ databaseType: effectiveDbType, + identifierQuote, schema: tableMeta.schema, tableName: tableMeta.tableName, tableType: tableMeta.tableType, diff --git a/crates/dbx-core/assets/agent-protocol-v1.json b/crates/dbx-core/assets/agent-protocol-v1.json index 29d9ca4359..db51f9650f 100644 --- a/crates/dbx-core/assets/agent-protocol-v1.json +++ b/crates/dbx-core/assets/agent-protocol-v1.json @@ -39,6 +39,7 @@ "connect", "test_connection", "validate_connection", + "connection_info", "list_databases", "list_schemas", "list_tables", diff --git a/crates/dbx-core/assets/agent-protocol-v2.json b/crates/dbx-core/assets/agent-protocol-v2.json index 33d2400f72..08900d5d76 100644 --- a/crates/dbx-core/assets/agent-protocol-v2.json +++ b/crates/dbx-core/assets/agent-protocol-v2.json @@ -46,6 +46,7 @@ "cancel_session", "test_connection", "validate_connection", + "connection_info", "list_databases", "list_schemas", "list_tables", diff --git a/crates/dbx-core/src/connection.rs b/crates/dbx-core/src/connection.rs index 0b83da55cd..b3c0ef3fd4 100644 --- a/crates/dbx-core/src/connection.rs +++ b/crates/dbx-core/src/connection.rs @@ -2317,6 +2317,35 @@ impl AppState { keys } + pub async fn connection_identifier_quote( + &self, + connection_id: &str, + database: Option<&str>, + ) -> Result, String> { + let config = self + .configs + .read() + .await + .get(connection_id) + .cloned() + .ok_or_else(|| format!("Connection config not found: {connection_id}"))?; + if !database_capabilities::is_agent_type(&config.db_type) { + return Ok(None); + } + + let pool_key = self.get_or_create_pool(connection_id, database).await?; + let client = { + let connections = self.connections.read().await; + match connections.get(&pool_key) { + Some(PoolKind::Agent(client)) => client.clone(), + _ => return Ok(None), + } + }; + let mut agent = client.lock().await; + let info = agent.connection_info(Some(db::connection_timeout())).await?; + Ok(Some(info.identifier_quote)) + } + pub async fn reset_connection_transport(&self, connection_id: &str) { let layer_count = { let configs = self.configs.read().await; diff --git a/crates/dbx-core/src/data_grid_sql.rs b/crates/dbx-core/src/data_grid_sql.rs index 7248501526..14eeec6fa1 100644 --- a/crates/dbx-core/src/data_grid_sql.rs +++ b/crates/dbx-core/src/data_grid_sql.rs @@ -217,6 +217,8 @@ pub struct DataGridCountSqlOptions { #[serde(default, skip_serializing_if = "Option::is_none")] pub database_type: Option, #[serde(default, skip_serializing_if = "Option::is_none")] + pub identifier_quote: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] pub catalog: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub schema: Option, @@ -721,12 +723,21 @@ pub fn build_data_grid_column_distinct_values_sql(options: DataGridColumnDistinc } pub fn build_data_grid_count_sql(options: DataGridCountSqlOptions) -> String { - let table = crate::sql_dialect::qualified_table_name_with_catalog( - options.database_type, - options.catalog.as_deref(), - options.schema.as_deref(), - &options.table_name, - ); + let table = if options.database_type == Some(DatabaseType::Kingbase) { + crate::sql_dialect::table_data_qualified_table_name( + options.database_type, + options.schema.as_deref(), + &options.table_name, + options.identifier_quote.as_deref(), + ) + } else { + crate::sql_dialect::qualified_table_name_with_catalog( + options.database_type, + options.catalog.as_deref(), + options.schema.as_deref(), + &options.table_name, + ) + }; let predicate = crate::sql_dialect::normalize_where_input(options.where_input.as_deref()); let where_clause = if predicate.is_empty() { String::new() } else { format!(" WHERE ({predicate})") }; format!("SELECT COUNT(*) AS cnt FROM {table}{where_clause}") @@ -3060,6 +3071,7 @@ mod tests { assert_eq!( build_data_grid_count_sql(DataGridCountSqlOptions { database_type: Some(DatabaseType::Postgres), + identifier_quote: None, catalog: None, schema: Some("public".to_string()), table_name: "users".to_string(), @@ -3070,6 +3082,7 @@ mod tests { assert_eq!( build_data_grid_count_sql(DataGridCountSqlOptions { database_type: Some(DatabaseType::Doris), + identifier_quote: None, catalog: Some("iceberg_catalog".to_string()), schema: Some("sales".to_string()), table_name: "orders".to_string(), @@ -3080,6 +3093,7 @@ mod tests { assert_eq!( build_data_grid_count_sql(DataGridCountSqlOptions { database_type: Some(DatabaseType::StarRocks), + identifier_quote: None, catalog: Some("hive_catalog".to_string()), schema: None, table_name: "orders".to_string(), @@ -3087,6 +3101,17 @@ mod tests { }), "SELECT COUNT(*) AS cnt FROM `hive_catalog`.`orders`" ); + assert_eq!( + build_data_grid_count_sql(DataGridCountSqlOptions { + database_type: Some(DatabaseType::Kingbase), + identifier_quote: Some("`".to_string()), + catalog: None, + schema: Some("cqbq_ls".to_string()), + table_name: "ANALYZE".to_string(), + where_input: None, + }), + "SELECT COUNT(*) AS cnt FROM `cqbq_ls`.`ANALYZE`" + ); } #[test] diff --git a/crates/dbx-core/src/db/agent_driver.rs b/crates/dbx-core/src/db/agent_driver.rs index be1dc3bf40..47a6532136 100644 --- a/crates/dbx-core/src/db/agent_driver.rs +++ b/crates/dbx-core/src/db/agent_driver.rs @@ -373,6 +373,7 @@ pub enum AgentMethod { CancelSession, TestConnection, ValidateConnection, + ConnectionInfo, ListDatabases, ListSchemas, ListTables, @@ -400,7 +401,7 @@ pub enum AgentMethod { } impl AgentMethod { - pub const ALL: [Self; 32] = [ + pub const ALL: [Self; 33] = [ Self::Handshake, Self::Connect, Self::OpenSession, @@ -409,6 +410,7 @@ impl AgentMethod { Self::CancelSession, Self::TestConnection, Self::ValidateConnection, + Self::ConnectionInfo, Self::ListDatabases, Self::ListSchemas, Self::ListTables, @@ -445,6 +447,7 @@ impl AgentMethod { Self::CancelSession => "cancel_session", Self::TestConnection => "test_connection", Self::ValidateConnection => "validate_connection", + Self::ConnectionInfo => "connection_info", Self::ListDatabases => "list_databases", Self::ListSchemas => "list_schemas", Self::ListTables => "list_tables", @@ -473,6 +476,12 @@ impl AgentMethod { } } +#[derive(Debug, Clone, PartialEq, Eq, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct AgentConnectionInfo { + pub identifier_quote: String, +} + #[derive(Debug, Clone, Serialize)] #[serde(rename_all = "camelCase")] pub struct AgentTableReadStartParams { @@ -910,6 +919,10 @@ impl AgentDriverClient { self.call_method_with_timeout(AgentMethod::ValidateConnection, serde_json::json!({}), timeout_duration).await } + pub async fn connection_info(&mut self, timeout_duration: Option) -> Result { + self.call_method_with_timeout(AgentMethod::ConnectionInfo, serde_json::json!({}), timeout_duration).await + } + pub async fn disconnect(&mut self) -> Result { if self.shared_runtime.is_some() { let session_id = self.agent_session_id.as_ref().ok_or("Shared Agent session id is missing")?.clone(); diff --git a/crates/dbx-core/src/sql_dialect.rs b/crates/dbx-core/src/sql_dialect.rs index b013050b3b..7c6b21fefd 100644 --- a/crates/dbx-core/src/sql_dialect.rs +++ b/crates/dbx-core/src/sql_dialect.rs @@ -14,5 +14,6 @@ pub use identifiers::{ normalize_where_input, qualified_table_name, qualified_table_name_with_catalog, quote_table_identifier, }; pub(crate) use identifiers::{parse_sqlserver_linked_schema_ref, qualified_transfer_table, quote_transfer_identifier}; +pub(crate) use table_select::table_data_qualified_table_name; pub use table_select::{build_count_table_sql, build_table_data_select_sql, build_table_select_sql}; pub use types::*; diff --git a/crates/dbx-core/src/sql_dialect/table_select.rs b/crates/dbx-core/src/sql_dialect/table_select.rs index b749edbdfd..394272dfc3 100644 --- a/crates/dbx-core/src/sql_dialect/table_select.rs +++ b/crates/dbx-core/src/sql_dialect/table_select.rs @@ -21,12 +21,21 @@ pub fn build_table_data_select_sql(options: TableDataSelectSqlOptions) -> String } // Doris / StarRocks multi-catalog: prefix the catalog for external-catalog tables. - let table = qualified_table_name_with_catalog( - database_type, - options.catalog.as_deref(), - options.schema.as_deref(), - &options.table_name, - ); + let table = if database_type == Some(DatabaseType::Kingbase) { + table_data_qualified_table_name( + database_type, + options.schema.as_deref(), + &options.table_name, + options.identifier_quote.as_deref(), + ) + } else { + qualified_table_name_with_catalog( + database_type, + options.catalog.as_deref(), + options.schema.as_deref(), + &options.table_name, + ) + }; let predicate = normalize_where_input(options.where_input.as_deref()); let where_clause = if predicate.is_empty() { String::new() } else { format!(" WHERE ({predicate})") }; let default_order_by = if database_type == Some(DatabaseType::InfluxDb) { @@ -143,6 +152,40 @@ pub fn build_table_data_select_sql(options: TableDataSelectSqlOptions) -> String } } +pub(crate) fn table_data_qualified_table_name( + database_type: Option, + schema: Option<&str>, + table_name: &str, + identifier_quote: Option<&str>, +) -> String { + if database_type != Some(DatabaseType::Kingbase) { + return qualified_table_name(database_type, schema, table_name); + } + let table = quote_table_data_identifier(database_type, table_name, identifier_quote); + schema + .map(str::trim) + .filter(|schema| !schema.is_empty()) + .map(|schema| format!("{}.{}", quote_table_data_identifier(database_type, schema, identifier_quote), table)) + .unwrap_or(table) +} + +fn quote_table_data_identifier( + database_type: Option, + name: &str, + identifier_quote: Option<&str>, +) -> String { + if database_type != Some(DatabaseType::Kingbase) { + return quote_table_identifier(database_type, name); + } + let Some(quote) = identifier_quote else { + return quote_table_identifier(database_type, name); + }; + if quote.is_empty() { + return name.to_string(); + } + format!("{quote}{}{quote}", name.replace(quote, &format!("{quote}{quote}"))) +} + fn is_view_table_type(table_type: Option<&str>) -> bool { table_type.is_some_and(|value| value.to_ascii_uppercase().contains("VIEW")) } @@ -418,6 +461,7 @@ mod tests { fn opts(database_type: DatabaseType, catalog: Option<&str>, table: &str) -> TableDataSelectSqlOptions { TableDataSelectSqlOptions { database_type: Some(database_type), + identifier_quote: None, schema: None, table_name: table.to_string(), catalog: catalog.map(|c| c.to_string()), diff --git a/crates/dbx-core/src/sql_dialect/tests.rs b/crates/dbx-core/src/sql_dialect/tests.rs index c25cc95df4..369d756084 100644 --- a/crates/dbx-core/src/sql_dialect/tests.rs +++ b/crates/dbx-core/src/sql_dialect/tests.rs @@ -21,6 +21,11 @@ fn quotes_identifiers_by_database_type() { assert_eq!(quote_table_identifier(Some(DatabaseType::StarRocks), "user`name"), "`user``name`"); assert_eq!(quote_table_identifier(Some(DatabaseType::SqlServer), "user]name"), "[user]]name]"); assert_eq!(quote_table_identifier(Some(DatabaseType::Postgres), "user\"name"), "\"user\"\"name\""); + assert_eq!(quote_table_identifier(Some(DatabaseType::Kingbase), "cqbq_ls"), "\"cqbq_ls\""); + assert_eq!(quote_table_identifier(Some(DatabaseType::Kingbase), "actionlogs"), "\"actionlogs\""); + assert_eq!(quote_table_identifier(Some(DatabaseType::Kingbase), "order"), "\"order\""); + assert_eq!(quote_table_identifier(Some(DatabaseType::Kingbase), "MixedCase"), "\"MixedCase\""); + assert_eq!(quote_table_identifier(Some(DatabaseType::Kingbase), "order detail"), "\"order detail\""); assert_eq!(quote_table_identifier(Some(DatabaseType::Informix), "users_1"), "users_1"); assert_eq!(quote_table_identifier(Some(DatabaseType::Jdbc), "users_1"), "users_1"); assert_eq!(quote_table_identifier(Some(DatabaseType::Jdbc), "user name"), "user name"); @@ -31,6 +36,10 @@ fn quotes_identifiers_by_database_type() { fn qualifies_schema_only_for_schema_aware_databases() { assert_eq!(qualified_table_name(Some(DatabaseType::Postgres), Some("public"), "users"), "\"public\".\"users\""); assert_eq!(qualified_table_name(Some(DatabaseType::Kwdb), Some("public"), "users"), "\"public\".\"users\""); + assert_eq!( + qualified_table_name(Some(DatabaseType::Kingbase), Some("cqbq_ls"), "actionlogs"), + "\"cqbq_ls\".\"actionlogs\"" + ); assert_eq!(qualified_table_name(Some(DatabaseType::Mysql), Some("public"), "users"), "`public`.`users`"); assert_eq!(qualified_table_name(Some(DatabaseType::Goldendb), Some("public"), "users"), "`public`.`users`"); assert_eq!(qualified_table_name(Some(DatabaseType::StarRocks), Some("warehouse"), "users"), "`warehouse`.`users`"); @@ -244,6 +253,10 @@ fn builds_select_sql_with_limit_syntax_for_database_type() { #[test] fn builds_table_data_where_and_schema_queries() { + assert_eq!( + build_count_table_sql(Some(DatabaseType::Kingbase), Some("cqbq_ls"), "actionlogs"), + "SELECT COUNT(*) AS row_count FROM \"cqbq_ls\".\"actionlogs\"" + ); assert_eq!( build_table_data_select_sql(TableDataSelectSqlOptions { database_type: Some(DatabaseType::Mysql), @@ -298,6 +311,36 @@ fn builds_table_data_where_and_schema_queries() { }), "SELECT * FROM \"public\".\"orders\" WHERE (amount > 10) LIMIT 50 OFFSET 100;" ); + assert_eq!( + build_table_data_select_sql(TableDataSelectSqlOptions { + database_type: Some(DatabaseType::Kingbase), + identifier_quote: Some("`".to_string()), + schema: Some("cqbq_ls".to_string()), + table_name: "actionlogs".to_string(), + table_type: None, + primary_keys: Vec::new(), + columns: Vec::new(), + fallback_order_columns: Vec::new(), + order_by: None, + limit: Some(100), + offset: None, + where_input: None, + include_row_id: false, + ..Default::default() + }), + "SELECT * FROM `cqbq_ls`.`actionlogs` LIMIT 100;" + ); + assert_eq!( + build_table_data_select_sql(TableDataSelectSqlOptions { + database_type: Some(DatabaseType::Kingbase), + identifier_quote: Some("\"".to_string()), + schema: Some("App Schema".to_string()), + table_name: "ANALYZE".to_string(), + limit: Some(100), + ..Default::default() + }), + "SELECT * FROM \"App Schema\".\"ANALYZE\" LIMIT 100;" + ); assert_eq!( build_table_data_select_sql(TableDataSelectSqlOptions { database_type: Some(DatabaseType::Xugu), diff --git a/crates/dbx-core/src/sql_dialect/types.rs b/crates/dbx-core/src/sql_dialect/types.rs index 2f3b523111..21aca4629d 100644 --- a/crates/dbx-core/src/sql_dialect/types.rs +++ b/crates/dbx-core/src/sql_dialect/types.rs @@ -22,6 +22,8 @@ pub struct TableDataSelectSqlOptions { #[serde(default, skip_serializing_if = "Option::is_none")] pub database_type: Option, #[serde(default, skip_serializing_if = "Option::is_none")] + pub identifier_quote: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] pub schema: Option, pub table_name: String, /// Doris / StarRocks multi-catalog: when set to a non-`internal` catalog, diff --git a/crates/dbx-web/src/main.rs b/crates/dbx-web/src/main.rs index 4e327a6dfc..28a615d099 100644 --- a/crates/dbx-web/src/main.rs +++ b/crates/dbx-web/src/main.rs @@ -222,6 +222,7 @@ async fn main() { .route("/connection/final-proxy-port", post(routes::connection::connection_final_proxy_port)) .route("/connection/disconnect", post(routes::connection::disconnect_db)) .route("/connection/check-health", post(routes::connection::check_connection_health)) + .route("/connection/identifier-quote", post(routes::connection::connection_identifier_quote)) .route("/connection/close-database", post(routes::connection::close_database_connection)) .route("/connection/save", post(routes::connection::save_connections)) .route("/connection/list", get(routes::connection::load_connections)) diff --git a/crates/dbx-web/src/routes/connection.rs b/crates/dbx-web/src/routes/connection.rs index 714fd43a34..ca1cb37ddd 100644 --- a/crates/dbx-web/src/routes/connection.rs +++ b/crates/dbx-web/src/routes/connection.rs @@ -30,6 +30,13 @@ pub struct CloseDatabaseConnectionRequest { pub database: String, } +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct ConnectionIdentifierQuoteRequest { + pub connection_id: String, + pub database: Option, +} + #[derive(Deserialize)] #[serde(rename_all = "camelCase")] pub struct SaveConnectionsRequest { @@ -133,6 +140,18 @@ pub async fn check_connection_health( Ok(Json(())) } +pub async fn connection_identifier_quote( + State(state): State>, + Json(body): Json, +) -> Result>, AppError> { + state + .app + .connection_identifier_quote(&body.connection_id, body.database.as_deref()) + .await + .map(Json) + .map_err(AppError) +} + pub async fn close_database_connection( State(state): State>, Json(body): Json, diff --git a/src-tauri/src/commands/connection.rs b/src-tauri/src/commands/connection.rs index 5e48a7981d..a08090a32d 100644 --- a/src-tauri/src/commands/connection.rs +++ b/src-tauri/src/commands/connection.rs @@ -1245,6 +1245,15 @@ pub async fn check_connection_health(state: State<'_, Arc>, connection state.check_connection_health(&connection_id).await } +#[tauri::command] +pub async fn connection_identifier_quote( + state: State<'_, Arc>, + connection_id: String, + database: Option, +) -> Result, String> { + state.connection_identifier_quote(&connection_id, database.as_deref()).await +} + /// Check whether a connection has read-only protection enabled. /// Returns an error if the connection is read-only, preventing write operations. pub async fn ensure_connection_writable( diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index 78daad4259..6af956c05e 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -915,6 +915,7 @@ pub fn run() { commands::connection::close_database_connection, commands::connection::refresh_connections, commands::connection::check_connection_health, + commands::connection::connection_identifier_quote, commands::connection::save_connections, commands::connection::load_connections, commands::connection::save_sidebar_layout,