diff --git a/apps/desktop/src/App.vue b/apps/desktop/src/App.vue index 0fd1e52d19..c4bcea55c0 100644 --- a/apps/desktop/src/App.vue +++ b/apps/desktop/src/App.vue @@ -128,6 +128,7 @@ import { savedSqlDefaultTargetForWrite } from "@/lib/savedSql/savedSqlExecutionT import { countActiveUpdateBlockingTasks } from "@/lib/app/appUpdateTaskGuard"; import { initSavedSqlEditorPositions } from "@/lib/app/savedSqlEditorPosition"; import { hasTreeNodeDatabaseContext } from "@/lib/sidebar/treeNodeContext"; +import { objectBrowserTablesToAiTreeNodes } from "@/lib/ai/objectBrowserToAiTargets"; import { isSchemaAware, isSingleDatabase, usesTreeSchemaMode } from "@/lib/database/databaseFeatureSupport"; import { codeMirrorSqlDialect, connectionUsesDatabaseObjectTreeMode, effectiveDatabaseTypeForConnection } from "@/lib/database/jdbcDialect"; import { canFormatSqlForDatabaseType, formatSqlForEditing, sqlFormatDialectForDbType } from "@/lib/sql/sqlFormatter"; @@ -3705,6 +3706,13 @@ onUnmounted(() => { @object-schema-change="(tabId: string, schema: string | undefined) => queryStore.updateSchema(tabId, schema)" @object-browser-viewport-change="(tabId: string, viewport: any) => queryStore.updateObjectBrowserViewport(tabId, viewport)" @object-browser-search-change="(tabId: string, query: string) => queryStore.updateObjectBrowserSearch(tabId, query)" + @add-object-table-to-ai=" + (tabId: string, tables: Array<{ name: string; schema?: string }>) => { + const tab = queryStore.tabs.find((candidate) => candidate.id === tabId) ?? activeTab; + if (!tab) return; + addToAi(objectBrowserTablesToAiTreeNodes(tab, tables)); + } + " @structure-editor-saved=" (tabId: string, commentChanged: boolean) => { const tab = queryStore.tabs.find((candidate) => candidate.id === tabId); diff --git a/apps/desktop/src/components/layout/ContentArea.vue b/apps/desktop/src/components/layout/ContentArea.vue index 3c0085b129..719e9cd6c0 100644 --- a/apps/desktop/src/components/layout/ContentArea.vue +++ b/apps/desktop/src/components/layout/ContentArea.vue @@ -2507,6 +2507,7 @@ defineExpose({ @schema-change="emit('objectSchemaChange', activeTab.id, $event)" @viewport-change="emit('objectBrowserViewportChange', activeTab.id, $event)" @search-change="emit('objectBrowserSearchChange', activeTab.id, $event)" + @add-to-ai="emit('addObjectTableToAi', activeTab.id, $event)" /> diff --git a/apps/desktop/src/components/layout/querySurfaces.ts b/apps/desktop/src/components/layout/querySurfaces.ts index acd60810f3..316a4b8e08 100644 --- a/apps/desktop/src/components/layout/querySurfaces.ts +++ b/apps/desktop/src/components/layout/querySurfaces.ts @@ -90,6 +90,7 @@ export interface ContentAreaSurfaceEmits { objectSchemaChange: [tabId: string, schema: string | undefined]; objectBrowserViewportChange: [tabId: string, viewport: ObjectBrowserViewport]; objectBrowserSearchChange: [tabId: string, query: string]; + addObjectTableToAi: [tabId: string, tables: Array<{ name: string; schema?: string }>]; structureEditorSaved: [tabId: string, commentChanged: boolean]; structureEditorClose: [tabId: string]; previewStatement: [tabId: string, range: StatementRange | null]; diff --git a/apps/desktop/src/components/objects/ObjectBrowser.vue b/apps/desktop/src/components/objects/ObjectBrowser.vue index 6fd6ec5309..31f27f1317 100644 --- a/apps/desktop/src/components/objects/ObjectBrowser.vue +++ b/apps/desktop/src/components/objects/ObjectBrowser.vue @@ -42,6 +42,7 @@ import { Search, ScrollText, ShieldCheck, + Sparkles, Square, Table, Table2, @@ -70,7 +71,7 @@ import * as api from "@/lib/backend/api"; import type { ColumnInfo, ConnectionConfig, ConstraintInfo, ForeignKeyInfo, IndexInfo, ObjectBrowserViewMode, ObjectBrowserViewport, ObjectInfo, ObjectSourceKind, ObjectStatistics, TableInfoTab, TreeNode, TriggerInfo } from "@/types/database"; import { sortTablesByFkDependency, type TableWithFk } from "@/lib/table/tableDependencySort"; import { isSchemaAware, supportsTableVacuum, supportsTransfer } from "@/lib/database/databaseCapabilities"; -import { supportsSchemaDiagram, supportsTableImport, supportsTableStructureEditing, supportsTableTruncate } from "@/lib/database/databaseFeatureSupport"; +import { supportsAiAssistantContext, supportsSchemaDiagram, supportsTableImport, supportsTableStructureEditing, supportsTableTruncate } from "@/lib/database/databaseFeatureSupport"; import { codeMirrorSqlDialect, connectionObjectTreeNodeSchema, connectionUsesDatabaseObjectTreeMode, effectiveDatabaseTypeForConnection, objectListSchemaForConnection, tableStructureDatabaseTypeForConnection } from "@/lib/database/jdbcDialect"; import { getTableMetadataCapabilities, type TableMetadataCapabilities } from "@/lib/table/tableMetadataCapabilities"; import { constraintsForConstraintsTab } from "@/lib/table/constraintPresentation"; @@ -188,6 +189,7 @@ const emit = defineEmits<{ schemaChange: [schema: string | undefined]; viewportChange: [viewport: ObjectBrowserViewport]; searchChange: [query: string]; + addToAi: [tables: Array<{ name: string; schema?: string }>]; }>(); const { t } = useI18n(); @@ -3108,6 +3110,22 @@ function isSelectedBatchTableContext(item: ObjectBrowserRow): boolean { return item.type === "TABLE" && selectedTableCount.value > 1 && selectedTableIds.value.has(item.id); } +function addToAiMenuItem(item: ObjectBrowserRow): ContextMenuItem { + const useBatch = isSelectedBatchTableContext(item); + const count = selectedTableCount.value; + // Schema stays per-row: the consumer (App.vue addToAi) resolves the final + // schema with its own fallback chain (table.schema || tab.schema — no + // database fallback, mirroring the sidebar tree path). ObjectBrowser is a + // single-schema view, but keeping row-level schemas makes the payload honest + // and future-proof if multi-schema selection ever appears. + const targets = useBatch ? selectedTableRows.value.map((row) => ({ name: row.name, schema: row.schema })) : [{ name: item.name, schema: item.schema }]; + return { + label: useBatch ? t("contextMenu.addToAiMultiple", { count }) : t("contextMenu.addToAi"), + action: () => emit("addToAi", targets), + icon: Sparkles, + }; +} + function selectedBatchTableCountLabel(key: "batchDrop" | "batchTruncate" | "batchEmpty"): string { return t(`contextMenu.${key}`, { count: selectedTableCount.value }); } @@ -3117,6 +3135,7 @@ function getTableMenuItems(item: ObjectBrowserRow): ContextMenuItem[] { return [ { label: t("contextMenu.viewData"), action: () => openViewData(item), icon: Table2 }, { label: t("contextMenu.newQuery"), action: () => openNewQuery(item), icon: TerminalSquare }, + ...(supportsAiAssistantContext(effectiveDatabaseType.value) ? [addToAiMenuItem(item)] : []), { label: "", separator: true }, exportDataSubmenu(item), { label: "", separator: true }, @@ -3160,6 +3179,7 @@ function getTableMenuItems(item: ObjectBrowserRow): ContextMenuItem[] { ...(canOpenStructureEditor.value ? [{ label: t("contextMenu.editStructure"), action: () => openStructureEditor(item), icon: PencilRuler }] : []), ...(canRename(item) ? [{ label: t("contextMenu.renameObject"), action: () => requestRename(item), icon: Pencil }] : []), { label: t("contextMenu.newQuery"), action: () => openNewQuery(item), icon: TerminalSquare }, + ...(supportsAiAssistantContext(effectiveDatabaseType.value) ? [addToAiMenuItem(item)] : []), ...(canOpenDiagram.value ? [{ label: t("diagram.open"), action: () => openDiagram(item), icon: Network }] : []), ...(canOpenTableImport.value ? [{ label: t("contextMenu.importData"), action: () => openTableImport(item), icon: Download }] : []), { label: t("dataCompare.title"), action: () => openDataCompare(item), icon: ArrowRightLeft }, diff --git a/apps/desktop/src/components/objects/__tests__/ObjectBrowserAddToAi.spec.ts b/apps/desktop/src/components/objects/__tests__/ObjectBrowserAddToAi.spec.ts new file mode 100644 index 0000000000..e414d78dfd --- /dev/null +++ b/apps/desktop/src/components/objects/__tests__/ObjectBrowserAddToAi.spec.ts @@ -0,0 +1,59 @@ +import { readFileSync } from "node:fs"; +import { describe, expect, it } from "vitest"; + +const source = readFileSync(new URL("../ObjectBrowser.vue", import.meta.url), "utf8"); + +function functionBody(name: string): string { + const signature = new RegExp(`(?:async\\s+)?function\\s+${name}\\s*\\([^)]*\\)\\s*(?::\\s*[^\\{]+)?\\{`, "m").exec(source); + if (!signature) throw new Error(`Missing function ${name}`); + const bodyStart = signature.index + signature[0].length; + let depth = 1; + for (let index = bodyStart; index < source.length; index += 1) { + if (source[index] === "{") depth += 1; + else if (source[index] === "}") depth -= 1; + if (depth === 0) return source.slice(bodyStart, index); + } + throw new Error(`Unclosed function ${name}`); +} + +describe("ObjectBrowser add to AI", () => { + it("declares the addToAi emit for table context hand-off", () => { + const emitDeclaration = source.slice(source.indexOf("const emit = defineEmits"), source.indexOf(">();", source.indexOf("const emit = defineEmits"))); + expect(emitDeclaration).toContain("addToAi"); + }); + + it("adds an 'Add to AI' menu item after 'New Query' in the table menu", () => { + const tableMenu = functionBody("getTableMenuItems"); + const newQueryIndex = tableMenu.indexOf('t("contextMenu.newQuery")'); + const addToAiIndex = tableMenu.indexOf("addToAiMenuItem(item)"); + expect(newQueryIndex).toBeGreaterThan(-1); + expect(addToAiIndex).toBeGreaterThan(newQueryIndex); + }); + + it("guards the menu item behind supportsAiAssistantContext", () => { + const tableMenu = functionBody("getTableMenuItems"); + expect(tableMenu).toContain("supportsAiAssistantContext(effectiveDatabaseType.value)"); + expect(tableMenu).toContain("[addToAiMenuItem(item)]"); + }); + + it("offers 'Add to AI' in the VictoriaMetrics early-return branch too", () => { + const tableMenu = functionBody("getTableMenuItems"); + // The VM branch is the first block inside getTableMenuItems; verify it guards + // the same AI entry so ObjectBrowser stays consistent with the sidebar tree. + const vmBranchStart = tableMenu.indexOf("if (isVictoriaMetrics.value)"); + const vmBranchEnd = tableMenu.indexOf("const useBatchActions", vmBranchStart); + const vmBranch = tableMenu.slice(vmBranchStart, vmBranchEnd); + + expect(vmBranch).toContain("supportsAiAssistantContext(effectiveDatabaseType.value)"); + expect(vmBranch).toContain("[addToAiMenuItem(item)]"); + expect(vmBranch.indexOf("addToAiMenuItem(item)")).toBeGreaterThan(vmBranch.indexOf('t("contextMenu.newQuery")')); + }); + + it("emits selected tables (batch) or the single table with per-row schema", () => { + const addToAi = functionBody("addToAiMenuItem"); + expect(addToAi).toContain("selectedTableRows.value.map((row) => ({ name: row.name, schema: row.schema }))"); + expect(addToAi).toContain('emit("addToAi", targets)'); + expect(addToAi).toContain('t("contextMenu.addToAiMultiple", { count })'); + expect(addToAi).toContain('t("contextMenu.addToAi")'); + }); +}); diff --git a/apps/desktop/src/lib/ai/__tests__/objectBrowserToAiTargets.spec.ts b/apps/desktop/src/lib/ai/__tests__/objectBrowserToAiTargets.spec.ts new file mode 100644 index 0000000000..fd0fedf73f --- /dev/null +++ b/apps/desktop/src/lib/ai/__tests__/objectBrowserToAiTargets.spec.ts @@ -0,0 +1,92 @@ +import { describe, expect, it } from "vitest"; +import { objectBrowserTablesToAiTreeNodes } from "@/lib/ai/objectBrowserToAiTargets"; + +function tab(overrides: Partial[0]> = {}) { + return { + id: "tab-1", + connectionId: "conn-1", + database: "db-1", + schema: undefined, + catalog: undefined, + ...overrides, + } as any; +} + +describe("objectBrowserTablesToAiTreeNodes", () => { + it("maps table names and row-level schemas into AI tree nodes", () => { + const nodes = objectBrowserTablesToAiTreeNodes(tab(), [ + { name: "users", schema: "public" }, + { name: "orders", schema: "sales" }, + ]); + + expect(nodes).toEqual([ + { + id: "conn-1:db-1:public:users", + label: "users", + type: "table", + connectionId: "conn-1", + database: "db-1", + catalog: undefined, + schema: "public", + }, + { + id: "conn-1:db-1:sales:orders", + label: "orders", + type: "table", + connectionId: "conn-1", + database: "db-1", + catalog: undefined, + schema: "sales", + }, + ]); + }); + + it("falls back to the tab schema when a row has no schema", () => { + const nodes = objectBrowserTablesToAiTreeNodes(tab({ schema: "app" }), [{ name: "users" }]); + + expect(nodes[0].schema).toBe("app"); + expect(nodes[0].id).toBe("conn-1:db-1:app:users"); + }); + + it("keeps the schema undefined instead of falling back to the database name", () => { + const nodes = objectBrowserTablesToAiTreeNodes(tab(), [{ name: "users" }, { name: "orders", schema: undefined }]); + + expect(nodes.map((node) => node.schema)).toEqual([undefined, undefined]); + expect(nodes[0].id).toBe("conn-1:db-1::users"); + }); + + it("keeps the schema undefined for schema-less engines so addToAi reuses the current tab", () => { + // MySQL/MariaDB-style connection: object-browser rows and the tab both + // carry schema: undefined (rows via connectionObjectTreeNodeSchema, which + // returns undefined for schema-less engines). The synthesized node must + // keep schema undefined so App.vue's addToAi() tab matching + // ((tab.schema || "") === (target.schema || "")) reuses the current tab + // instead of spawning a stray query tab and clearing accumulated AI + // context via clearContextReferences(). + const nodes = objectBrowserTablesToAiTreeNodes(tab(), [{ name: "users" }]); + + expect(nodes[0]).toEqual({ + id: "conn-1:db-1::users", + label: "users", + type: "table", + connectionId: "conn-1", + database: "db-1", + catalog: undefined, + schema: undefined, + }); + }); + + it("carries the tab catalog through for external-catalog contexts (Doris/StarRocks)", () => { + const nodes = objectBrowserTablesToAiTreeNodes(tab({ catalog: "hive_catalog" }), [{ name: "orders", schema: "tpch" }]); + + expect(nodes[0]).toEqual({ + id: "conn-1:db-1:tpch:orders", + label: "orders", + type: "table", + connectionId: "conn-1", + database: "db-1", + catalog: "hive_catalog", + schema: "tpch", + }); + }); +}); diff --git a/apps/desktop/src/lib/ai/objectBrowserToAiTargets.ts b/apps/desktop/src/lib/ai/objectBrowserToAiTargets.ts new file mode 100644 index 0000000000..a415603ea7 --- /dev/null +++ b/apps/desktop/src/lib/ai/objectBrowserToAiTargets.ts @@ -0,0 +1,38 @@ +import type { QueryTab, TreeNode } from "@/types/database"; + +/** A table reference emitted by the object browser's "Add to AI" action. */ +export interface ObjectBrowserAiTableTarget { + name: string; + schema?: string; +} + +/** + * Converts object-browser table selections into the TreeNode[] shape that + * App.vue's addToAi() consumes. Schema resolution mirrors the sidebar path, + * which passes real tree nodes through as-is: the row's own schema, then the + * tab's schema. The database name is deliberately NOT a schema fallback — + * rows are already built with the engine-aware connectionObjectTreeNodeSchema() + * (ObjectBrowser.vue), which applies the database-as-schema fallback for the + * engines where that is correct (Oracle/Dameng/Hive family, sqlite) and keeps + * schema undefined for schema-less engines (MySQL/MariaDB etc., whose tabs + * also carry schema: undefined). Substituting the database name here would + * break addToAi()'s tab matching for those engines, spawning a stray query + * tab and clearing the accumulated AI context. The tab's catalog + * (Doris/StarRocks external catalogs, etc.) is carried through so addToAi() + * can locate the correct query context instead of falling back to a + * catalog-less tab. + */ +export function objectBrowserTablesToAiTreeNodes(tab: QueryTab, tables: ObjectBrowserAiTableTarget[]): TreeNode[] { + return tables.map((table) => { + const schema = table.schema || tab.schema; + return { + id: `${tab.connectionId}:${tab.database}:${schema || ""}:${table.name}`, + label: table.name, + type: "table" as const, + connectionId: tab.connectionId, + database: tab.database, + catalog: tab.catalog, + schema, + }; + }); +} diff --git a/apps/desktop/src/lib/tabs/contentSurfaceEvents.ts b/apps/desktop/src/lib/tabs/contentSurfaceEvents.ts index a8e47693d6..9adcb242c9 100644 --- a/apps/desktop/src/lib/tabs/contentSurfaceEvents.ts +++ b/apps/desktop/src/lib/tabs/contentSurfaceEvents.ts @@ -35,6 +35,7 @@ export const contentSurfaceEventNames = [ "objectSchemaChange", "objectBrowserViewportChange", "objectBrowserSearchChange", + "addObjectTableToAi", "structureEditorSaved", "structureEditorClose", "previewStatement",