From 025a4fdcce2da1905154389839741dbe018e4ca7 Mon Sep 17 00:00:00 2001
From: t8y2 <1156263951@qq.com>
Date: Tue, 9 Jun 2026 00:44:04 +0800
Subject: [PATCH] fix: support tree-schema default DB selection in AI assistant

---
 .../src/components/editor/AiAssistant.vue     | 54 +++++++++++++++----
 .../src/components/layout/EditorToolbar.vue   | 12 ++---
 apps/desktop/src/lib/defaultDatabase.ts       | 27 +++++++++-
 packages/app-tests/databaseOptions.test.ts    |  2 +-
 packages/app-tests/defaultDatabase.test.ts    | 37 ++++++++++++-
 5 files changed, 112 insertions(+), 20 deletions(-)

diff --git a/apps/desktop/src/components/editor/AiAssistant.vue b/apps/desktop/src/components/editor/AiAssistant.vue
index 4cac802fb..bc755fa72 100644
--- a/apps/desktop/src/components/editor/AiAssistant.vue
+++ b/apps/desktop/src/components/editor/AiAssistant.vue
@@ -61,7 +61,12 @@ import {
 import type { AiMessage } from "@/lib/api";
 import type { ConnectionConfig, QueryTab, TableInfo } from "@/types/database";
 import { useDatabaseOptions } from "@/composables/useDatabaseOptions";
-import { resolveDefaultDatabase } from "@/lib/defaultDatabase";
+import {
+  decodeSelectableDatabaseValue,
+  encodeSelectableDatabaseValue,
+  formatDatabaseLabel,
+  resolveDefaultDatabase,
+} from "@/lib/defaultDatabase";
 import { isSchemaAware } from "@/lib/databaseCapabilities";
 import { copyToClipboard } from "@/lib/clipboard";
 import { formatAiTableMention, parseAiTableMentions, type AiTableMention } from "@/lib/aiTableMentions";
@@ -191,6 +196,32 @@ const dbOptions = computed(() => {
   return allDbOptions.value[props.connection.id] || [];
 });
 
+const dbSelectOptions = computed(() => {
+  const connection = props.connection;
+  if (!connection) return [];
+  return dbOptions.value.map((database) => ({
+    database,
+    value: encodeSelectableDatabaseValue(connection.db_type, database),
+    label: formatDatabaseLabel(connection, database, {
+      defaultDatabase: t("editor.defaultDatabase"),
+      noDatabase: t("editor.noDatabase"),
+    }),
+  }));
+});
+
+const selectedDatabaseSelectValue = computed(() =>
+  props.connection ? encodeSelectableDatabaseValue(props.connection.db_type, props.tab?.database || "") : "",
+);
+
+const selectedDatabaseLabel = computed(() => {
+  if (!props.connection) return t("editor.selectDatabase");
+  if (!props.tab) return t("editor.selectDatabase");
+  return formatDatabaseLabel(props.connection, props.tab.database || "", {
+    defaultDatabase: t("editor.defaultDatabase"),
+    noDatabase: t("editor.noDatabase"),
+  });
+});
+
 async function loadDatabases() {
   if (!props.connection) return;
   await loadDatabaseOptions(props.connection.id);
@@ -217,10 +248,11 @@ async function changeConnection(connectionId: string) {
   }
 }
 
-function changeDatabase(database: string) {
+function changeDatabase(value: string) {
   const tab = props.tab;
-  if (!tab) return;
-  queryStore.updateDatabase(tab.id, database);
+  const connection = props.connection;
+  if (!tab || !connection) return;
+  queryStore.updateDatabase(tab.id, decodeSelectableDatabaseValue(connection.db_type, value));
 }
 
 function appendAssistantDelta(assistantIdx: number, delta: string) {
@@ -885,7 +917,7 @@ const messageRenderer = computed(() => {
           <template v-if="connection">
             <Database class="h-3 w-3 shrink-0 text-foreground/40" />
             <Select
-              :model-value="tab?.database || ''"
+              :model-value="selectedDatabaseSelectValue"
               @update:model-value="(v: any) => changeDatabase(v)"
               @update:open="
                 (open: boolean) => {
@@ -896,14 +928,14 @@ const messageRenderer = computed(() => {
               <SelectTrigger
                 class="h-5 w-auto border-0 rounded-md bg-transparent dark:bg-transparent p-0 px-1 text-xs text-foreground/80 shadow-none focus:ring-0 focus-visible:ring-0 [&_svg]:size-3"
               >
-                <SelectValue :placeholder="t('editor.selectDatabase')">{{
-                  tab?.database || t("editor.selectDatabase")
-                }}</SelectValue>
+                <SelectValue :placeholder="t('editor.selectDatabase')">{{ selectedDatabaseLabel }}</SelectValue>
               </SelectTrigger>
               <SelectContent>
-                <SelectItem v-for="db in dbOptions" :key="db" :value="db">{{ db }}</SelectItem>
-                <SelectItem v-if="!dbOptions.length && tab?.database" :value="tab.database">{{
-                  tab.database
+                <SelectItem v-for="option in dbSelectOptions" :key="option.value" :value="option.value">{{
+                  option.label
+                }}</SelectItem>
+                <SelectItem v-if="!dbSelectOptions.length && connection && tab" :value="selectedDatabaseSelectValue">{{
+                  selectedDatabaseLabel
                 }}</SelectItem>
               </SelectContent>
             </Select>
diff --git a/apps/desktop/src/components/layout/EditorToolbar.vue b/apps/desktop/src/components/layout/EditorToolbar.vue
index 9102f6d25..e97b7b014 100644
--- a/apps/desktop/src/components/layout/EditorToolbar.vue
+++ b/apps/desktop/src/components/layout/EditorToolbar.vue
@@ -23,9 +23,9 @@ import { useConnectionStore } from "@/stores/connectionStore";
 import { useDatabaseOptions } from "@/composables/useDatabaseOptions";
 import { useSchemaOptions } from "@/composables/useSchemaOptions";
 import { connectionIconType } from "@/lib/connectionPresentation";
-import { isDefaultDatabase } from "@/lib/defaultDatabase";
+import { formatDatabaseLabel, isDefaultDatabase } from "@/lib/defaultDatabase";
 import { connectionDisplayName } from "@/lib/tabPresentation";
-import { isSingleDatabase, usesTreeSchemaMode } from "@/lib/databaseCapabilities";
+import { isSingleDatabase } from "@/lib/databaseCapabilities";
 import { hexToRgba } from "@/lib/color";
 import type { QueryTab, ConnectionConfig } from "@/types/database";
 
@@ -103,10 +103,10 @@ const toolbarStyle = computed(() => {
 });
 
 function databaseDisplayName(database: string): string {
-  const connection = props.activeConnection;
-  if (connection?.db_type === "redis" && database !== "") return `db${database}`;
-  if (database === "" && usesTreeSchemaMode(connection?.db_type)) return t("editor.defaultDatabase");
-  return database || t("editor.noDatabase");
+  return formatDatabaseLabel(props.activeConnection, database, {
+    defaultDatabase: t("editor.defaultDatabase"),
+    noDatabase: t("editor.noDatabase"),
+  });
 }
 
 function connectionById(connectionId: string): ConnectionConfig | undefined {
diff --git a/apps/desktop/src/lib/defaultDatabase.ts b/apps/desktop/src/lib/defaultDatabase.ts
index c0c6cccb7..18891d184 100644
--- a/apps/desktop/src/lib/defaultDatabase.ts
+++ b/apps/desktop/src/lib/defaultDatabase.ts
@@ -1,9 +1,34 @@
-import type { ConnectionConfig } from "@/types/database";
+import type { ConnectionConfig, DatabaseType } from "@/types/database";
+import { usesTreeSchemaMode } from "@/lib/databaseCapabilities";
+
+export const TREE_SCHEMA_DEFAULT_DATABASE_SELECT_VALUE = "__dbx_tree_schema_default_database__";
 
 export function resolveDefaultDatabase(connection: Pick<ConnectionConfig, "database">, options: string[]): string {
   return connection.database || options[0] || "";
 }
 
+export function isTreeSchemaDefaultDatabase(dbType: DatabaseType | undefined, database: string): boolean {
+  return database === "" && usesTreeSchemaMode(dbType);
+}
+
+export function encodeSelectableDatabaseValue(dbType: DatabaseType | undefined, database: string): string {
+  return isTreeSchemaDefaultDatabase(dbType, database) ? TREE_SCHEMA_DEFAULT_DATABASE_SELECT_VALUE : database;
+}
+
+export function decodeSelectableDatabaseValue(dbType: DatabaseType | undefined, value: string): string {
+  return value === TREE_SCHEMA_DEFAULT_DATABASE_SELECT_VALUE && usesTreeSchemaMode(dbType) ? "" : value;
+}
+
+export function formatDatabaseLabel(
+  connection: Pick<ConnectionConfig, "db_type"> | undefined,
+  database: string,
+  labels: { defaultDatabase: string; noDatabase: string },
+): string {
+  if (connection?.db_type === "redis" && database !== "") return `db${database}`;
+  if (isTreeSchemaDefaultDatabase(connection?.db_type, database)) return labels.defaultDatabase;
+  return database || labels.noDatabase;
+}
+
 export function isDefaultDatabase(
   connection: Pick<ConnectionConfig, "database"> | undefined,
   database: string,
diff --git a/packages/app-tests/databaseOptions.test.ts b/packages/app-tests/databaseOptions.test.ts
index 3a2b373f2..6ae2d3f55 100644
--- a/packages/app-tests/databaseOptions.test.ts
+++ b/packages/app-tests/databaseOptions.test.ts
@@ -3,7 +3,7 @@ import { test } from "vitest";
 import { databaseOptionsForConnection } from "../../apps/desktop/src/composables/useDatabaseOptions.ts";
 
 test("tree-schema connections include the default database when no catalogs are returned", () => {
-  assert.deepEqual(databaseOptionsForConnection([], { db_type: "jdbc" }), [""]);
+  assert.deepEqual(databaseOptionsForConnection([], { db_type: "saphana" }), [""]);
 });
 
 test("non tree-schema connections keep an empty database option list", () => {
diff --git a/packages/app-tests/defaultDatabase.test.ts b/packages/app-tests/defaultDatabase.test.ts
index 71933562f..c108dac84 100644
--- a/packages/app-tests/defaultDatabase.test.ts
+++ b/packages/app-tests/defaultDatabase.test.ts
@@ -1,6 +1,13 @@
 import assert from "node:assert/strict";
 import { test } from "vitest";
-import { isDefaultDatabase, resolveDefaultDatabase } from "../../apps/desktop/src/lib/defaultDatabase.ts";
+import {
+  TREE_SCHEMA_DEFAULT_DATABASE_SELECT_VALUE,
+  decodeSelectableDatabaseValue,
+  encodeSelectableDatabaseValue,
+  formatDatabaseLabel,
+  isDefaultDatabase,
+  resolveDefaultDatabase,
+} from "../../apps/desktop/src/lib/defaultDatabase.ts";
 
 test("优先使用连接上已保存的默认数据库", () => {
   assert.equal(resolveDefaultDatabase({ database: "analytics" }, ["app", "analytics"]), "analytics");
@@ -20,3 +27,31 @@ test("判断当前数据库是否为默认数据库", () => {
   assert.equal(isDefaultDatabase(undefined, "analytics"), false);
   assert.equal(isDefaultDatabase({ database: "analytics" }, ""), false);
 });
+
+test("tree-schema 默认数据库会编码成可选中的稳定值并可解码回空字符串", () => {
+  assert.equal(encodeSelectableDatabaseValue("saphana", ""), TREE_SCHEMA_DEFAULT_DATABASE_SELECT_VALUE);
+  assert.equal(decodeSelectableDatabaseValue("saphana", TREE_SCHEMA_DEFAULT_DATABASE_SELECT_VALUE), "");
+});
+
+test("命名数据库保持原始值不变", () => {
+  assert.equal(encodeSelectableDatabaseValue("saphana", "SALES"), "SALES");
+  assert.equal(decodeSelectableDatabaseValue("saphana", "SALES"), "SALES");
+});
+
+test("数据库标签复用默认库显示语义", () => {
+  assert.equal(
+    formatDatabaseLabel({ db_type: "saphana" }, "", { defaultDatabase: "Default", noDatabase: "No database selected" }),
+    "Default",
+  );
+  assert.equal(
+    formatDatabaseLabel({ db_type: "postgres" }, "analytics", {
+      defaultDatabase: "Default",
+      noDatabase: "No database selected",
+    }),
+    "analytics",
+  );
+  assert.equal(
+    formatDatabaseLabel({ db_type: "redis" }, "3", { defaultDatabase: "Default", noDatabase: "No database selected" }),
+    "db3",
+  );
+});