feat(ai): 支持聊天附件解析并优化数据库对象操作

This commit is contained in:
Syngnat
2026-06-12 12:30:28 +08:00
parent d5688d31f6
commit d1aa06d537
21 changed files with 908 additions and 134 deletions

View File

@@ -17,6 +17,7 @@
"@types/react-syntax-highlighter": "^15.5.13",
"antd": "^5.12.0",
"clsx": "^2.1.0",
"fflate": "^0.8.3",
"mermaid": "^11.13.0",
"react": "^18.2.0",
"react-dom": "^18.2.0",
@@ -3347,6 +3348,12 @@
}
}
},
"node_modules/fflate": {
"version": "0.8.3",
"resolved": "https://registry.npmjs.org/fflate/-/fflate-0.8.3.tgz",
"integrity": "sha512-tbZNuJrLwGUp3zshBtdy4W+ORxZuIh8a5ilyIEQDC5rY1f3U20JMry0Ll3WBzU58EZKsEuJFXhb5gwv8CsPvgA==",
"license": "MIT"
},
"node_modules/format": {
"version": "0.2.2",
"resolved": "https://registry.npmjs.org/format/-/format-0.2.2.tgz",

View File

@@ -19,6 +19,7 @@
"@types/react-syntax-highlighter": "^15.5.13",
"antd": "^5.12.0",
"clsx": "^2.1.0",
"fflate": "^0.8.3",
"mermaid": "^11.13.0",
"react": "^18.2.0",
"react-dom": "^18.2.0",

View File

@@ -3,6 +3,7 @@ import { createPortal } from 'react-dom';
import { useStore } from '../store';
import type { OverlayWorkbenchTheme } from '../utils/overlayWorkbenchTheme';
import type {
AIChatAttachment,
AIChatMessage,
JVMAIPlanContext,
JVMDiagnosticPlanContext,
@@ -61,7 +62,7 @@ export const AIChatPanel: React.FC<AIChatPanelProps> = ({
width = 380, darkMode, bgColor, onClose, onOpenSettings, onWidthChange, overlayTheme
}) => {
const [input, setInput] = useState('');
const [draftImages, setDraftImages] = useState<string[]>([]);
const [draftAttachments, setDraftAttachments] = useState<AIChatAttachment[]>([]);
const [sending, setSending] = useState(false);
const [showScrollBottom, setShowScrollBottom] = useState(false);
const [historyOpen, setHistoryOpen] = useState(false);
@@ -209,6 +210,7 @@ export const AIChatPanel: React.FC<AIChatPanelProps> = ({
truncateAIChatMessages(sid, msg.id);
deleteAIChatMessage(sid, msg.id);
setInput(msg.content);
setDraftAttachments(msg.attachments || []);
setTimeout(() => textareaRef.current?.focus(), 50);
}, [sid, truncateAIChatMessages, deleteAIChatMessage]);
@@ -342,7 +344,7 @@ export const AIChatPanel: React.FC<AIChatPanelProps> = ({
const handleSend = useCallback(async () => {
const text = input.trim();
if ((!text && draftImages.length === 0) || sending) return;
if ((!text && draftAttachments.length === 0) || sending) return;
const connectionKey = activeContext?.connectionId ? `${activeContext.connectionId}:${activeContext.dbName || ''}` : 'default';
const readiness = buildAIChatReadinessSnapshot({
@@ -375,9 +377,13 @@ export const AIChatPanel: React.FC<AIChatPanelProps> = ({
pendingJVMPlanContextRef.current = currentJVMPlanContext;
pendingJVMDiagnosticPlanContextRef.current = currentJVMDiagnosticPlanContext;
const currentImages = [...draftImages];
const currentAttachments = [...draftAttachments];
const currentImages = currentAttachments
.filter((attachment) => attachment.kind === 'image' && attachment.dataUrl)
.map((attachment) => attachment.dataUrl as string);
const currentFileAttachments = currentAttachments.filter((attachment) => attachment.kind !== 'image');
setInput('');
setDraftImages([]);
setDraftAttachments([]);
setSending(true);
if (textareaRef.current) {
@@ -387,6 +393,7 @@ export const AIChatPanel: React.FC<AIChatPanelProps> = ({
const userMsg: AIChatMessage = {
id: genId(), role: 'user', content: text, timestamp: Date.now(),
images: currentImages.length > 0 ? currentImages : undefined,
attachments: currentFileAttachments.length > 0 ? currentFileAttachments : undefined,
};
addAIChatMessage(sid, userMsg);
@@ -419,7 +426,7 @@ export const AIChatPanel: React.FC<AIChatPanelProps> = ({
useStore.getState().replaceAIChatHistory(sid, [compressedMsg, userMsg, connectingMsg]);
finalMessagesPayload = [
{ role: 'assistant', content: compressedMsg.content },
{ role: 'user', content: userMsg.content, images: userMsg.images }
toAIRequestMessage(userMsg),
];
}
@@ -449,7 +456,7 @@ export const AIChatPanel: React.FC<AIChatPanelProps> = ({
});
}, [
input,
draftImages,
draftAttachments,
sending,
messages,
addAIChatMessage,
@@ -642,8 +649,8 @@ export const AIChatPanel: React.FC<AIChatPanelProps> = ({
<AIChatInput
input={input}
setInput={setInput}
draftImages={draftImages}
setDraftImages={setDraftImages}
draftAttachments={draftAttachments}
setDraftAttachments={setDraftAttachments}
sending={sending}
onSend={handleSend}
onStop={handleStop}

View File

@@ -232,6 +232,10 @@ describe('Sidebar locate toolbar', () => {
title: 'users',
dataRef: { tableName: 'public.users' },
})).toBe('public.users');
expect(resolveSidebarTableNameForCopy({
title: 'v_users',
dataRef: { viewName: 'reporting.v_users' },
})).toBe('reporting.v_users');
expect(resolveSidebarTableNameForCopy({
title: 'users',
dataRef: {},

View File

@@ -2803,17 +2803,25 @@ const Sidebar: React.FC<{
}
};
const resolveCopyObjectNameLabel = (node: any): string => {
if (node?.type === 'view') return '视图名称';
if (node?.type === 'materialized-view') return '物化视图名称';
if (node?.type === 'db-event') return '事件名称';
return '表名';
};
const handleCopyTableName = async (node: any) => {
const tableName = resolveSidebarTableNameForCopy(node);
if (!tableName) {
message.warning('表名为空,无法复制');
const objectName = resolveSidebarTableNameForCopy(node);
const label = resolveCopyObjectNameLabel(node);
if (!objectName) {
message.warning(`${label}为空,无法复制`);
return;
}
try {
await navigator.clipboard.writeText(tableName);
message.success('表名已复制到剪贴板');
await navigator.clipboard.writeText(objectName);
message.success(`${label}已复制到剪贴板`);
} catch (e: any) {
message.error('复制表名失败: ' + (e?.message || String(e)));
message.error(`复制${label}失败: ` + (e?.message || String(e)));
}
};
@@ -6938,6 +6946,12 @@ const Sidebar: React.FC<{
icon: <CodeOutlined />,
onClick: () => openViewDefinition(node)
},
{
key: 'copy-view-name',
label: '复制名称',
icon: <CopyOutlined />,
onClick: () => handleCopyTableName(node)
},
{ type: 'divider' },
{
key: 'edit-view',
@@ -7000,6 +7014,12 @@ const Sidebar: React.FC<{
icon: <CodeOutlined />,
onClick: () => openViewDefinition(node)
},
{
key: 'copy-materialized-view-name',
label: '复制名称',
icon: <CopyOutlined />,
onClick: () => handleCopyTableName(node)
},
{
key: 'new-query',
label: '新建查询',

View File

@@ -1,38 +1,117 @@
import React from 'react';
import { FileTextOutlined, WarningOutlined } from '@ant-design/icons';
import type { AIChatAttachment } from '../../types';
import type { OverlayWorkbenchTheme } from '../../utils/overlayWorkbenchTheme';
import { formatAIChatAttachmentSize } from './aiChatAttachments';
interface AIChatAttachmentStripProps {
draftImages: string[];
attachments: AIChatAttachment[];
onRemove: (index: number) => void;
overlayTheme: OverlayWorkbenchTheme;
variant: 'legacy' | 'v2';
}
const formatAttachmentKind = (attachment: AIChatAttachment): string => {
if (attachment.kind === 'markdown') return 'MD';
if (attachment.kind === 'pdf') return 'PDF';
if (attachment.kind === 'word') return 'Word';
if (attachment.kind === 'excel') return 'Excel';
if (attachment.kind === 'text') return 'Text';
if (attachment.kind === 'image') return 'Image';
return 'File';
};
const AttachmentFileChip: React.FC<{
attachment: AIChatAttachment;
onRemove: () => void;
overlayTheme: OverlayWorkbenchTheme;
variant: 'legacy' | 'v2';
}> = ({ attachment, onRemove, overlayTheme, variant }) => {
if (variant === 'v2') {
return (
<div className={`gn-v2-ai-attachment-file${attachment.extractWarning ? ' has-warning' : ''}`}>
<FileTextOutlined />
<span className="gn-v2-ai-attachment-file-name" title={attachment.name}>{attachment.name}</span>
<span className="gn-v2-ai-attachment-file-meta">
{formatAttachmentKind(attachment)} · {formatAIChatAttachmentSize(attachment.size)}
</span>
{attachment.extractWarning ? <WarningOutlined title={attachment.extractWarning} /> : null}
<button type="button" onClick={onRemove} aria-label="移除附件">×</button>
</div>
);
}
return (
<div
style={{
display: 'flex',
alignItems: 'center',
gap: 6,
maxWidth: 220,
minHeight: 34,
border: overlayTheme.shellBorder,
borderRadius: 8,
padding: '4px 8px',
color: overlayTheme.titleText,
background: 'rgba(0,0,0,0.03)',
}}
title={attachment.extractWarning || attachment.name}
>
<FileTextOutlined />
<span style={{ minWidth: 0, overflow: 'hidden', textOverflow: 'ellipsis', whiteSpace: 'nowrap', fontSize: 12 }}>
{attachment.name}
</span>
<span style={{ color: overlayTheme.mutedText, fontSize: 11, flexShrink: 0 }}>
{formatAttachmentKind(attachment)}
</span>
{attachment.extractWarning ? <WarningOutlined style={{ color: '#faad14', flexShrink: 0 }} /> : null}
<button
type="button"
onClick={onRemove}
aria-label="移除附件"
style={{ border: 'none', background: 'transparent', cursor: 'pointer', color: overlayTheme.mutedText, padding: 0 }}
>
×
</button>
</div>
);
};
export const AIChatAttachmentStrip: React.FC<AIChatAttachmentStripProps> = ({
draftImages,
attachments,
onRemove,
overlayTheme,
variant,
}) => {
if (draftImages.length === 0) {
if (attachments.length === 0) {
return null;
}
if (variant === 'v2') {
return (
<div className="gn-v2-ai-attachment-row">
{draftImages.map((b64, index) => (
<div key={index} className="gn-v2-ai-attachment-thumb">
<img src={b64} alt={`Draft ${index}`} />
<button
type="button"
onClick={() => onRemove(index)}
aria-label="移除图片"
>
</button>
</div>
{attachments.map((attachment, index) => (
attachment.kind === 'image' && attachment.dataUrl ? (
<div key={attachment.id || index} className="gn-v2-ai-attachment-thumb">
<img src={attachment.dataUrl} alt={`Draft ${index}`} />
<button
type="button"
onClick={() => onRemove(index)}
aria-label="移除图片"
>
×
</button>
</div>
) : (
<AttachmentFileChip
key={attachment.id || index}
attachment={attachment}
overlayTheme={overlayTheme}
variant="v2"
onRemove={() => onRemove(index)}
/>
)
))}
</div>
);
@@ -40,16 +119,28 @@ export const AIChatAttachmentStrip: React.FC<AIChatAttachmentStripProps> = ({
return (
<>
{draftImages.map((b64, index) => (
<div key={index} style={{ position: 'relative', width: 60, height: 60, borderRadius: 6, overflow: 'hidden', border: overlayTheme.shellBorder }}>
<img src={b64} style={{ width: '100%', height: '100%', objectFit: 'cover' }} alt={`Draft ${index}`} />
<div
onClick={() => onRemove(index)}
style={{ position: 'absolute', top: 2, right: 2, background: 'rgba(0,0,0,0.5)', color: '#fff', borderRadius: '50%', width: 16, height: 16, display: 'flex', alignItems: 'center', justifyContent: 'center', cursor: 'pointer', fontSize: 10 }}
>
{attachments.map((attachment, index) => (
attachment.kind === 'image' && attachment.dataUrl ? (
<div key={attachment.id || index} style={{ position: 'relative', width: 60, height: 60, borderRadius: 6, overflow: 'hidden', border: overlayTheme.shellBorder }}>
<img src={attachment.dataUrl} style={{ width: '100%', height: '100%', objectFit: 'cover' }} alt={`Draft ${index}`} />
<button
type="button"
onClick={() => onRemove(index)}
aria-label="移除图片"
style={{ position: 'absolute', top: 2, right: 2, background: 'rgba(0,0,0,0.5)', color: '#fff', borderRadius: '50%', width: 16, height: 16, display: 'flex', alignItems: 'center', justifyContent: 'center', cursor: 'pointer', fontSize: 10, border: 'none', padding: 0 }}
>
×
</button>
</div>
</div>
) : (
<AttachmentFileChip
key={attachment.id || index}
attachment={attachment}
overlayTheme={overlayTheme}
variant="legacy"
onRemove={() => onRemove(index)}
/>
)
))}
</>
);

View File

@@ -3,18 +3,19 @@ import { Button, Tooltip } from 'antd';
import { CodeOutlined, PictureOutlined, SendOutlined, StopOutlined, TableOutlined } from '@ant-design/icons';
import type { OverlayWorkbenchTheme } from '../../utils/overlayWorkbenchTheme';
import { AI_CHAT_ATTACHMENT_ACCEPT } from './aiChatAttachments';
interface AIChatComposerActionsProps {
variant: 'legacy' | 'v2';
input: string;
draftImageCount: number;
draftAttachmentCount: number;
sending: boolean;
darkMode: boolean;
textColor: string;
mutedColor: string;
overlayTheme: OverlayWorkbenchTheme;
fileInputRef: React.RefObject<HTMLInputElement>;
onImageUpload: React.ChangeEventHandler<HTMLInputElement>;
onAttachmentUpload: React.ChangeEventHandler<HTMLInputElement>;
onOpenContext: () => void;
onOpenSlashMenu?: () => void;
onSend: () => void;
@@ -26,20 +27,20 @@ const buttonIconStyle = { fontSize: 16 };
const AIChatComposerActions: React.FC<AIChatComposerActionsProps> = ({
variant,
input,
draftImageCount,
draftAttachmentCount,
sending,
darkMode,
textColor,
mutedColor,
overlayTheme,
fileInputRef,
onImageUpload,
onAttachmentUpload,
onOpenContext,
onOpenSlashMenu,
onSend,
onStop,
}) => {
const canSend = input.trim().length > 0 || draftImageCount > 0;
const canSend = input.trim().length > 0 || draftAttachmentCount > 0;
const isV2 = variant === 'v2';
const legacyIconButtonStyle: React.CSSProperties = {
color: overlayTheme.mutedText,
@@ -61,13 +62,13 @@ const AIChatComposerActions: React.FC<AIChatComposerActionsProps> = ({
>
<input
type="file"
accept="image/*"
accept={AI_CHAT_ATTACHMENT_ACCEPT}
multiple
ref={fileInputRef}
style={{ display: 'none' }}
onChange={onImageUpload}
onChange={onAttachmentUpload}
/>
<Tooltip title="上传图片/截图">
<Tooltip title="上传附件图片、Markdown、Word、Excel、PDF、文本">
<Button
type="text"
icon={<PictureOutlined style={isV2 ? undefined : buttonIconStyle} />}

View File

@@ -37,8 +37,8 @@ const renderAIChatInput = (overrides: Partial<React.ComponentProps<typeof AIChat
<AIChatInput
input=""
setInput={() => {}}
draftImages={[]}
setDraftImages={() => {}}
draftAttachments={[]}
setDraftAttachments={() => {}}
sending={false}
onSend={() => {}}
onStop={() => {}}
@@ -69,8 +69,8 @@ describe('AIChatInput notice layout', () => {
<AIChatInput
input=""
setInput={() => {}}
draftImages={[]}
setDraftImages={() => {}}
draftAttachments={[]}
setDraftAttachments={() => {}}
sending={false}
onSend={() => {}}
onStop={() => {}}
@@ -115,8 +115,8 @@ describe('AIChatInput notice layout', () => {
<AIChatInput
input=""
setInput={() => {}}
draftImages={[]}
setDraftImages={() => {}}
draftAttachments={[]}
setDraftAttachments={() => {}}
sending={false}
onSend={() => {}}
onStop={() => {}}
@@ -149,8 +149,8 @@ describe('AIChatInput notice layout', () => {
<AIChatInput
input=""
setInput={() => {}}
draftImages={[]}
setDraftImages={() => {}}
draftAttachments={[]}
setDraftAttachments={() => {}}
sending={false}
onSend={() => {}}
onStop={() => {}}

View File

@@ -17,14 +17,15 @@ import AIChatContextPreview from './AIChatContextPreview';
import AIChatProviderModelSelect from './AIChatProviderModelSelect';
import { buildAIChatReadinessSnapshot } from './aiChatReadiness';
import { useAIChatContextBinding } from './useAIChatContextBinding';
import { useAIChatDraftImages } from './useAIChatDraftImages';
import { useAIChatDraftAttachments } from './useAIChatDraftAttachments';
import { useAISlashCommandMenu } from './useAISlashCommandMenu';
import type { AIChatAttachment } from '../../types';
interface AIChatInputProps {
input: string;
setInput: (val: string) => void;
draftImages: string[];
setDraftImages: React.Dispatch<React.SetStateAction<string[]>>;
draftAttachments: AIChatAttachment[];
setDraftAttachments: React.Dispatch<React.SetStateAction<AIChatAttachment[]>>;
sending: boolean;
onSend: () => void;
onStop: () => void;
@@ -51,7 +52,7 @@ interface AIChatInputProps {
}
export const AIChatInput: React.FC<AIChatInputProps> = ({
input, setInput, draftImages, setDraftImages, sending, onSend, onStop, handleKeyDown,
input, setInput, draftAttachments, setDraftAttachments, sending, onSend, onStop, handleKeyDown,
activeConnName, activeContext, activeProvider, dynamicModels, loadingModels,
sendShortcutBinding, shortcutPlatform = 'windows', composerNotice, onComposerAction,
onModelChange, onFetchModels, textareaRef, darkMode, textColor, mutedColor, overlayTheme,
@@ -98,11 +99,11 @@ export const AIChatInput: React.FC<AIChatInputProps> = ({
const {
fileInputRef,
handleImageUpload,
handleAttachmentUpload,
handlePasteImages,
handleRemoveDraftImage,
} = useAIChatDraftImages({
setDraftImages,
handleRemoveDraftAttachment,
} = useAIChatDraftAttachments({
setDraftAttachments,
});
const {
@@ -151,9 +152,9 @@ export const AIChatInput: React.FC<AIChatInputProps> = ({
/>
<AIChatAttachmentStrip
variant="legacy"
draftImages={draftImages}
attachments={draftAttachments}
overlayTheme={overlayTheme}
onRemove={handleRemoveDraftImage}
onRemove={handleRemoveDraftAttachment}
/>
</div>
<AIChatComposerNotice
@@ -245,14 +246,14 @@ export const AIChatInput: React.FC<AIChatInputProps> = ({
<AIChatComposerActions
variant="legacy"
input={input}
draftImageCount={draftImages.length}
draftAttachmentCount={draftAttachments.length}
sending={sending}
darkMode={darkMode}
textColor={textColor}
mutedColor={mutedColor}
overlayTheme={overlayTheme}
fileInputRef={fileInputRef}
onImageUpload={handleImageUpload}
onAttachmentUpload={handleAttachmentUpload}
onOpenContext={handleOpenContext}
onSend={onSend}
onStop={onStop}
@@ -305,9 +306,9 @@ export const AIChatInput: React.FC<AIChatInputProps> = ({
/>
<AIChatAttachmentStrip
variant="v2"
draftImages={draftImages}
attachments={draftAttachments}
overlayTheme={overlayTheme}
onRemove={handleRemoveDraftImage}
onRemove={handleRemoveDraftAttachment}
/>
<AIChatComposerNotice
composerNotice={composerNotice}
@@ -353,14 +354,14 @@ export const AIChatInput: React.FC<AIChatInputProps> = ({
<AIChatComposerActions
variant="v2"
input={input}
draftImageCount={draftImages.length}
draftAttachmentCount={draftAttachments.length}
sending={sending}
darkMode={darkMode}
textColor={textColor}
mutedColor={mutedColor}
overlayTheme={overlayTheme}
fileInputRef={fileInputRef}
onImageUpload={handleImageUpload}
onAttachmentUpload={handleAttachmentUpload}
onOpenContext={handleOpenContext}
onOpenSlashMenu={handleOpenSlashMenu}
onSend={onSend}

View File

@@ -5,7 +5,9 @@ import {
CopyOutlined,
DeleteOutlined,
EditOutlined,
FileTextOutlined,
ReloadOutlined,
WarningOutlined,
RobotOutlined,
UserOutlined,
} from '@ant-design/icons';
@@ -20,6 +22,7 @@ import {
} from '../../utils/jvmDiagnosticPlan';
import { AIMessageMarkdown } from './messageBubble/AIMessageMarkdown';
import { AIThinkingBlock, AIToolCallingBlock } from './messageBubble/AIMessageStatusBlocks';
import { formatAIChatAttachmentSize } from './aiChatAttachments';
interface AIMessageBubbleProps {
msg: AIChatMessage;
@@ -47,6 +50,43 @@ interface AIMessageActionBarProps {
onCopy: () => void;
}
const AIMessageAttachmentSummary: React.FC<{
msg: AIChatMessage;
overlayTheme: OverlayWorkbenchTheme;
}> = ({ msg, overlayTheme }) => {
const fileAttachments = (msg.attachments || []).filter((attachment) => attachment.kind !== 'image');
if (fileAttachments.length === 0) {
return null;
}
return (
<div style={{ display: 'flex', gap: 6, flexWrap: 'wrap', marginBottom: 10 }}>
{fileAttachments.map((attachment) => (
<div
key={attachment.id}
title={attachment.extractWarning || attachment.name}
style={{
display: 'inline-flex',
alignItems: 'center',
gap: 6,
maxWidth: 260,
padding: '4px 8px',
borderRadius: 8,
border: overlayTheme.shellBorder,
color: overlayTheme.titleText,
background: 'rgba(0,0,0,0.03)',
fontSize: 12,
}}
>
<FileTextOutlined />
<span style={{ overflow: 'hidden', textOverflow: 'ellipsis', whiteSpace: 'nowrap' }}>{attachment.name}</span>
<span style={{ color: overlayTheme.mutedText, flexShrink: 0 }}>{formatAIChatAttachmentSize(attachment.size)}</span>
{attachment.extractWarning ? <WarningOutlined style={{ color: '#faad14', flexShrink: 0 }} /> : null}
</div>
))}
</div>
);
};
const AIMessageActionBar: React.FC<AIMessageActionBarProps> = ({
msg,
isUser,
@@ -286,6 +326,7 @@ export const AIMessageBubble: React.FC<AIMessageBubbleProps> = React.memo(({
))}
</div>
)}
<AIMessageAttachmentSummary msg={msg} overlayTheme={overlayTheme} />
{!isUser && parsedThinking && (
<AIThinkingBlock

View File

@@ -0,0 +1,81 @@
import { describe, expect, it } from 'vitest';
import { strToU8, zipSync } from 'fflate';
import {
appendAIChatAttachmentsToContent,
buildAIChatAttachmentPromptText,
createAIChatAttachmentFromFile,
resolveAIChatAttachmentKind,
} from './aiChatAttachments';
const makeFile = (parts: BlobPart[], name: string, type: string): File => {
const blob = new Blob(parts, { type });
return Object.assign(blob, { name, lastModified: 0 }) as File;
};
describe('aiChatAttachments', () => {
it('extracts markdown text so it can be sent to AI', async () => {
const attachment = await createAIChatAttachmentFromFile(makeFile(['# Report\n\nhello'], 'report.md', 'text/markdown'));
expect(attachment.kind).toBe('markdown');
expect(attachment.text).toContain('# Report');
expect(buildAIChatAttachmentPromptText([attachment])).toContain('hello');
});
it('extracts docx document text from Office Open XML', async () => {
const bytes = zipSync({
'word/document.xml': strToU8('<w:document><w:body><w:p><w:r><w:t>用户增长</w:t></w:r></w:p><w:p><w:r><w:t>GMV &amp; 留存</w:t></w:r></w:p></w:body></w:document>'),
});
const attachment = await createAIChatAttachmentFromFile(makeFile([bytes], 'plan.docx', 'application/vnd.openxmlformats-officedocument.wordprocessingml.document'));
expect(attachment.kind).toBe('word');
expect(attachment.text).toContain('用户增长');
expect(attachment.text).toContain('GMV & 留存');
});
it('extracts xlsx worksheet rows with shared strings', async () => {
const bytes = zipSync({
'xl/sharedStrings.xml': strToU8('<sst><si><t>姓名</t></si><si><t>张三</t></si></sst>'),
'xl/worksheets/sheet1.xml': strToU8('<worksheet><sheetData><row><c t="s"><v>0</v></c><c><v>100</v></c></row><row><c t="s"><v>1</v></c><c><v>88</v></c></row></sheetData></worksheet>'),
});
const attachment = await createAIChatAttachmentFromFile(makeFile([bytes], 'score.xlsx', 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet'));
expect(attachment.kind).toBe('excel');
expect(attachment.text).toContain('姓名\t100');
expect(attachment.text).toContain('张三\t88');
});
it('extracts lightweight PDF literal text and keeps a warning about limitations', async () => {
const attachment = await createAIChatAttachmentFromFile(makeFile(['%PDF-1.4\nBT (Invoice total 42) Tj ET\n%%EOF'], 'invoice.pdf', 'application/pdf'));
expect(attachment.kind).toBe('pdf');
expect(attachment.text).toContain('Invoice total 42');
expect(attachment.extractWarning).toContain('轻量文本提取');
});
it('appends non-image attachments to the upstream user content', () => {
const content = appendAIChatAttachmentsToContent('帮我总结', [{
id: 'att-1',
name: 'report.txt',
mimeType: 'text/plain',
size: 12,
kind: 'text',
text: '核心指标下降',
}]);
expect(content).toContain('帮我总结');
expect(content).toContain('<用户上传附件>');
expect(content).toContain('核心指标下降');
});
it('keeps images out of the prompt text because they are sent through multimodal payload fields', () => {
expect(resolveAIChatAttachmentKind({ name: 'screen.png', type: 'image/png' })).toBe('image');
expect(buildAIChatAttachmentPromptText([{
id: 'att-img',
name: 'screen.png',
mimeType: 'image/png',
size: 10,
kind: 'image',
dataUrl: 'data:image/png;base64,abc',
}])).toBe('');
});
});

View File

@@ -0,0 +1,311 @@
import { strFromU8, unzipSync } from 'fflate';
import type { AIChatAttachment, AIChatAttachmentKind } from '../../types';
export const AI_CHAT_ATTACHMENT_ACCEPT = [
'image/*',
'.md',
'.markdown',
'.txt',
'.csv',
'.tsv',
'.json',
'.sql',
'.log',
'.xml',
'.yaml',
'.yml',
'.pdf',
'.doc',
'.docx',
'.xls',
'.xlsx',
].join(',');
const MAX_ATTACHMENT_BYTES = 15 * 1024 * 1024;
const MAX_ATTACHMENT_TEXT_CHARS = 60000;
const MAX_PROMPT_TEXT_CHARS = 50000;
const textExtensions = new Set([
'md',
'markdown',
'txt',
'csv',
'tsv',
'json',
'sql',
'log',
'xml',
'yaml',
'yml',
]);
const nextAttachmentId = () => `att-${Date.now()}-${Math.random().toString(36).slice(2, 8)}`;
export const formatAIChatAttachmentSize = (size: number): string => {
if (!Number.isFinite(size) || size <= 0) return '0 B';
if (size < 1024) return `${size} B`;
if (size < 1024 * 1024) return `${(size / 1024).toFixed(1)} KB`;
return `${(size / 1024 / 1024).toFixed(1)} MB`;
};
const getFileExtension = (name: string): string => {
const match = String(name || '').trim().toLowerCase().match(/\.([a-z0-9]+)$/);
return match?.[1] || '';
};
export const resolveAIChatAttachmentKind = (file: Pick<File, 'name' | 'type'>): AIChatAttachmentKind => {
const mimeType = String(file.type || '').toLowerCase();
const extension = getFileExtension(file.name);
if (mimeType.startsWith('image/')) return 'image';
if (extension === 'md' || extension === 'markdown') return 'markdown';
if (extension === 'pdf' || mimeType === 'application/pdf') return 'pdf';
if (extension === 'doc' || extension === 'docx' || mimeType.includes('word')) return 'word';
if (extension === 'xls' || extension === 'xlsx' || mimeType.includes('spreadsheet') || mimeType.includes('excel')) return 'excel';
if (textExtensions.has(extension) || mimeType.startsWith('text/')) return 'text';
return 'document';
};
const clampExtractedText = (raw: string): { text: string; truncated: boolean } => {
const normalized = raw
.replace(/\r\n?/g, '\n')
.replace(/[\u0000-\u0008\u000B\u000C\u000E-\u001F]+/g, ' ')
.replace(/[ \t]+\n/g, '\n')
.replace(/\n{3,}/g, '\n\n')
.trim();
if (normalized.length <= MAX_ATTACHMENT_TEXT_CHARS) {
return { text: normalized, truncated: false };
}
return {
text: normalized.slice(0, MAX_ATTACHMENT_TEXT_CHARS).trimEnd(),
truncated: true,
};
};
const decodeXmlEntities = (value: string): string => value
.replace(/&lt;/g, '<')
.replace(/&gt;/g, '>')
.replace(/&amp;/g, '&')
.replace(/&quot;/g, '"')
.replace(/&apos;/g, "'");
const collectXmlTextTags = (xml: string): string[] => {
const values: string[] = [];
const tagPattern = /<(?:[a-zA-Z0-9_]+:)?t\b[^>]*>([\s\S]*?)<\/(?:[a-zA-Z0-9_]+:)?t>/g;
let match: RegExpExecArray | null;
while ((match = tagPattern.exec(xml))) {
values.push(decodeXmlEntities(match[1].replace(/<[^>]+>/g, '')));
}
return values;
};
const extractDocxText = (entries: Record<string, Uint8Array>): string => {
const xmlPaths = Object.keys(entries)
.filter((path) => /^word\/(?:document|header\d+|footer\d+|footnotes|endnotes)\.xml$/i.test(path))
.sort((left, right) => {
if (left === 'word/document.xml') return -1;
if (right === 'word/document.xml') return 1;
return left.localeCompare(right);
});
return xmlPaths.map((path) => {
const xml = strFromU8(entries[path]);
const prepared = xml
.replace(/<w:tab\s*\/>/g, '\t')
.replace(/<w:(?:br|cr)\b[^>]*\/>/g, '\n')
.replace(/<\/w:p>/g, '\n');
return collectXmlTextTags(prepared).join('');
}).filter(Boolean).join('\n\n');
};
const extractSharedStrings = (entries: Record<string, Uint8Array>): string[] => {
const sharedStrings = entries['xl/sharedStrings.xml'];
if (!sharedStrings) return [];
const xml = strFromU8(sharedStrings);
const blocks = xml.match(/<si\b[\s\S]*?<\/si>/g) || [];
return blocks.map((block) => collectXmlTextTags(block).join(''));
};
const extractCellValue = (cellXml: string, sharedStrings: string[]): string => {
const type = cellXml.match(/\bt="([^"]+)"/)?.[1] || '';
if (type === 'inlineStr') {
return collectXmlTextTags(cellXml).join('');
}
const value = decodeXmlEntities(cellXml.match(/<v>([\s\S]*?)<\/v>/)?.[1] || '').trim();
if (type === 's') {
const index = Number.parseInt(value, 10);
return Number.isFinite(index) ? (sharedStrings[index] || '') : value;
}
return value;
};
const extractXlsxText = (entries: Record<string, Uint8Array>): string => {
const sharedStrings = extractSharedStrings(entries);
const sheetPaths = Object.keys(entries)
.filter((path) => /^xl\/worksheets\/sheet\d+\.xml$/i.test(path))
.sort((left, right) => left.localeCompare(right, undefined, { numeric: true }));
return sheetPaths.map((path) => {
const xml = strFromU8(entries[path]);
const rows = xml.match(/<row\b[\s\S]*?<\/row>/g) || [];
const lines = rows.map((rowXml) => {
const cells = rowXml.match(/<c\b[\s\S]*?<\/c>/g) || [];
return cells.map((cellXml) => extractCellValue(cellXml, sharedStrings)).join('\t').trimEnd();
}).filter((line) => line.trim().length > 0);
if (lines.length === 0) return '';
const sheetName = path.replace(/^xl\/worksheets\//i, '').replace(/\.xml$/i, '');
return `[工作表: ${sheetName}]\n${lines.join('\n')}`;
}).filter(Boolean).join('\n\n');
};
const decodePdfLiteralString = (value: string): string => value
.slice(1, -1)
.replace(/\\n/g, '\n')
.replace(/\\r/g, '\n')
.replace(/\\t/g, '\t')
.replace(/\\([()\\])/g, '$1')
.replace(/\\([0-7]{1,3})/g, (_, octal) => String.fromCharCode(Number.parseInt(octal, 8)));
const bytesToBinaryString = (bytes: Uint8Array): string => {
const chunks: string[] = [];
const chunkSize = 0x8000;
for (let offset = 0; offset < bytes.length; offset += chunkSize) {
chunks.push(String.fromCharCode(...bytes.slice(offset, offset + chunkSize)));
}
return chunks.join('');
};
const extractPdfText = (bytes: Uint8Array): { text: string; warning?: string } => {
const raw = bytesToBinaryString(bytes);
const values: string[] = [];
const literalPattern = /\((?:\\.|[^\\)]){1,2000}\)/g;
let match: RegExpExecArray | null;
while ((match = literalPattern.exec(raw)) && values.length < 5000) {
const decoded = decodePdfLiteralString(match[0]).trim();
if (decoded && /[\p{L}\p{N}\u4e00-\u9fa5]/u.test(decoded)) {
values.push(decoded);
}
}
const text = values.join('\n');
const warning = text
? 'PDF 已使用轻量文本提取;扫描件或压缩字体内容可能无法完整读取。'
: '未从 PDF 中提取到可读文本;如果是扫描件或复杂编码 PDF请复制正文后再发送。';
return { text, warning };
};
const extractLegacyOfficeText = (bytes: Uint8Array): { text: string; warning: string } => {
const raw = bytesToBinaryString(bytes);
const matches = raw.match(/[A-Za-z0-9\u4e00-\u9fa5][\x20-\x7E\u4e00-\u9fa5]{3,}/g) || [];
return {
text: Array.from(new Set(matches)).join('\n'),
warning: '旧版 Office 二进制格式仅做轻量文本片段提取;建议转为 docx/xlsx 后上传以获得更完整正文。',
};
};
const extractOfficeOpenXmlText = (bytes: Uint8Array, kind: AIChatAttachmentKind): string => {
const entries = unzipSync(bytes);
if (kind === 'word') return extractDocxText(entries);
if (kind === 'excel') return extractXlsxText(entries);
return '';
};
const readFileAsDataUrl = (file: File): Promise<string> => new Promise((resolve, reject) => {
const reader = new FileReader();
reader.onload = () => resolve(String(reader.result || ''));
reader.onerror = () => reject(reader.error || new Error('读取文件失败'));
reader.readAsDataURL(file);
});
export const createAIChatAttachmentFromFile = async (file: File): Promise<AIChatAttachment> => {
const kind = resolveAIChatAttachmentKind(file);
const base: Omit<AIChatAttachment, 'kind'> = {
id: nextAttachmentId(),
name: file.name || 'unnamed',
mimeType: file.type || 'application/octet-stream',
size: file.size || 0,
};
if (kind === 'image') {
return { ...base, kind, dataUrl: await readFileAsDataUrl(file) };
}
if (file.size > MAX_ATTACHMENT_BYTES) {
return {
...base,
kind,
extractWarning: `文件超过 ${formatAIChatAttachmentSize(MAX_ATTACHMENT_BYTES)},已附加文件信息但未读取正文。`,
};
}
try {
if (kind === 'text' || kind === 'markdown') {
const { text, truncated } = clampExtractedText(await file.text());
return { ...base, kind, text, textTruncated: truncated };
}
const bytes = new Uint8Array(await file.arrayBuffer());
const extension = getFileExtension(file.name);
if ((kind === 'word' || kind === 'excel') && (extension === 'docx' || extension === 'xlsx')) {
const { text, truncated } = clampExtractedText(extractOfficeOpenXmlText(bytes, kind));
return { ...base, kind, text, textTruncated: truncated };
}
if (kind === 'pdf') {
const extracted = extractPdfText(bytes);
const { text, truncated } = clampExtractedText(extracted.text);
return { ...base, kind, text, textTruncated: truncated, extractWarning: extracted.warning };
}
if ((kind === 'word' || kind === 'excel') && (extension === 'doc' || extension === 'xls')) {
const extracted = extractLegacyOfficeText(bytes);
const { text, truncated } = clampExtractedText(extracted.text);
return { ...base, kind, text, textTruncated: truncated, extractWarning: extracted.warning };
}
return {
...base,
kind,
extractWarning: '当前文件类型已附加,但暂未提取正文;如需模型分析内容,请改用 markdown、txt、docx、xlsx 或 pdf。',
};
} catch (error: any) {
return {
...base,
kind,
extractWarning: `附件正文提取失败:${error?.message || String(error)}`,
};
}
};
export const buildAIChatAttachmentPromptText = (attachments: AIChatAttachment[] = []): string => {
const documentAttachments = attachments.filter((attachment) => attachment.kind !== 'image');
if (documentAttachments.length === 0) return '';
return documentAttachments.map((attachment, index) => {
const content = String(attachment.text || '').trim();
const truncatedContent = content.length > MAX_PROMPT_TEXT_CHARS
? `${content.slice(0, MAX_PROMPT_TEXT_CHARS).trimEnd()}\n\n[附件正文过长,已截断]`
: content;
const fence = truncatedContent.includes('```') ? '~~~' : '```';
const lines = [
`### 附件 ${index + 1}: ${attachment.name}`,
`- 类型: ${attachment.kind}`,
`- MIME: ${attachment.mimeType || 'unknown'}`,
`- 大小: ${formatAIChatAttachmentSize(attachment.size)}`,
];
if (attachment.extractWarning) {
lines.push(`- 提取说明: ${attachment.extractWarning}`);
}
if (attachment.textTruncated) {
lines.push('- 提取说明: 附件正文较长,已截断后发送。');
}
if (truncatedContent) {
lines.push('', fence, truncatedContent, fence);
} else {
lines.push('', '未提取到可发送的附件正文。');
}
return lines.join('\n');
}).join('\n\n');
};
export const appendAIChatAttachmentsToContent = (content: string, attachments: AIChatAttachment[] = []): string => {
const attachmentPrompt = buildAIChatAttachmentPromptText(attachments);
if (!attachmentPrompt) return content;
const userContent = String(content || '').trim();
return [
userContent || '请根据以下附件内容继续处理。',
'',
'<用户上传附件>',
attachmentPrompt,
'</用户上传附件>',
].join('\n');
};

View File

@@ -0,0 +1,65 @@
import React from 'react';
import { message } from 'antd';
import type { AIChatAttachment } from '../../types';
import { createAIChatAttachmentFromFile } from './aiChatAttachments';
interface UseAIChatDraftAttachmentsParams {
setDraftAttachments: React.Dispatch<React.SetStateAction<AIChatAttachment[]>>;
}
export const useAIChatDraftAttachments = ({
setDraftAttachments,
}: UseAIChatDraftAttachmentsParams) => {
const fileInputRef = React.useRef<HTMLInputElement>(null);
const appendDraftFiles = React.useCallback(async (files: File[]) => {
for (const file of files) {
const attachment = await createAIChatAttachmentFromFile(file);
setDraftAttachments((prev) => [...prev, attachment]);
if (attachment.extractWarning) {
message.warning(`${attachment.name}: ${attachment.extractWarning}`);
}
}
}, [setDraftAttachments]);
const handleAttachmentUpload = React.useCallback((event: React.ChangeEvent<HTMLInputElement>) => {
const files = Array.from(event.target.files || []);
if (files.length > 0) {
void appendDraftFiles(files);
}
if (fileInputRef.current) {
fileInputRef.current.value = '';
}
}, [appendDraftFiles]);
const handlePasteImages = React.useCallback((event: React.ClipboardEvent<HTMLTextAreaElement>) => {
const items = event.clipboardData?.items;
if (!items) {
return;
}
const imageFiles: File[] = [];
for (let index = 0; index < items.length; index += 1) {
if (items[index].type.includes('image')) {
const blob = items[index].getAsFile();
if (blob) {
imageFiles.push(blob);
}
}
}
if (imageFiles.length > 0) {
event.preventDefault();
void appendDraftFiles(imageFiles);
}
}, [appendDraftFiles]);
const handleRemoveDraftAttachment = React.useCallback((index: number) => {
setDraftAttachments((prev) => prev.filter((_, currentIndex) => currentIndex !== index));
}, [setDraftAttachments]);
return {
fileInputRef,
handleAttachmentUpload,
handlePasteImages,
handleRemoveDraftAttachment,
};
};

View File

@@ -1,60 +0,0 @@
import React from 'react';
interface UseAIChatDraftImagesParams {
setDraftImages: React.Dispatch<React.SetStateAction<string[]>>;
}
export const useAIChatDraftImages = ({
setDraftImages,
}: UseAIChatDraftImagesParams) => {
const fileInputRef = React.useRef<HTMLInputElement>(null);
const appendDraftImage = React.useCallback((blob: Blob) => {
const reader = new FileReader();
reader.onload = (event) => {
if (event.target?.result) {
setDraftImages((prev) => [...prev, event.target!.result as string]);
}
};
reader.readAsDataURL(blob);
}, [setDraftImages]);
const handleImageUpload = React.useCallback((event: React.ChangeEvent<HTMLInputElement>) => {
const files = Array.from(event.target.files || []);
files.forEach((file) => {
if (file.type.includes('image')) {
appendDraftImage(file);
}
});
if (fileInputRef.current) {
fileInputRef.current.value = '';
}
}, [appendDraftImage]);
const handlePasteImages = React.useCallback((event: React.ClipboardEvent<HTMLTextAreaElement>) => {
const items = event.clipboardData?.items;
if (!items) {
return;
}
for (let index = 0; index < items.length; index += 1) {
if (items[index].type.includes('image')) {
event.preventDefault();
const blob = items[index].getAsFile();
if (blob) {
appendDraftImage(blob);
}
}
}
}, [appendDraftImage]);
const handleRemoveDraftImage = React.useCallback((index: number) => {
setDraftImages((prev) => prev.filter((_, currentIndex) => currentIndex !== index));
}, [setDraftImages]);
return {
fileInputRef,
handleImageUpload,
handlePasteImages,
handleRemoveDraftImage,
};
};

View File

@@ -66,7 +66,7 @@ export const shouldLoadSidebarNodeOnExpand = (
export const resolveSidebarTableNameForCopy = (
node: Pick<SidebarTreeNode, 'title' | 'dataRef'> | null | undefined,
): string => {
return String(node?.dataRef?.tableName || node?.title || '').trim();
return String(node?.dataRef?.tableName || node?.dataRef?.viewName || node?.dataRef?.eventName || node?.title || '').trim();
};
type SidebarTableSortPreference = 'name' | 'frequency';

View File

@@ -646,6 +646,20 @@ export interface AIToolCall {
};
}
export type AIChatAttachmentKind = "image" | "markdown" | "text" | "pdf" | "word" | "excel" | "document";
export interface AIChatAttachment {
id: string;
name: string;
mimeType: string;
size: number;
kind: AIChatAttachmentKind;
dataUrl?: string;
text?: string;
textTruncated?: boolean;
extractWarning?: string;
}
export type ChatPhase =
| "idle"
| "connecting"
@@ -663,6 +677,7 @@ export interface AIChatMessage {
timestamp: number;
loading?: boolean;
images?: string[]; // base64 encoded images with data URI prefix
attachments?: AIChatAttachment[];
tool_calls?: AIToolCall[];
tool_call_id?: string;
tool_name?: string; // used for UI display

View File

@@ -75,4 +75,24 @@ describe('toAIRequestMessage', () => {
images: ['data:image/png;base64,abc'],
});
});
it('appends extracted file attachment content to the user request payload', () => {
const payload = toAIRequestMessage(message({
role: 'user',
content: '帮我看附件',
attachments: [{
id: 'att-1',
name: 'report.md',
mimeType: 'text/markdown',
size: 24,
kind: 'markdown',
text: '# 周报\n收入下降',
}],
}));
expect(payload.content).toContain('帮我看附件');
expect(payload.content).toContain('<用户上传附件>');
expect(payload.content).toContain('report.md');
expect(payload.content).toContain('收入下降');
});
});

View File

@@ -1,4 +1,5 @@
import type { AIChatMessage, AIToolCall } from '../types';
import { appendAIChatAttachmentsToContent } from '../components/ai/aiChatAttachments';
export interface AIRequestMessage {
role: AIChatMessage['role'];
@@ -12,7 +13,7 @@ export interface AIRequestMessage {
export const toAIRequestMessage = (message: AIChatMessage): AIRequestMessage => {
const payload: AIRequestMessage = {
role: message.role,
content: message.content,
content: appendAIChatAttachmentsToContent(message.content, message.attachments),
};
if (message.images && message.images.length > 0) {

View File

@@ -6110,6 +6110,48 @@ body[data-ui-version="v2"] .gn-v2-ai-attachment-thumb button {
font-size: 10px;
}
body[data-ui-version="v2"] .gn-v2-ai-attachment-file {
display: inline-flex;
align-items: center;
gap: 6px;
max-width: min(100%, 280px);
min-height: 28px;
padding: 4px 8px;
border: 0.5px solid var(--gn-br-2);
border-radius: 8px;
background: var(--gn-bg-panel);
color: var(--gn-fg-1);
font-size: 11.5px;
}
body[data-ui-version="v2"] .gn-v2-ai-attachment-file.has-warning {
border-color: rgba(217, 119, 6, 0.45);
background: rgba(217, 119, 6, 0.08);
}
body[data-ui-version="v2"] .gn-v2-ai-attachment-file-name {
min-width: 0;
overflow: hidden;
text-overflow: ellipsis;
white-space: nowrap;
}
body[data-ui-version="v2"] .gn-v2-ai-attachment-file-meta {
color: var(--gn-fg-3);
flex-shrink: 0;
}
body[data-ui-version="v2"] .gn-v2-ai-attachment-file button {
width: 16px;
height: 16px;
padding: 0;
border: none;
border-radius: 50%;
background: var(--gn-bg-subtle);
color: var(--gn-fg-3);
line-height: 1;
}
body[data-ui-version="v2"] .gn-v2-ai-input-box {
min-height: 72px;
padding: 0;

View File

@@ -733,16 +733,18 @@ func (a *App) DBQueryMulti(config connection.ConnectionConfig, dbName string, qu
// 注意:原生 conn.Query() 执行写操作UPDATE/INSERT/DELETE
// sql.Rows 不暴露 RowsAffected导致影响行数丢失。
// 因此仅在全部语句皆为读操作时才使用原生路径。
statements := splitSQLStatements(query)
allReadOnly := true
for _, stmt := range splitSQLStatements(query) {
for _, stmt := range statements {
if strings.TrimSpace(stmt) != "" && !isReadOnlySQLQuery(runConfig.Type, stmt) {
allReadOnly = false
break
}
}
useNativeMultiResult := shouldUseNativeMultiResultBatch(runConfig.Type, statements, allReadOnly)
runMultiQuery := func(inst db.Database) ([]connection.ResultSetData, []string, error) {
if !allReadOnly {
if !useNativeMultiResult {
return nil, nil, nil // 包含写操作,走逐条执行路径
}
if q, ok := inst.(db.MultiResultQueryMessageExecer); ok {
@@ -781,7 +783,6 @@ func (a *App) DBQueryMulti(config connection.ConnectionConfig, dbName string, qu
}
// 驱动不支持多结果集,回退到逐条执行
statements := splitSQLStatements(query)
if len(statements) == 0 {
return connection.QueryResult{
Success: true,
@@ -1005,6 +1006,26 @@ func (a *App) DBQueryMulti(config connection.ConnectionConfig, dbName string, qu
return connection.QueryResult{Success: true, Data: resultSets, QueryID: queryID, Message: fallbackMsg}
}
func shouldUseNativeMultiResultBatch(dbType string, statements []string, allReadOnly bool) bool {
if allReadOnly {
return true
}
if !strings.EqualFold(strings.TrimSpace(dbType), "sqlserver") {
return false
}
for _, stmt := range statements {
stmt = strings.TrimSpace(stmt)
if stmt == "" {
continue
}
if isReadOnlySQLQuery(dbType, stmt) || shouldTryQueryResultFirst(dbType, stmt) {
continue
}
return false
}
return true
}
func shouldTryQueryResultFirst(dbType string, query string) bool {
isSQLServer := strings.EqualFold(strings.TrimSpace(dbType), "sqlserver")
if keyword, withHasWrite := sqlDataOperationInfo(query); withHasWrite && keyword == "select" {

View File

@@ -4,6 +4,7 @@ import (
"context"
"database/sql"
"errors"
"strings"
"testing"
"GoNavi-Wails/internal/connection"
@@ -28,6 +29,44 @@ type fakeBatchWriteDB struct {
session *fakeBatchWriteSession
}
type fakeNativeMultiResultDB struct {
*fakeBatchWriteDB
multiCalls int
}
func (f *fakeNativeMultiResultDB) QueryMulti(query string) ([]connection.ResultSetData, error) {
results, _, err := f.QueryMultiWithMessages(query)
return results, err
}
func (f *fakeNativeMultiResultDB) QueryMultiWithMessages(query string) ([]connection.ResultSetData, []string, error) {
return f.QueryMultiContextWithMessages(context.Background(), query)
}
func (f *fakeNativeMultiResultDB) QueryMultiContext(ctx context.Context, query string) ([]connection.ResultSetData, error) {
results, _, err := f.QueryMultiContextWithMessages(ctx, query)
return results, err
}
func (f *fakeNativeMultiResultDB) QueryMultiContextWithMessages(ctx context.Context, query string) ([]connection.ResultSetData, []string, error) {
f.multiCalls++
if err := f.queryErr[query]; err != nil {
return nil, nil, err
}
if multi := f.multiResult[query]; len(multi) > 0 {
return cloneResultSets(multi), append([]string(nil), f.messageMap[query]...), nil
}
rows, columns, messages, err := f.QueryContextWithMessages(ctx, query)
if err != nil {
return nil, nil, err
}
return []connection.ResultSetData{{
Rows: rows,
Columns: columns,
Messages: append([]string(nil), messages...),
}}, append([]string(nil), messages...), nil
}
func (f *fakeBatchWriteDB) Connect(config connection.ConnectionConfig) error {
return nil
}
@@ -1227,6 +1266,72 @@ func TestDBQueryMultiDoesNotBatchExecStoredProcedureAsWriteStatement(t *testing.
}
}
func TestDBQueryMultiRunsSQLServerStatisticsBatchNatively(t *testing.T) {
originalNewDatabaseFunc := newDatabaseFunc
t.Cleanup(func() {
newDatabaseFunc = originalNewDatabaseFunc
})
query := "SET STATISTICS IO, TIME ON;\nSELECT 1 AS value;"
baseDB := &fakeBatchWriteDB{
multiResult: map[string][]connection.ResultSetData{
query: {
{
Rows: []map[string]interface{}{},
Columns: []string{},
Messages: []string{"SQL Server parse and compile time: CPU time = 0 ms."},
},
{
Rows: []map[string]interface{}{{"value": 1}},
Columns: []string{"value"},
Messages: []string{"Table 'users'. Scan count 1, logical reads 3."},
},
},
},
messageMap: map[string][]string{
query: {"Table 'users'. Scan count 1, logical reads 3."},
},
queryErr: map[string]error{},
}
fakeDB := &fakeNativeMultiResultDB{fakeBatchWriteDB: baseDB}
newDatabaseFunc = func(dbType string) (db.Database, error) {
return fakeDB, nil
}
app := NewAppWithSecretStore(secretstore.NewUnavailableStore("test"))
config := connection.ConnectionConfig{Type: "sqlserver", Host: "127.0.0.1", Port: 1433, User: "sa"}
result := app.DBQueryMulti(config, "master", query, "sqlserver-statistics-native-batch-test")
if !result.Success {
t.Fatalf("expected DBQueryMulti success, got failure: %s", result.Message)
}
if strings.Contains(result.Message, "不支持原生多语句执行") {
t.Fatalf("expected SQL Server statistics batch to avoid sequential fallback warning, got %q", result.Message)
}
if fakeDB.multiCalls != 1 {
t.Fatalf("expected one native multi-result batch call, got %d", fakeDB.multiCalls)
}
if baseDB.session != nil {
t.Fatal("expected native SQL Server batch to avoid sequential session fallback")
}
if baseDB.queryCalls != 0 || baseDB.execCalls != 0 {
t.Fatalf("expected native batch to avoid per-statement query/exec calls, queryCalls=%d execCalls=%d", baseDB.queryCalls, baseDB.execCalls)
}
resultSets, ok := result.Data.([]connection.ResultSetData)
if !ok {
t.Fatalf("expected []connection.ResultSetData, got %T", result.Data)
}
if len(resultSets) != 2 {
t.Fatalf("expected two native result sets, got %#v", resultSets)
}
if got := resultSets[1].Rows[0]["value"]; got != 1 {
t.Fatalf("expected SELECT result value=1, got %#v", got)
}
if len(result.Messages) != 1 || !strings.Contains(result.Messages[0], "logical reads") {
t.Fatalf("expected SQL Server statistics message to be returned, got %#v", result.Messages)
}
}
func TestDBQueryMultiUsesPinnedSessionForSequentialFallback(t *testing.T) {
originalNewDatabaseFunc := newDatabaseFunc
t.Cleanup(func() {