From 1ae2b742795f6a791885ad93fc9feac0031f4cfa Mon Sep 17 00:00:00 2001 From: Syngnat Date: Mon, 13 Jul 2026 12:52:56 +0800 Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9B=20fix(sql-parser):=20=E4=BF=AE?= =?UTF-8?q?=E5=A4=8D=E5=B0=BE=E9=9A=8F=E6=B3=A8=E9=87=8A=E5=AF=BC=E8=87=B4?= =?UTF-8?q?=E4=BA=8B=E5=8A=A1=E4=B8=8E=E5=8F=AA=E8=AF=BB=E5=88=A4=E5=AE=9A?= =?UTF-8?q?=E5=BC=82=E5=B8=B8?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 过滤分号后的纯注释语句并保留数据库可执行版本注释 - 前后端按数据库方言统一处理双横线、井号与块注释 - 修复事务选择、只读保护、SQL 审计及 AI 风险分析误判 - 补充流式 SQL、事务执行与方言解析回归测试 --- .../QueryEditor.external-sql-save.test.tsx | 39 ++++++++ frontend/src/components/QueryEditor.tsx | 49 ++++++---- .../ai/aiSqlEditorTransactionInsights.ts | 20 +++-- .../components/ai/aiSqlRiskInsights.test.ts | 18 ++++ .../src/components/ai/aiSqlRiskInsights.ts | 10 ++- frontend/src/utils/connectionReadOnly.test.ts | 16 ++++ frontend/src/utils/connectionReadOnly.ts | 47 +++++++--- .../src/utils/sqlEditorTransaction.test.ts | 20 +++++ .../src/utils/sqlStatementSelection.test.ts | 45 ++++++++++ frontend/src/utils/sqlStatementSelection.ts | 89 +++++++++++++++++-- internal/app/connection_readonly.go | 5 +- internal/app/methods_db.go | 4 +- internal/app/methods_db_multi_test.go | 47 ++++++++++ internal/app/methods_db_transaction.go | 8 +- internal/app/methods_db_transaction_test.go | 14 +++ internal/app/methods_sql_audit.go | 6 +- internal/app/sql_inspect.go | 2 +- internal/app/sql_split.go | 62 ++++++++++++- internal/app/sql_split_stream.go | 30 ++++--- internal/app/sql_split_stream_test.go | 41 +++++++++ internal/app/sql_split_test.go | 59 ++++++++++++ 21 files changed, 559 insertions(+), 72 deletions(-) create mode 100644 internal/app/sql_split_stream_test.go diff --git a/frontend/src/components/QueryEditor.external-sql-save.test.tsx b/frontend/src/components/QueryEditor.external-sql-save.test.tsx index 07618305..fa67ea36 100644 --- a/frontend/src/components/QueryEditor.external-sql-save.test.tsx +++ b/frontend/src/components/QueryEditor.external-sql-save.test.tsx @@ -7495,6 +7495,45 @@ describe('QueryEditor external SQL save', () => { expect(textContent(renderer!.root)).not.toContain('未提交'); }); + it('keeps DML with a trailing line comment in a pending managed transaction', async () => { + backendApp.DBQueryMultiTransactional.mockResolvedValueOnce({ + success: true, + transactionId: 'tx-comment', + transactionPending: true, + data: [ + { columns: ['affectedRows'], rows: [{ affectedRows: 1 }], statementIndex: 1 }, + ], + }); + + let renderer!: ReactTestRenderer; + await act(async () => { + renderer = create(); + }); + + await act(async () => { + await findButton(renderer!, '运行').props.onClick(); + }); + await act(async () => { + await Promise.resolve(); + await Promise.resolve(); + }); + + expect(backendApp.DBQueryMultiTransactional).toHaveBeenCalledWith( + expect.anything(), + 'main', + 'DELETE FROM users WHERE id = 1', + 'query-1', + ); + expect(backendApp.DBQueryMulti).not.toHaveBeenCalled(); + expect(storeState.sqlEditorPendingTransactions['tab-1']).toMatchObject({ + id: 'tx-comment', + dbType: 'mysql', + statements: ['DELETE FROM users WHERE id = 1'], + }); + }); + it('keeps TDengine insert on the regular query path because it has no managed transaction support', async () => { storeState.connections[0].config.type = 'tdengine'; backendApp.DBQueryMulti.mockResolvedValueOnce({ diff --git a/frontend/src/components/QueryEditor.tsx b/frontend/src/components/QueryEditor.tsx index 6a4d2801..32e78139 100644 --- a/frontend/src/components/QueryEditor.tsx +++ b/frontend/src/components/QueryEditor.tsx @@ -4935,7 +4935,7 @@ const QueryEditor: React.FC<{ tab: TabData; isActive?: boolean }> = ({ tab, isAc lineNumber: Number(position?.lineNumber || 1), column: Number(position?.column || 1), }); - const currentStatementRange = resolveCurrentSqlStatementRange(fullText, cursorOffset); + const currentStatementRange = resolveCurrentSqlStatementRange(fullText, cursorOffset, activeDialect); // 获取当前行光标前的内容 const linePrefix = model.getLineContent(position.lineNumber).slice(0, position.column - 1); @@ -5649,8 +5649,8 @@ const QueryEditor: React.FC<{ tab: TabData; isActive?: boolean }> = ({ tab, isAc }, ]; - const splitSQLStatements = (sql: string): string[] => { - return findSqlStatementRanges(sql).map((range) => range.text); + const splitSQLStatements = (sql: string, dbType = ''): string[] => { + return findSqlStatementRanges(sql, dbType).map((range) => range.text); }; const containsOraclePlsqlDefinition = (statements: string[]): boolean => ( @@ -5776,19 +5776,19 @@ const QueryEditor: React.FC<{ tab: TabData; isActive?: boolean }> = ({ tab, isAc || ''; }; - const resolveExecutableSQLAtEditorPosition = (model: any, sqlText: string, position: any): string => { + const resolveExecutableSQLAtEditorPosition = (model: any, sqlText: string, position: any, dbType = ''): string => { const normalizedPosition = normalizeEditorPosition(position); if (!normalizedPosition) return ''; const cursorOffset = getNormalizedOffsetAtPosition(sqlText, normalizedPosition); - const resolved = resolveExecutableSql(sqlText, cursorOffset, ''); + const resolved = resolveExecutableSql(sqlText, cursorOffset, '', dbType); return resolved?.sql || ''; }; - const getExecutableSQLAtCurrentCursor = (model: any, sqlText: string): string => { + const getExecutableSQLAtCurrentCursor = (model: any, sqlText: string, dbType = ''): string => { const editor = editorRef.current; const liveSelection = normalizeEditorPosition(editor?.getSelection?.()); if (liveSelection) { - return resolveExecutableSQLAtEditorPosition(model, sqlText, liveSelection); + return resolveExecutableSQLAtEditorPosition(model, sqlText, liveSelection, dbType); } const livePosition = normalizeEditorPosition(editor?.getPosition?.()); @@ -5802,12 +5802,12 @@ const QueryEditor: React.FC<{ tab: TabData; isActive?: boolean }> = ({ tab, isAc const key = `${position.lineNumber}:${position.column}`; if (seen.has(key)) continue; seen.add(key); - const sql = resolveExecutableSQLAtEditorPosition(model, sqlText, position); + const sql = resolveExecutableSQLAtEditorPosition(model, sqlText, position, dbType); if (sql.trim()) return sql; } const fallbackPosition = cachedPosition || livePosition; - return resolveExecutableSQLAtEditorPosition(model, sqlText, fallbackPosition); + return resolveExecutableSQLAtEditorPosition(model, sqlText, fallbackPosition, dbType); }; const getExecutableSQL = (): string => { @@ -5829,7 +5829,13 @@ const QueryEditor: React.FC<{ tab: TabData; isActive?: boolean }> = ({ tab, isAc if (selected) { return selectedSQL; } - return getExecutableSQLAtCurrentCursor(model, String(model.getValue?.() ?? currentQuery)); + const activeConnection = connections.find((connection) => connection.id === currentConnectionId); + const activeDialect = resolveSqlDialect( + String(activeConnection?.config?.type || ''), + String(activeConnection?.config?.driver || ''), + { oceanBaseProtocol: activeConnection?.config?.oceanBaseProtocol }, + ); + return getExecutableSQLAtCurrentCursor(model, String(model.getValue?.() ?? currentQuery), activeDialect); }; const captureEditorCursorPosition = (event?: React.MouseEvent) => { @@ -5847,9 +5853,10 @@ const QueryEditor: React.FC<{ tab: TabData; isActive?: boolean }> = ({ tab, isAc sql: string, queryId: string, sourceStatements: string[], + dbType = String(config.type || ''), ) => { const pendingTransaction = pendingSqlTransactionRef.current; - if (pendingTransaction && canReusePendingSqlEditorTransactionForType(String(config.type || ''), sourceStatements)) { + if (pendingTransaction && canReusePendingSqlEditorTransactionForType(dbType, sourceStatements)) { return DBQueryMultiInTransaction(pendingTransaction.id, sql, queryId); } return DBQueryMulti(buildRpcConnectionConfig(config) as any, dbName, sql, queryId); @@ -5871,6 +5878,11 @@ const QueryEditor: React.FC<{ tab: TabData; isActive?: boolean }> = ({ tab, isAc useSSH: conn.config.useSSH || false, ssh: conn.config.ssh || { host: "", port: 22, user: "", password: "", keyPath: "" } }; + const normalizedDbType = String(resolveSqlDialect( + String(config.type || ''), + String((config as any).driver || ''), + { oceanBaseProtocol: String((config as any).oceanBaseProtocol || '') }, + )).trim().toLowerCase(); try { setLoading(true); @@ -5886,7 +5898,8 @@ const QueryEditor: React.FC<{ tab: TabData; isActive?: boolean }> = ({ tab, isAc currentDb, sql, queryId, - splitSQLStatements(sql), + splitSQLStatements(sql, normalizedDbType), + normalizedDbType, ); if (!res?.success) { message.error(translate('query_editor.message.refresh_failed', { @@ -5989,7 +6002,8 @@ const QueryEditor: React.FC<{ tab: TabData; isActive?: boolean }> = ({ tab, isAc currentDb, pageSql, queryId, - splitSQLStatements(pageSql), + splitSQLStatements(pageSql, normalizedDbType), + normalizedDbType, ); if (!res?.success) { message.error(translate('query_editor.message.page_query_failed', { @@ -6128,13 +6142,13 @@ const QueryEditor: React.FC<{ tab: TabData; isActive?: boolean }> = ({ tab, isAc const splitInput = normalizedRawSQL .replace(/^\s*\/\/.*$/gm, '') .replace(/^\s*#.*$/gm, ''); - const statements = splitSQLStatements(splitInput); + const statements = splitSQLStatements(splitInput, normalizedDbType); const didExecuteAppendedSql = resultSets.length > 0 && lastExecutedEditorQueryRef.current && currentQuery.startsWith(lastExecutedEditorQueryRef.current) && normalizedRawSQL.trim() === currentQuery.slice(lastExecutedEditorQueryRef.current.length).replace(/；/g, ';').trim(); const didExecuteWholeEditor = areSqlStatementListsEqual( - splitSQLStatements(currentQuery.replace(/；/g, ';')), + splitSQLStatements(currentQuery.replace(/；/g, ';'), normalizedDbType), statements, ); if (statements.length === 0) { @@ -6289,13 +6303,13 @@ const QueryEditor: React.FC<{ tab: TabData; isActive?: boolean }> = ({ tab, isAc } else { // 非 MongoDB：使用 DBQueryMulti 一次性执行多条 SQL，后端返回多结果集 - const sourceStatements = splitSQLStatements(normalizedRawSQL); + const sourceStatements = splitSQLStatements(normalizedRawSQL, normalizedDbType); const didExecuteAppendedSql = resultSets.length > 0 && lastExecutedEditorQueryRef.current && currentQuery.startsWith(lastExecutedEditorQueryRef.current) && normalizedRawSQL.trim() === currentQuery.slice(lastExecutedEditorQueryRef.current.length).replace(/；/g, ';').trim(); const didExecuteWholeEditor = areSqlStatementListsEqual( - splitSQLStatements(currentQuery.replace(/；/g, ';')), + splitSQLStatements(currentQuery.replace(/；/g, ';'), normalizedDbType), sourceStatements, ); if (sourceStatements.length === 0) { @@ -6451,6 +6465,7 @@ const QueryEditor: React.FC<{ tab: TabData; isActive?: boolean }> = ({ tab, isAc fullSQL, queryId, executableStatements, + normalizedDbType, ); const duration = Date.now() - startTime; diff --git a/frontend/src/components/ai/aiSqlEditorTransactionInsights.ts b/frontend/src/components/ai/aiSqlEditorTransactionInsights.ts index d06501b3..e8ae9ead 100644 --- a/frontend/src/components/ai/aiSqlEditorTransactionInsights.ts +++ b/frontend/src/components/ai/aiSqlEditorTransactionInsights.ts @@ -2,7 +2,11 @@ import type { SqlLog } from '../../store'; import type { I18nParams } from '../../i18n'; import type { SavedConnection, TabData } from '../../types'; import { findSqlStatementRanges } from '../../utils/sqlStatementSelection'; -import { shouldUseSqlEditorManagedTransaction } from '../../utils/sqlEditorTransaction'; +import { + shouldUseSqlEditorManagedTransaction, + shouldUseSqlEditorManagedTransactionForType, +} from '../../utils/sqlEditorTransaction'; +import { resolveSqlDialect } from '../../utils/sqlDialect'; import type { AISqlEditorPendingTransactionRuntimeState, AISqlEditorTransactionRuntimeState, @@ -33,8 +37,8 @@ const normalizeDelayMs = (value: unknown): number => { return Number.isFinite(delayMs) && delayMs > 0 ? delayMs : DEFAULT_AUTO_COMMIT_DELAY_MS; }; -const splitStatements = (sql: string): string[] => - findSqlStatementRanges(String(sql || '')) +const splitStatements = (sql: string, dbType = ''): string[] => + findSqlStatementRanges(String(sql || ''), dbType) .map((range) => String(range.text || '').trim()) .filter(Boolean); @@ -88,9 +92,15 @@ const buildActiveSqlTabSnapshot = (params: { } const sql = String(activeTab.query || '').trim(); - const statements = splitStatements(sql); + const connection = connections.find((item) => item.id === activeTab.connectionId); + const dbType = resolveSqlDialect( + String(connection?.config?.type || ''), + String(connection?.config?.driver || ''), + { oceanBaseProtocol: connection?.config?.oceanBaseProtocol }, + ); + const statements = splitStatements(sql, dbType); const hasExplicitTransactionControl = statements.some(hasTransactionControlStatement); - const usesManagedTransaction = shouldUseSqlEditorManagedTransaction(statements); + const usesManagedTransaction = shouldUseSqlEditorManagedTransactionForType(dbType, statements); return { hasActiveTab: true, diff --git a/frontend/src/components/ai/aiSqlRiskInsights.test.ts b/frontend/src/components/ai/aiSqlRiskInsights.test.ts index bbf9b48d..3fc46b1b 100644 --- a/frontend/src/components/ai/aiSqlRiskInsights.test.ts +++ b/frontend/src/components/ai/aiSqlRiskInsights.test.ts @@ -106,4 +106,22 @@ describe('aiSqlRiskInsights', () => { expect(snapshot.warnings).toContain('UPDATE is missing a WHERE clause and may update the entire table.'); expect(snapshot.warnings).toContain('The current AI safety policy does not allow UPDATE SQL.'); }); + + it('uses the active connection dialect when splitting SQL comments', () => { + const snapshot = buildSqlRiskSnapshot({ + tabs: [{ + id: 'tab-1', + title: '用户删除', + type: 'query', + connectionId: 'conn-1', + dbName: 'crm', + query: 'DELETE FROM users WHERE id = 1;--compact', + }], + activeTabId: 'tab-1', + connections, + }); + + expect(snapshot.statementCount).toBe(2); + expect(snapshot.warnings.join('\n')).toContain('Confirm the impact scope of each statement'); + }); }); diff --git a/frontend/src/components/ai/aiSqlRiskInsights.ts b/frontend/src/components/ai/aiSqlRiskInsights.ts index b30a7e51..afdf7edc 100644 --- a/frontend/src/components/ai/aiSqlRiskInsights.ts +++ b/frontend/src/components/ai/aiSqlRiskInsights.ts @@ -1,5 +1,6 @@ import type { SavedConnection, TabData } from '../../types'; import type { I18nParams } from '../../i18n'; +import { resolveSqlDialect } from '../../utils/sqlDialect'; import { findSqlStatementRanges } from '../../utils/sqlStatementSelection'; import type { AIInspectionTranslator } from './aiInspectionI18n'; import { translateInspectionCopy } from './aiInspectionI18n'; @@ -277,7 +278,14 @@ export const buildSqlRiskSnapshot = (params: { }; } - const statements = findSqlStatementRanges(sql).map((range) => range.text.trim()).filter(Boolean); + const dbType = connection + ? resolveSqlDialect( + String(connection.config?.type || ''), + String(connection.config?.driver || ''), + { oceanBaseProtocol: connection.config?.oceanBaseProtocol }, + ) + : ''; + const statements = findSqlStatementRanges(sql, dbType).map((range) => range.text.trim()).filter(Boolean); const statementRisks = statements.map((statement) => buildStatementRisk(statement, params.translate)); let riskLevel: SqlRiskLevel = statements.length > 0 ? 'low' : 'none'; const warnings: string[] = []; diff --git a/frontend/src/utils/connectionReadOnly.test.ts b/frontend/src/utils/connectionReadOnly.test.ts index fe35136e..3e88d487 100644 --- a/frontend/src/utils/connectionReadOnly.test.ts +++ b/frontend/src/utils/connectionReadOnly.test.ts @@ -54,4 +54,20 @@ describe('connectionReadOnly', () => { }, }, "UPDATE users SET name = 'next';")).toEqual([]); }); + + it('uses the connection dialect when filtering comment-only statements', () => { + expect(findConnectionMutatingStatements({ + type: 'postgres', + protection: { + restrictScriptExecution: true, + }, + }, 'SELECT * FROM users; /*! MySQL-only comment */')).toEqual([]); + + expect(findConnectionMutatingStatements({ + type: 'mysql', + protection: { + restrictScriptExecution: true, + }, + }, 'SELECT * FROM users;--compact')).toEqual(['--compact']); + }); }); diff --git a/frontend/src/utils/connectionReadOnly.ts b/frontend/src/utils/connectionReadOnly.ts index cbc519fe..4e44e062 100644 --- a/frontend/src/utils/connectionReadOnly.ts +++ b/frontend/src/utils/connectionReadOnly.ts @@ -134,20 +134,45 @@ const MONGO_META_KEYS = new Set([ "writeconcern", ]); -const stripLeadingSqlComments = (statement: string): string => { +const MYSQL_COMMENT_DIALECTS = new Set([ + "mysql", + "goldendb", + "mariadb", + "oceanbase", + "diros", + "starrocks", + "sphinx", + "tidb", +]); + +const isDashLineCommentStart = (text: string, dbType: string): boolean => { + if (!MYSQL_COMMENT_DIALECTS.has(dbType)) return true; + const next = text[2] || ""; + return !next || /\s/.test(next); +}; + +const supportsHashLineComment = (dbType: string): boolean => + !dbType || dbType === "clickhouse" || MYSQL_COMMENT_DIALECTS.has(dbType); + +const isExecutableBlockComment = (text: string, dbType: string): boolean => { + if (text.slice(0, 4).toLowerCase() === "/*m!") return dbType === "mariadb"; + return text.startsWith("/*!") && MYSQL_COMMENT_DIALECTS.has(dbType); +}; + +const stripLeadingSqlComments = (statement: string, dbType = ""): string => { let text = String(statement || "").trim(); while (text) { - if (text.startsWith("--")) { + if (text.startsWith("--") && isDashLineCommentStart(text, dbType)) { const next = text.indexOf("\n"); text = next >= 0 ? text.slice(next + 1).trimStart() : ""; continue; } - if (text.startsWith("#")) { + if (text.startsWith("#") && supportsHashLineComment(dbType)) { const next = text.indexOf("\n"); text = next >= 0 ? text.slice(next + 1).trimStart() : ""; continue; } - if (text.startsWith("/*")) { + if (text.startsWith("/*") && !isExecutableBlockComment(text, dbType)) { const next = text.indexOf("*/"); text = next >= 0 ? text.slice(next + 2).trimStart() : ""; continue; @@ -157,8 +182,8 @@ const stripLeadingSqlComments = (statement: string): string => { return text; }; -const extractLeadingSqlKeyword = (statement: string): string => { - const text = stripLeadingSqlComments(statement); +const extractLeadingSqlKeyword = (statement: string, dbType = ""): string => { + const text = stripLeadingSqlComments(statement, dbType); const match = text.match(/^[A-Za-z_][A-Za-z0-9_]*/); return match ? match[0].toLowerCase() : ""; }; @@ -176,10 +201,10 @@ const resolveConnectionReadOnlyType = ( .toLowerCase(); }; -const isReadOnlySqlStatement = (statement: string): boolean => { - const text = stripLeadingSqlComments(statement); +const isReadOnlySqlStatement = (statement: string, dbType: string): boolean => { + const text = stripLeadingSqlComments(statement, dbType); if (!text) return true; - const keyword = extractLeadingSqlKeyword(text); + const keyword = extractLeadingSqlKeyword(text, dbType); if (!keyword || !SQL_READ_ONLY_KEYWORDS.has(keyword)) { return false; } @@ -252,7 +277,7 @@ const isConnectionReadOnlyStatement = ( if (dialect === "mongodb") { return isReadOnlyMongoStatement(statement); } - return isReadOnlySqlStatement(statement); + return isReadOnlySqlStatement(statement, dialect); }; export const supportsConnectionReadOnlyMode = ( @@ -366,7 +391,7 @@ export const findConnectionMutatingStatements = ( if (!isConnectionScriptExecutionRestricted(config)) { return []; } - return findSqlStatementRanges(String(sql || "")) + return findSqlStatementRanges(String(sql || ""), resolveConnectionReadOnlyType(config)) .map((range) => range.text.trim()) .filter((statement) => statement.length > 0) .filter((statement) => !isConnectionReadOnlyStatement(config, statement)); diff --git a/frontend/src/utils/sqlEditorTransaction.test.ts b/frontend/src/utils/sqlEditorTransaction.test.ts index 1cd2a714..fa5c3520 100644 --- a/frontend/src/utils/sqlEditorTransaction.test.ts +++ b/frontend/src/utils/sqlEditorTransaction.test.ts @@ -6,6 +6,7 @@ import { shouldUseSqlEditorManagedTransaction, shouldUseSqlEditorManagedTransactionForType, } from './sqlEditorTransaction'; +import { findSqlStatementRanges } from './sqlStatementSelection'; describe('sqlEditorTransaction', () => { it('keeps regular DML in a managed transaction', () => { @@ -14,6 +15,25 @@ describe('sqlEditorTransaction', () => { expect(shouldUseSqlEditorManagedTransaction(['DELETE FROM users WHERE id = 1'])).toBe(true); }); + it('keeps DML with a trailing line comment in a managed transaction', () => { + const sql = 'DELETE FROM users WHERE id = 1; -- keep this operation pending'; + const statements = findSqlStatementRanges(sql).map((range) => range.text); + + expect(statements).toEqual(['DELETE FROM users WHERE id = 1']); + expect(shouldUseSqlEditorManagedTransactionForType('mysql', statements)).toBe(true); + }); + + it('uses dialect-specific rules for compact line comments', () => { + const sql = 'DELETE FROM users WHERE id = 1;--comment'; + const postgresStatements = findSqlStatementRanges(sql, 'postgres').map((range) => range.text); + const mysqlStatements = findSqlStatementRanges(sql, 'mysql').map((range) => range.text); + + expect(postgresStatements).toEqual(['DELETE FROM users WHERE id = 1']); + expect(shouldUseSqlEditorManagedTransactionForType('postgres', postgresStatements)).toBe(true); + expect(mysqlStatements).toEqual(['DELETE FROM users WHERE id = 1', '--comment']); + expect(shouldUseSqlEditorManagedTransactionForType('mysql', mysqlStatements)).toBe(false); + }); + it('classifies WITH statements by their top-level operation', () => { expect(resolveSqlEditorOperationKeyword('WITH target AS (SELECT id FROM users) SELECT * FROM target')).toBe('select'); expect(resolveSqlEditorOperationKeyword('WITH target AS (SELECT id FROM users) UPDATE users SET synced = 1')).toBe('update'); diff --git a/frontend/src/utils/sqlStatementSelection.test.ts b/frontend/src/utils/sqlStatementSelection.test.ts index e0e73301..b38a361f 100644 --- a/frontend/src/utils/sqlStatementSelection.test.ts +++ b/frontend/src/utils/sqlStatementSelection.test.ts @@ -28,6 +28,51 @@ describe('sqlStatementSelection', () => { ]); }); + it('drops comment-only ranges after a terminated statement', () => { + const sql = [ + 'DELETE FROM users WHERE id = 1; -- keep this operation pending', + '/* trailing explanation */', + ].join('\n'); + + expect(findSqlStatementRanges(sql).map((range) => range.text)).toEqual([ + 'DELETE FROM users WHERE id = 1', + ]); + expect(findSqlStatementRanges('DELETE FROM users WHERE id = 1;--').map((range) => range.text)).toEqual([ + 'DELETE FROM users WHERE id = 1', + ]); + }); + + it('keeps executable MySQL comments as statements', () => { + const sql = '/*!40101 SET @OLD_CHARACTER_SET_CLIENT=@@CHARACTER_SET_CLIENT */;'; + + expect(findSqlStatementRanges(sql, 'mysql').map((range) => range.text)).toEqual([ + '/*!40101 SET @OLD_CHARACTER_SET_CLIENT=@@CHARACTER_SET_CLIENT */', + ]); + }); + + it('uses dialect-specific executable block comment rules', () => { + const mysqlComment = '/*!40101 SET @OLD_CHARACTER_SET_CLIENT=@@CHARACTER_SET_CLIENT */;'; + const mariaDbComment = '/*M!100100 SET @OLD_CHARACTER_SET_CLIENT=@@CHARACTER_SET_CLIENT */;'; + + expect(findSqlStatementRanges(mysqlComment, 'postgres')).toEqual([]); + expect(findSqlStatementRanges(mariaDbComment, 'mysql')).toEqual([]); + expect(findSqlStatementRanges(mariaDbComment, 'mariadb').map((range) => range.text)).toEqual([ + '/*M!100100 SET @OLD_CHARACTER_SET_CLIENT=@@CHARACTER_SET_CLIENT */', + ]); + }); + + it('uses dialect-specific hash comment rules', () => { + const sql = 'DELETE FROM users WHERE id = 1; #comment'; + + expect(findSqlStatementRanges(sql, 'mysql').map((range) => range.text)).toEqual([ + 'DELETE FROM users WHERE id = 1', + ]); + expect(findSqlStatementRanges(sql, 'postgres').map((range) => range.text)).toEqual([ + 'DELETE FROM users WHERE id = 1', + '#comment', + ]); + }); + it('keeps Oracle anonymous PL/SQL blocks as one executable statement', () => { const plsql = [ 'BEGIN', diff --git a/frontend/src/utils/sqlStatementSelection.ts b/frontend/src/utils/sqlStatementSelection.ts index bb68fe64..bc92f8bf 100644 --- a/frontend/src/utils/sqlStatementSelection.ts +++ b/frontend/src/utils/sqlStatementSelection.ts @@ -23,6 +23,73 @@ const isSqlIdentifierStart = (ch: string): boolean => /^[A-Za-z_]$/.test(ch); const isSqlIdentifierPart = (ch: string): boolean => /^[A-Za-z0-9_$#]$/.test(ch); +const normalizeSqlLexicalDbType = (dbType: string): string => { + const normalized = String(dbType || '').trim().toLowerCase(); + if (normalized === 'doris') return 'diros'; + if (normalized === 'greatdb' || normalized === 'gdb') return 'goldendb'; + return normalized; +}; + +const MYSQL_DASH_COMMENT_DIALECTS = new Set([ + 'mysql', 'mariadb', 'oceanbase', 'diros', 'starrocks', 'goldendb', 'sphinx', 'tidb', +]); + +const supportsSqlHashLineComment = (dbType: string): boolean => { + const normalized = normalizeSqlLexicalDbType(dbType); + return !normalized || normalized === 'clickhouse' || MYSQL_DASH_COMMENT_DIALECTS.has(normalized); +}; + +const isSqlDashLineCommentStart = (dbType: string, next2: string): boolean => { + const normalized = normalizeSqlLexicalDbType(dbType); + return !MYSQL_DASH_COMMENT_DIALECTS.has(normalized) || !next2 || isWhitespace(next2); +}; + +const isExecutableSqlBlockComment = (sql: string, index: number, dbType: string): boolean => { + const isMySqlVersionComment = sql.startsWith('/*!', index); + const isMariaDbVersionComment = sql.slice(index, index + 4).toLowerCase() === '/*m!'; + if (!isMySqlVersionComment && !isMariaDbVersionComment) { + return false; + } + const normalized = normalizeSqlLexicalDbType(dbType); + if (!normalized) { + return true; + } + if (isMariaDbVersionComment) { + return normalized === 'mariadb'; + } + return MYSQL_DASH_COMMENT_DIALECTS.has(normalized); +}; + +const hasExecutableSqlStatementContent = (sql: string, dbType = ''): boolean => { + const text = String(sql || ''); + let index = 0; + while (index < text.length) { + const ch = text[index]; + const next = index + 1 < text.length ? text[index + 1] : ''; + const next2 = index + 2 < text.length ? text[index + 2] : ''; + if (isWhitespace(ch)) { + index++; + continue; + } + if ((ch === '#' && supportsSqlHashLineComment(dbType)) + || (ch === '-' && next === '-' && isSqlDashLineCommentStart(dbType, next2))) { + const lineEnd = text.indexOf('\n', index + (ch === '#' ? 1 : 2)); + index = lineEnd < 0 ? text.length : lineEnd + 1; + continue; + } + if (ch === '/' && next === '*') { + if (isExecutableSqlBlockComment(text, index, dbType)) { + return true; + } + const blockEnd = text.indexOf('*/', index + 2); + index = blockEnd < 0 ? text.length : blockEnd + 2; + continue; + } + return true; + } + return false; +}; + const skipSqlWhitespaceAndComments = (text: string, position: number): number => { let index = position; while (index < text.length) { @@ -206,7 +273,7 @@ const isPlsqlControlEnd = (text: string, tokenEnd: number): boolean => ( ['if', 'loop', 'case'].includes(nextSqlSignificantToken(text, tokenEnd)) ); -const trimStatementRange = (sql: string, start: number, end: number): SqlStatementRange | null => { +const trimStatementRange = (sql: string, start: number, end: number, dbType = ''): SqlStatementRange | null => { let nextStart = Math.max(0, start); let nextEnd = Math.min(sql.length, Math.max(start, end)); @@ -221,6 +288,10 @@ const trimStatementRange = (sql: string, start: number, end: number): SqlStateme return null; } + if (!hasExecutableSqlStatementContent(sql.slice(nextStart, nextEnd), dbType)) { + return null; + } + return { start: nextStart, end: nextEnd, @@ -228,7 +299,7 @@ const trimStatementRange = (sql: string, start: number, end: number): SqlStateme }; }; -export const findSqlStatementRanges = (sql: string): SqlStatementRange[] => { +export const findSqlStatementRanges = (sql: string, dbType = ''): SqlStatementRange[] => { const text = String(sql || '').replace(/\r\n/g, '\n'); const ranges: SqlStatementRange[] = []; @@ -247,7 +318,7 @@ export const findSqlStatementRanges = (sql: string): SqlStatementRange[] => { let justClosedPLSQLBlock = false; const push = (end: number) => { - const range = trimStatementRange(text, statementStart, end); + const range = trimStatementRange(text, statementStart, end, dbType); if (range) { ranges.push(range); } @@ -256,7 +327,6 @@ export const findSqlStatementRanges = (sql: string): SqlStatementRange[] => { for (let index = 0; index < text.length; index++) { const ch = text[index]; const next = index + 1 < text.length ? text[index + 1] : ''; - const prev = index > 0 ? text[index - 1] : ''; const next2 = index + 2 < text.length ? text[index + 2] : ''; if (dollarTag) { @@ -300,11 +370,11 @@ export const findSqlStatementRanges = (sql: string): SqlStatementRange[] => { continue; } } - if (ch === '#') { + if (ch === '#' && supportsSqlHashLineComment(dbType)) { inLineComment = true; continue; } - if (ch === '-' && next === '-' && (index === 0 || isWhitespace(prev)) && (next2 === '' || isWhitespace(next2))) { + if (ch === '-' && next === '-' && isSqlDashLineCommentStart(dbType, next2)) { index++; inLineComment = true; continue; @@ -409,10 +479,10 @@ export const findSqlStatementRanges = (sql: string): SqlStatementRange[] => { return ranges; }; -export const resolveCurrentSqlStatementRange = (sql: string, cursorOffset: number): SqlStatementRange | null => { +export const resolveCurrentSqlStatementRange = (sql: string, cursorOffset: number, dbType = ''): SqlStatementRange | null => { const text = String(sql || '').replace(/\r\n/g, '\n'); const offset = Math.max(0, Math.min(text.length, Number.isFinite(cursorOffset) ? cursorOffset : 0)); - const ranges = findSqlStatementRanges(text); + const ranges = findSqlStatementRanges(text, dbType); if (ranges.length === 0) { return null; } @@ -439,6 +509,7 @@ export const resolveExecutableSql = ( sql: string, cursorOffset: number, selectedSql = '', + dbType = '', ): SqlExecutionSelection | null => { const selected = String(selectedSql || '').trim(); if (selected) { @@ -447,7 +518,7 @@ export const resolveExecutableSql = ( const text = String(sql || '').replace(/\r\n/g, '\n'); const offset = Math.max(0, Math.min(text.length, Number.isFinite(cursorOffset) ? cursorOffset : 0)); - const ranges = findSqlStatementRanges(text); + const ranges = findSqlStatementRanges(text, dbType); const statement = ranges.find((range) => offset >= range.start && offset <= range.end); if (statement?.text.trim()) { return { sql: statement.text, source: 'statement' }; diff --git a/internal/app/connection_readonly.go b/internal/app/connection_readonly.go index 39226386..e168fab1 100644 --- a/internal/app/connection_readonly.go +++ b/internal/app/connection_readonly.go @@ -266,8 +266,9 @@ func ensureConnectionAllowsQueryWithText(config connection.ConnectionConfig, que if !isConnectionScriptExecutionRestricted(config) { return nil } - for _, statement := range splitSQLStatements(query) { - if trimmed := strings.TrimSpace(statement); trimmed != "" && !isReadOnlySQLQuery(resolveDDLDBType(config), trimmed) { + dbType := resolveDDLDBType(config) + for _, statement := range splitSQLStatementsForDialect(dbType, query) { + if trimmed := strings.TrimSpace(statement); trimmed != "" && !isReadOnlySQLQuery(dbType, trimmed) { return errors.New(readOnlyConnectionQueryBlockedMessageWithText(text)) } } diff --git a/internal/app/methods_db.go b/internal/app/methods_db.go index 5eff733f..f7d5481e 100644 --- a/internal/app/methods_db.go +++ b/internal/app/methods_db.go @@ -909,7 +909,7 @@ type dbQueryMultiAuditOptions struct { } func containsSQLAuditWrite(dbType string, query string) bool { - statements := splitSQLStatements(query) + statements := splitSQLStatementsForDialect(dbType, query) if len(statements) == 0 { return !isReadOnlySQLQuery(dbType, query) } @@ -1204,7 +1204,7 @@ func (a *App) dbQueryMulti( // 注意：原生 conn.Query() 执行写操作（UPDATE/INSERT/DELETE）时， // sql.Rows 不暴露 RowsAffected，导致影响行数丢失。 // 因此仅在全部语句皆为读操作时才使用原生路径。 - statements := splitSQLStatements(query) + statements := splitSQLStatementsForDialect(resolvedDBType, query) statementCount := 0 for _, statement := range statements { if strings.TrimSpace(statement) != "" { diff --git a/internal/app/methods_db_multi_test.go b/internal/app/methods_db_multi_test.go index d6bb520a..ba64f676 100644 --- a/internal/app/methods_db_multi_test.go +++ b/internal/app/methods_db_multi_test.go @@ -1075,6 +1075,53 @@ func TestDBQueryMultiTransactionalKeepsDMLTransactionOpenUntilCommit(t *testing. } } +func TestDBQueryMultiTransactionalKeepsTrailingCommentInsideManagedTransaction(t *testing.T) { + originalNewDatabaseFunc := newDatabaseFunc + t.Cleanup(func() { + newDatabaseFunc = originalNewDatabaseFunc + }) + + statement := "DELETE FROM users WHERE id = 1" + fakeDB := &fakeBatchWriteDB{ + execAffected: map[string]int64{statement: 1}, + } + newDatabaseFunc = func(dbType string) (db.Database, error) { + return fakeDB, nil + } + + app := NewAppWithSecretStore(secretstore.NewUnavailableStore("test")) + config := connection.ConnectionConfig{Type: "mysql", Host: "127.0.0.1", Port: 3306, User: "root"} + result := app.DBQueryMultiTransactional( + config, + "main", + statement+"; -- keep this operation pending", + "tx-trailing-comment", + ) + + if !result.Success { + t.Fatalf("expected transactional query success, got failure: %s", result.Message) + } + if result.TransactionID == "" || !result.TransactionPending { + t.Fatalf("expected trailing comment to preserve pending transaction, got id=%q pending=%v", result.TransactionID, result.TransactionPending) + } + if strings.Contains(result.Message, "逐条执行") { + t.Fatalf("expected trailing comment to avoid sequential fallback message, got %q", result.Message) + } + wantExecs := []string{"START TRANSACTION", statement} + if !reflect.DeepEqual(fakeDB.execQueries, wantExecs) { + t.Fatalf("expected exec queries %#v, got %#v", wantExecs, fakeDB.execQueries) + } + resultSets, ok := result.Data.([]connection.ResultSetData) + if !ok || len(resultSets) != 1 { + t.Fatalf("expected one DML result set, got %T %#v", result.Data, result.Data) + } + + rollbackResult := app.DBRollbackTransaction(result.TransactionID) + if !rollbackResult.Success { + t.Fatalf("expected rollback success, got failure: %s", rollbackResult.Message) + } +} + func TestDBQueryMultiInTransactionReusesPendingManagedSessionForReadQueries(t *testing.T) { originalNewDatabaseFunc := newDatabaseFunc t.Cleanup(func() { diff --git a/internal/app/methods_db_transaction.go b/internal/app/methods_db_transaction.go index 90d4357a..feec402f 100644 --- a/internal/app/methods_db_transaction.go +++ b/internal/app/methods_db_transaction.go @@ -104,7 +104,7 @@ func (a *App) DBQueryMultiTransactional(config connection.ConnectionConfig, dbNa CommitMode: "pending", BoundaryMode: transactionBoundaryMode, SQL: query, - StatementCount: countSQLAuditStatements(query), + StatementCount: countSQLAuditStatements(transactionDBType, query), Err: sqlAuditErrorFromResult(result), }) }() @@ -243,7 +243,7 @@ func (a *App) DBQueryMultiTransactional(config connection.ConnectionConfig, dbNa BoundaryMode: transactionBoundaryMode, }) - statements := splitSQLStatements(query) + statements := splitSQLStatementsForDialect(transactionDBType, query) queryStartedAt := time.Now() statementAuditEvents := make([]sqlaudit.Event, 0, len(statements)) resultSets, err := executeManagedSQLTransactionStatementsWithObserver( @@ -363,7 +363,7 @@ func (a *App) DBQueryMultiInTransaction(transactionID string, query string, quer a.recordQueryExecution(runConfig, "", tx.dbType, query, durationMs, 0, queryResultRowsReturned(result)) }() query = sanitizeSQLForPgLike(tx.dbType, query) - statements := splitSQLStatements(query) + statements := splitSQLStatementsForDialect(tx.dbType, query) ctx, cancel := newQueryExecutionContext(runConfig) defer cancel() @@ -623,7 +623,7 @@ func shouldUseManagedSQLTransaction(dbType string, query string) bool { if isManagedSQLTransactionUnsupportedType(dbType) { return false } - statements := splitSQLStatements(query) + statements := splitSQLStatementsForDialect(dbType, query) hasManagedWrite := false for _, stmt := range statements { stmt = strings.TrimSpace(stmt) diff --git a/internal/app/methods_db_transaction_test.go b/internal/app/methods_db_transaction_test.go index f57f0695..6d6e4b53 100644 --- a/internal/app/methods_db_transaction_test.go +++ b/internal/app/methods_db_transaction_test.go @@ -50,3 +50,17 @@ END;` t.Fatal("expected Oracle read-only anonymous block to stay unmanaged") } } + +func TestShouldUseManagedSQLTransaction_UsesDialectCommentRules(t *testing.T) { + t.Parallel() + + if !shouldUseManagedSQLTransaction("mysql", "DELETE FROM users WHERE id = 1; -- pending") { + t.Fatal("expected a valid MySQL trailing comment to preserve the managed transaction") + } + if shouldUseManagedSQLTransaction("mysql", "DELETE FROM users WHERE id = 1;--comment") { + t.Fatal("expected compact MySQL double-dash text to remain an executable statement") + } + if !shouldUseManagedSQLTransaction("postgres", "DELETE FROM users WHERE id = 1; /*! MySQL-only comment */") { + t.Fatal("expected PostgreSQL to ignore a MySQL-only block comment") + } +} diff --git a/internal/app/methods_sql_audit.go b/internal/app/methods_sql_audit.go index 5853b96c..4200dcea 100644 --- a/internal/app/methods_sql_audit.go +++ b/internal/app/methods_sql_audit.go @@ -361,7 +361,7 @@ func (a *App) recordSQLAuditQuery(input sqlAuditQueryInput) { status := sqlAuditStatusFromResult(input.Result) statementCount := input.StatementCount if statementCount <= 0 { - statementCount = countSQLAuditStatements(input.SQL) + statementCount = countSQLAuditStatements(input.DBType, input.SQL) } event := sqlaudit.Event{ EventType: "query", @@ -787,9 +787,9 @@ func durationMilliseconds(duration time.Duration) int64 { return milliseconds } -func countSQLAuditStatements(sql string) int { +func countSQLAuditStatements(dbType string, sql string) int { count := 0 - for _, statement := range splitSQLStatements(sql) { + for _, statement := range splitSQLStatementsForDialect(dbType, sql) { if strings.TrimSpace(statement) != "" { count++ } diff --git a/internal/app/sql_inspect.go b/internal/app/sql_inspect.go index 11c3848d..fbdc98aa 100644 --- a/internal/app/sql_inspect.go +++ b/internal/app/sql_inspect.go @@ -18,7 +18,7 @@ type SQLInspection struct { // InspectSQL 基于现有 SQL 拆分与只读判定逻辑，为外部调用方提供安全边界判断。 func InspectSQL(dbType string, sql string) SQLInspection { - statements := splitSQLStatements(sql) + statements := splitSQLStatementsForDialect(dbType, sql) result := SQLInspection{ ReadOnly: true, Statements: make([]SQLStatementInspection, 0, len(statements)), diff --git a/internal/app/sql_split.go b/internal/app/sql_split.go index d7cd49e6..c6b47942 100644 --- a/internal/app/sql_split.go +++ b/internal/app/sql_split.go @@ -34,7 +34,7 @@ func splitSQLStatementsForDialect(dbType, sql string) []string { push := func() { s := strings.TrimSpace(cur.String()) - if s != "" { + if s != "" && hasExecutableSQLStatementContent(dbType, s) { statements = append(statements, s) } cur.Reset() @@ -255,9 +255,65 @@ func splitSQLStatementsForDialect(dbType, sql string) []string { return statements } +func hasExecutableSQLStatementContent(dbType, statement string) bool { + for i := 0; i < len(statement); { + switch statement[i] { + case ' ', '\t', '\n', '\r', '\f': + i++ + continue + case '-': + if i+1 < len(statement) && statement[i+1] == '-' && isSQLDashLineCommentStart(dbType, statement, i) { + i = scanSQLLineCommentEnd(statement, i+2) + continue + } + case '#': + if supportsSQLHashLineComment(dbType) { + i = scanSQLLineCommentEnd(statement, i+1) + continue + } + case '/': + if i+1 < len(statement) && statement[i+1] == '*' { + remaining := statement[i:] + if supportsSQLExecutableBlockComment(dbType, remaining) { + return true + } + blockEnd := strings.Index(remaining[2:], "*/") + if blockEnd < 0 { + return false + } + i += blockEnd + 4 + continue + } + } + return true + } + return false +} + +func supportsSQLExecutableBlockComment(dbType, remaining string) bool { + isMySQLVersionComment := strings.HasPrefix(remaining, "/*!") + isMariaDBVersionComment := len(remaining) >= 4 && strings.EqualFold(remaining[:4], "/*m!") + if !isMySQLVersionComment && !isMariaDBVersionComment { + return false + } + normalized := normalizeExplainLexicalDBType(dbType) + if normalized == "" { + return true + } + if isMariaDBVersionComment { + return normalized == "mariadb" + } + switch normalized { + case "mysql", "mariadb", "oceanbase", "diros", "starrocks", "goldendb", "sphinx", "tidb": + return true + default: + return false + } +} + func isSQLDashLineCommentStart(dbType, text string, index int) bool { switch normalizeExplainLexicalDBType(dbType) { - case "mysql", "mariadb", "oceanbase", "diros", "starrocks", "goldendb": + case "mysql", "mariadb", "oceanbase", "diros", "starrocks", "goldendb", "sphinx", "tidb": return isMySQLDashCommentStart(text, index) default: return true @@ -270,7 +326,7 @@ func supportsSQLHashLineComment(dbType string) bool { return true } switch normalized { - case "mysql", "mariadb", "oceanbase", "diros", "starrocks", "goldendb", "clickhouse": + case "mysql", "mariadb", "oceanbase", "diros", "starrocks", "goldendb", "sphinx", "tidb", "clickhouse": return true default: return false diff --git a/internal/app/sql_split_stream.go b/internal/app/sql_split_stream.go index b2b4854d..f0b6caaa 100644 --- a/internal/app/sql_split_stream.go +++ b/internal/app/sql_split_stream.go @@ -26,6 +26,15 @@ type sqlStreamSplitter struct { closedPLSQL bool } +func (s *sqlStreamSplitter) takeStatement() string { + stmt := strings.TrimSpace(s.cur.String()) + s.cur.Reset() + if !hasExecutableSQLStatementContent("", stmt) { + return "" + } + return stmt +} + // Feed 将一个 chunk 喂入拆分器，返回在此 chunk 中完成的 SQL 语句列表。 func (s *sqlStreamSplitter) Feed(chunk []byte) []string { var statements []string @@ -209,11 +218,10 @@ func (s *sqlStreamSplitter) Feed(chunk []byte) []string { s.pending = text[i:] break } - stmt := strings.TrimSpace(s.cur.String()) + stmt := s.takeStatement() if stmt != "" { statements = append(statements, stmt) } - s.cur.Reset() s.closedPLSQL = false i = lineEnd continue @@ -254,19 +262,17 @@ func (s *sqlStreamSplitter) Feed(chunk []byte) []string { } if s.closedPLSQL { s.cur.WriteByte(ch) - stmt := strings.TrimSpace(s.cur.String()) + stmt := s.takeStatement() if stmt != "" { statements = append(statements, stmt) } - s.cur.Reset() s.closedPLSQL = false continue } - stmt := strings.TrimSpace(s.cur.String()) + stmt := s.takeStatement() if stmt != "" { statements = append(statements, stmt) } - s.cur.Reset() continue } // 全角分号 @@ -282,20 +288,18 @@ func (s *sqlStreamSplitter) Feed(chunk []byte) []string { } if s.closedPLSQL { s.cur.WriteString("；") - stmt := strings.TrimSpace(s.cur.String()) + stmt := s.takeStatement() if stmt != "" { statements = append(statements, stmt) } - s.cur.Reset() s.closedPLSQL = false i += 2 continue } - stmt := strings.TrimSpace(s.cur.String()) + stmt := s.takeStatement() if stmt != "" { statements = append(statements, stmt) } - s.cur.Reset() i += 2 continue } @@ -312,8 +316,7 @@ func (s *sqlStreamSplitter) Flush() string { if (s.closedPLSQL || strings.TrimSpace(s.cur.String()) == "") && sqlStreamCurrentLineWhitespaceOnly(&s.cur) { if _, standalone, _ := scanSQLStandaloneSlashLineSuffix(s.pending, 0); standalone { s.pending = "" - stmt := strings.TrimSpace(s.cur.String()) - s.cur.Reset() + stmt := s.takeStatement() s.closedPLSQL = false return stmt } @@ -321,8 +324,7 @@ func (s *sqlStreamSplitter) Flush() string { s.cur.WriteString(s.pending) s.pending = "" } - stmt := strings.TrimSpace(s.cur.String()) - s.cur.Reset() + stmt := s.takeStatement() if stmt == "/" { return "" } diff --git a/internal/app/sql_split_stream_test.go b/internal/app/sql_split_stream_test.go new file mode 100644 index 00000000..0209428e --- /dev/null +++ b/internal/app/sql_split_stream_test.go @@ -0,0 +1,41 @@ +package app + +import ( + "reflect" + "strings" + "testing" +) + +func TestStreamSQLFileDropsCommentOnlyTail(t *testing.T) { + t.Parallel() + + var statements []string + count, err := streamSQLFile( + strings.NewReader("DELETE FROM users WHERE id = 1; -- keep this operation pending"), + func(_ int, statement string) error { + statements = append(statements, statement) + return nil + }, + ) + if err != nil { + t.Fatalf("streamSQLFile returned error: %v", err) + } + if count != 1 { + t.Fatalf("expected one executable statement, got %d", count) + } + want := []string{"DELETE FROM users WHERE id = 1"} + if !reflect.DeepEqual(statements, want) { + t.Fatalf("expected statements %#v, got %#v", want, statements) + } +} + +func TestSQLStreamSplitterPreservesExecutableMySQLComment(t *testing.T) { + t.Parallel() + + splitter := &sqlStreamSplitter{} + got := splitter.Feed([]byte("/*!40101 SET @OLD_CHARACTER_SET_CLIENT=@@CHARACTER_SET_CLIENT */;")) + want := []string{"/*!40101 SET @OLD_CHARACTER_SET_CLIENT=@@CHARACTER_SET_CLIENT */"} + if !reflect.DeepEqual(got, want) { + t.Fatalf("expected statements %#v, got %#v", want, got) + } +} diff --git a/internal/app/sql_split_test.go b/internal/app/sql_split_test.go index 8b4f6225..dbe4aeff 100644 --- a/internal/app/sql_split_test.go +++ b/internal/app/sql_split_test.go @@ -32,6 +32,65 @@ func TestSplitSQLStatements_LineComment(t *testing.T) { } } +func TestSplitSQLStatements_DropsCommentOnlyTail(t *testing.T) { + t.Parallel() + + statement := "DELETE FROM users WHERE id = 1" + tests := []struct { + name string + query string + }{ + {name: "bare line comment marker", query: statement + ";--"}, + {name: "line comment", query: statement + "; -- keep this operation pending"}, + {name: "hash comment", query: statement + ";\n# keep this operation pending"}, + {name: "block comment", query: statement + ";\n/* keep this operation pending */"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := splitSQLStatementsForDialect("mysql", tt.query) + want := []string{statement} + if !reflect.DeepEqual(got, want) { + t.Fatalf("splitSQLStatementsForDialect(mysql, %q) = %#v, want %#v", tt.query, got, want) + } + }) + } +} + +func TestSplitSQLStatements_PreservesExecutableMySQLComment(t *testing.T) { + t.Parallel() + + query := "/*!40101 SET @OLD_CHARACTER_SET_CLIENT=@@CHARACTER_SET_CLIENT */;" + got := splitSQLStatementsForDialect("mysql", query) + want := []string{"/*!40101 SET @OLD_CHARACTER_SET_CLIENT=@@CHARACTER_SET_CLIENT */"} + if !reflect.DeepEqual(got, want) { + t.Fatalf("splitSQLStatementsForDialect(mysql, %q) = %#v, want %#v", query, got, want) + } +} + +func TestSplitSQLStatements_UsesDialectSpecificExecutableCommentRules(t *testing.T) { + t.Parallel() + + mysqlComment := "/*!40101 SET @OLD_CHARACTER_SET_CLIENT=@@CHARACTER_SET_CLIENT */;" + mariaDBComment := "/*M!100100 SET @OLD_CHARACTER_SET_CLIENT=@@CHARACTER_SET_CLIENT */;" + if got := splitSQLStatementsForDialect("postgres", mysqlComment); len(got) != 0 { + t.Fatalf("expected PostgreSQL to drop MySQL-only executable comment, got %#v", got) + } + if got := splitSQLStatementsForDialect("mysql", mariaDBComment); len(got) != 0 { + t.Fatalf("expected MySQL to drop MariaDB-only executable comment, got %#v", got) + } + want := []string{"/*M!100100 SET @OLD_CHARACTER_SET_CLIENT=@@CHARACTER_SET_CLIENT */"} + if got := splitSQLStatementsForDialect("mariadb", mariaDBComment); !reflect.DeepEqual(got, want) { + t.Fatalf("expected MariaDB statements %#v, got %#v", want, got) + } + + statement := "DELETE FROM users WHERE id = 1" + want = []string{statement, "#comment"} + if got := splitSQLStatementsForDialect("postgres", statement+"; #comment"); !reflect.DeepEqual(got, want) { + t.Fatalf("expected PostgreSQL statements %#v, got %#v", want, got) + } +} + func TestSplitSQLStatements_BlockComment(t *testing.T) { input := "SELECT /* ; */ 1; SELECT 2" got := splitSQLStatements(input)