fix: support tree-schema default DB selection in AI assistant

This commit is contained in:
t8y2 2026-06-09 00:44:04 +08:00
parent 03ade2ebba
commit 025a4fdcce
5 changed files with 112 additions and 20 deletions

View File

@ -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>

View File

@ -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 {

View File

@ -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,

View File

@ -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", () => {

View File

@ -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",
);
});