feat(ai): drag scoped table references into chat

This commit is contained in:
Allen authored and GitHub committed 2026-08-21 13:03:14 +08:00
1 parent cc78d4f528
commit 09b326558f
6 files changed
+197 -4

No files matched your search

@@ -1,4 +1,4 @@
<script setup lang="ts">
<script setup lang="ts">
import { computed, nextTick, onMounted, onUnmounted, ref, toRaw, watch, type Component } from "vue";
import { uuid } from "@/lib/common/utils";
import { useI18n } from "vue-i18n";
@@ -131,6 +131,8 @@ import ExplainPlanViewer from "@/components/explain/ExplainPlanViewer.vue";
import { parseExplainResult, parseOracleExplainText, type ParsedExplainPlan } from "@/lib/diagram/explainPlan";
import { copyToClipboard } from "@/lib/common/clipboard";
import { AI_TABLE_MENTION_CANDIDATE_LIMIT, AI_TABLE_MENTION_SCHEMA_LIMIT, filterAiTableMentionCandidates, formatAiTableMention, parseAiTableMentions, type AiTableMention } from "@/lib/ai/aiTableMentions";
import { handleAiTableReferenceDropEvent } from "@/lib/ai/aiTableReferenceDrop";
import { DBX_TABLE_REFERENCE_DROP_EVENT, clearActiveTableReferencePayload } from "@/lib/editor/queryEditorTableDrop";
import { isAiPromptImeCompositionEvent, shouldSubmitAiPromptOnKeydown } from "@/lib/ai/aiPromptKeyboard";
import { isActionableWriteProposalMessage, isActionableWriteSqlProposal, looksLikeActionProposal, looksLikeWriteSqlProposal, shouldGrantWriteSqlOnShortAffirmative } from "@/lib/ai/aiProposalDetect";
import { visibleToActualIndex } from "@/lib/ai/aiMessageEdit";
@@ -2286,6 +2288,22 @@ function onTauriFileDrop(event: Event) {
void addDroppedAttachmentPaths(payload.paths);
}
function onTableReferenceDropEvent(event: Event) {
handleAiTableReferenceDropEvent(event, {
context: {
connectionId: props.tab?.connectionId || props.connection?.id,
database: props.tab?.database || props.connection?.database || "",
},
assistantRoot: assistantRootRef.value,
elementFromPoint: (x, y) => document.elementFromPoint(x, y),
onMention: (mention, payload) => {
addSelectedMention({ kind: "table", schema: mention.schema, name: mention.table, tableType: "table" });
clearActiveTableReferencePayload(payload);
nextTick(() => promptTextareaRef.value?.focus());
},
});
}
async function send() {
const text = prompt.value.trim();
if ((!text && !selectedMentions.value.length && !selectedSqlFileMentions.value.length && !selectedCsvAttachments.value.length && !selectedImageAttachments.value.length) || isGenerating.value) return;
@@ -2891,6 +2909,7 @@ onMounted(async () => {
window.addEventListener("resize", handlePanelResize);
document.addEventListener("dbx:tauri-file-drop", onTauriFileDrop as EventListener);
window.addEventListener(DBX_TABLE_REFERENCE_DROP_EVENT, onTableReferenceDropEvent);
if (typeof ResizeObserver !== "undefined" && assistantRootRef.value) {
promptPanelResizeObserver = new ResizeObserver(handlePanelResize);
promptPanelResizeObserver.observe(assistantRootRef.value);
@@ -2969,6 +2988,7 @@ onUnmounted(() => {
document.body.style.cursor = "";
window.removeEventListener("resize", handlePanelResize);
document.removeEventListener("dbx:tauri-file-drop", onTauriFileDrop as EventListener);
window.removeEventListener(DBX_TABLE_REFERENCE_DROP_EVENT, onTableReferenceDropEvent);
promptPanelResizeObserver?.disconnect();
});
@@ -3042,7 +3062,7 @@ async function openExternalUrl(url: string) {
</script>
<template>
<div ref="assistantRootRef" class="flex h-full min-h-0 flex-col overflow-hidden" @dragenter="onAttachmentDragEnter" @dragover="onAttachmentDragOver" @dragleave="onAttachmentDragLeave" @drop="onAttachmentDrop">
<div ref="assistantRootRef" data-ai-assistant-root class="flex h-full min-h-0 flex-col overflow-hidden" @dragenter="onAttachmentDragEnter" @dragover="onAttachmentDragOver" @dragleave="onAttachmentDragLeave" @drop="onAttachmentDrop">
<div class="flex items-center gap-2 border-b px-3 shrink-0" :class="settings.editorSettings.appLayout === 'classic' ? 'h-9' : 'h-10'">
<span class="flex flex-1 self-stretch items-center truncate text-xs font-medium" data-tauri-drag-region>
{{ chatTitle }}
@@ -0,0 +1,53 @@
// @vitest-environment happy-dom
import { afterEach, describe, expect, it } from "vitest";
import { handleAiTableReferenceDropEvent, type AiTableReferenceDropContext } from "@/lib/ai/aiTableReferenceDrop";
import { createTableReferenceDropEvent, createTableReferencePayload, DBX_TABLE_REFERENCE_DROP_EVENT, type QueryEditorTableReferencePayload } from "@/lib/editor/queryEditorTableDrop";
function dispatchTableReferenceDrop(payload: QueryEditorTableReferencePayload, context: AiTableReferenceDropContext) {
const assistantRoot = document.createElement("div");
const target = document.createElement("span");
assistantRoot.append(target);
document.body.append(assistantRoot);
const mentions: string[] = [];
const listener = (event: Event) => {
handleAiTableReferenceDropEvent(event, {
context,
assistantRoot,
elementFromPoint: () => target,
onMention: (mention) => mentions.push(mention.raw),
});
};
window.addEventListener(DBX_TABLE_REFERENCE_DROP_EVENT, listener);
window.dispatchEvent(createTableReferenceDropEvent({ payload, clientX: 12, clientY: 24 }));
window.removeEventListener(DBX_TABLE_REFERENCE_DROP_EVENT, listener);
return mentions;
}
function tablePayload(overrides: Partial<QueryEditorTableReferencePayload> = {}) {
return createTableReferencePayload({
connectionId: overrides.connectionId ?? "conn-1",
database: overrides.database ?? "app-db",
schema: "public",
tableName: "users",
databaseType: "postgres",
})!;
}
afterEach(() => {
document.body.innerHTML = "";
});
describe("AI assistant table reference drop", () => {
it("accepts a table dropped from the active connection and database", () => {
expect(dispatchTableReferenceDrop(tablePayload(), { connectionId: "conn-1", database: "app-db" })).toEqual(["@public.users"]);
});
it("rejects a table dropped from another connection", () => {
expect(dispatchTableReferenceDrop(tablePayload({ connectionId: "conn-2" }), { connectionId: "conn-1", database: "app-db" })).toEqual([]);
});
it("rejects a table dropped from another database", () => {
expect(dispatchTableReferenceDrop(tablePayload({ database: "analytics" }), { connectionId: "conn-1", database: "app-db" })).toEqual([]);
});
});
@@ -54,6 +54,7 @@ import LightTooltip from "@/components/ui/LightTooltip.vue";
import type { ColumnInfo, ConnectionConfig, CustomTypeTreeMemberMeta, DatabaseType, TreeNode, TriggerInfo } from "@/types/database";
import { alignedCommentLeadingWidth, canTreeNodePin, canTreeNodeShowExpander, sidebarTreeNodeComment, trailingCommentAvailableWidth, trailingCommentGapPx, treeItemPaddingLeft, treeLabelWidthClass, usesFullWidthTreeLabel } from "@/lib/sidebar/sidebarTreeItemLayout";
import { clearActiveTableReferencePayload, createTableReferencePayload, createTableReferenceDropEvent, setActiveTableReferencePayload, type QueryEditorTableReferencePayload } from "@/lib/editor/queryEditorTableDrop";
import { AI_ASSISTANT_TABLE_DROP_ROOT_SELECTOR } from "@/lib/ai/aiTableReferenceDrop";
import { formatSidebarObjectStorage } from "@/lib/sidebar/sidebarDatabaseStorage";
import { dataTabOpenModeFromTreeClick } from "@/lib/sidebar/dataTabOpenPolicy";
import { effectiveDatabaseTypeForConnection } from "@/lib/database/jdbcDialect";
@@ -1189,7 +1190,7 @@ function onTableReferenceMouseUp(event: MouseEvent) {
if (payload) {
suppressNextTableReferenceClick = true;
const target = document.elementFromPoint(event.clientX, event.clientY);
if (target instanceof Element && target.closest("[data-query-editor-root]")) {
if (target instanceof Element && target.closest(`[data-query-editor-root], ${AI_ASSISTANT_TABLE_DROP_ROOT_SELECTOR}`)) {
window.dispatchEvent(
createTableReferenceDropEvent({
payload,
@@ -0,0 +1,67 @@
import { describe, expect, it } from "vitest";
import { aiTableMentionFromTableReference } from "@/lib/ai/aiTableReferenceDrop";
import { createTableReferencePayload } from "@/lib/editor/queryEditorTableDrop";
describe("ai table reference drop", () => {
const context = { connectionId: "conn-1", database: "app-db" };
it("maps a table payload to a table mention", () => {
const payload = createTableReferencePayload({
connectionId: "conn-1",
database: "app-db",
schema: "public",
tableName: "users",
databaseType: "postgres",
});
expect(aiTableMentionFromTableReference(payload, context)).toEqual({
raw: "@public.users",
schema: "public",
table: "users",
});
});
it("quotes mention parts that are not simple identifiers", () => {
const payload = createTableReferencePayload({
connectionId: "conn-1",
database: "app-db",
tableName: "order items",
databaseType: "mysql",
});
expect(aiTableMentionFromTableReference(payload, context)).toEqual({
raw: '@"order items"',
schema: undefined,
table: "order items",
});
});
it("ignores database references", () => {
const payload = createTableReferencePayload({
connectionId: "conn-1",
database: "app-db",
referenceType: "database",
databaseType: "mysql",
});
expect(aiTableMentionFromTableReference(payload, context)).toBeNull();
});
it("ignores column references", () => {
const payload = createTableReferencePayload({
connectionId: "conn-1",
database: "app-db",
schema: "public",
tableName: "users",
columnName: "email",
databaseType: "postgres",
});
expect(aiTableMentionFromTableReference(payload, context)).toBeNull();
});
it("ignores empty payloads", () => {
expect(aiTableMentionFromTableReference(null, context)).toBeNull();
expect(aiTableMentionFromTableReference(undefined, context)).toBeNull();
});
});
@@ -0,0 +1,43 @@
import { formatAiTableMention, type AiTableMention } from "@/lib/ai/aiTableMentions";
import type { QueryEditorTableReferenceDropDetail, QueryEditorTableReferencePayload } from "@/lib/editor/queryEditorTableDrop";
/** Selector marking the AI assistant panel root as a table-reference drop target. */
export const AI_ASSISTANT_TABLE_DROP_ROOT_SELECTOR = "[data-ai-assistant-root]";
export interface AiTableReferenceDropContext {
connectionId?: string;
database?: string;
}
export interface AiTableReferenceDropHandlerOptions {
context: AiTableReferenceDropContext;
assistantRoot: Element | null | undefined;
elementFromPoint: (x: number, y: number) => Element | null;
onMention: (mention: AiTableMention, payload: QueryEditorTableReferencePayload) => void;
}
/**
* Maps a sidebar table-reference drag payload to an AI table mention chip.
* Only plain table/view references become mentions; database and column
* references are not representable as table mentions and return null.
*/
export function aiTableMentionFromTableReference(payload: QueryEditorTableReferencePayload | null | undefined, context: AiTableReferenceDropContext): AiTableMention | null {
if (!payload || payload.referenceType === "database" || payload.columnName) return null;
if (!context.connectionId || context.database == null || payload.connectionId !== context.connectionId || payload.database !== context.database) return null;
const table = payload.tableName;
if (!table) return null;
const schema = payload.schema;
return { raw: formatAiTableMention(schema, table), schema, table };
}
export function handleAiTableReferenceDropEvent(event: Event, options: AiTableReferenceDropHandlerOptions): boolean {
if (!(event instanceof CustomEvent)) return false;
const detail = event.detail as QueryEditorTableReferenceDropDetail | undefined;
if (!detail?.payload) return false;
const target = options.elementFromPoint(detail.clientX, detail.clientY);
if (!target || !options.assistantRoot?.contains(target)) return false;
const mention = aiTableMentionFromTableReference(detail.payload, options.context);
if (!mention) return false;
options.onMention(mention, detail.payload);
return true;
}
@@ -24,7 +24,16 @@ export interface QueryEditorTableReferenceDropDetail {
let activeTableReferencePayload: QueryEditorTableReferencePayload | null = null;
export function createTableReferencePayload(options: { connectionId?: string; database?: string; schema?: string; tableName?: string; columnName?: string; referenceType?: "database" | "table" | "column"; databaseType?: DatabaseType; driverProfile?: string }): QueryEditorTableReferencePayload | null {
export function createTableReferencePayload(options: {
connectionId?: string;
database?: string;
schema?: string;
tableName?: string;
columnName?: string;
referenceType?: "database" | "table" | "column";
databaseType?: DatabaseType;
driverProfile?: string;
}): QueryEditorTableReferencePayload | null {
if (!options.connectionId || options.database == null) return null;
const referenceType = options.referenceType ?? (options.columnName ? "column" : "table");
if (referenceType !== "database" && !options.tableName) return null;