From 89639e36bc56fa1a2f3cc1d8bae7b44964a8bdb3 Mon Sep 17 00:00:00 2001
From: Syngnat <yangguofeng919@gmail.com>
Date: Wed, 10 Jun 2026 19:13:54 +0800
Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9B=20fix(query-editor):=20=E4=BF=AE?=
 =?UTF-8?q?=E6=AD=A3=20SQL=20=E7=BC=96=E8=BE=91=E5=99=A8=20DML=20=E4=BA=8B?=
 =?UTF-8?q?=E5=8A=A1=E8=AF=86=E5=88=AB?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- 统一前后端 DML 与数据修改 CTE 的受管事务判断

- 保留数据修改 CTE 返回行并补充事务回归测试

- 明确 SQL 编辑器事务提交策略文案
---
 .../QueryEditor.external-sql-save.test.tsx    | 55 +++++++++++++++++--
 frontend/src/components/QueryEditor.tsx       |  8 +--
 .../src/utils/sqlEditorTransaction.test.ts    |  6 ++
 frontend/src/utils/sqlEditorTransaction.ts    | 45 +++++++++++----
 frontend/src/v2-theme.css                     |  6 +-
 internal/app/methods_db.go                    |  3 +
 internal/app/methods_db_multi_test.go         | 53 ++++++++++++++++++
 internal/app/sql_sanitize.go                  | 52 +++++++++++++-----
 internal/app/sql_sanitize_test.go             | 15 +++++
 9 files changed, 206 insertions(+), 37 deletions(-)

diff --git a/frontend/src/components/QueryEditor.external-sql-save.test.tsx b/frontend/src/components/QueryEditor.external-sql-save.test.tsx
index a21a02b..0f062af 100644
--- a/frontend/src/components/QueryEditor.external-sql-save.test.tsx
+++ b/frontend/src/components/QueryEditor.external-sql-save.test.tsx
@@ -48,6 +48,8 @@ const storeState = vi.hoisted(() => ({
     autoCommitDelayMs: 5000,
   },
   setSqlEditorTransactionOptions: vi.fn(),
+  sqlEditorPendingTransactions: {} as Record<string, unknown>,
+  setSqlEditorPendingTransaction: vi.fn(),
   shortcutOptions: {
     runQuery: {
       mac: { enabled: false, combo: '' },
@@ -487,6 +489,15 @@ describe('QueryEditor external SQL save', () => {
     storeState.setSqlEditorTransactionOptions.mockImplementation((options: Record<string, unknown>) => {
       storeState.sqlEditorTransactionOptions = { ...storeState.sqlEditorTransactionOptions, ...options };
     });
+    storeState.sqlEditorPendingTransactions = {};
+    storeState.setSqlEditorPendingTransaction.mockReset();
+    storeState.setSqlEditorPendingTransaction.mockImplementation((tabId: string, transaction: unknown) => {
+      if (!transaction) {
+        delete storeState.sqlEditorPendingTransactions[tabId];
+        return;
+      }
+      storeState.sqlEditorPendingTransactions[tabId] = transaction;
+    });
     messageApi.success.mockReset();
     messageApi.error.mockReset();
     messageApi.warning.mockReset();
@@ -2257,6 +2268,40 @@ describe('QueryEditor external SQL save', () => {
     expect(backendApp.DBQueryMultiTransactional).not.toHaveBeenCalled();
   });
 
+  it('runs SQL editor data-changing CTEs through a pending managed transaction', async () => {
+    const sql = 'WITH moved AS (DELETE FROM audit_logs WHERE created_at < NOW() RETURNING id) SELECT * FROM moved';
+    backendApp.DBQueryMultiTransactional.mockResolvedValueOnce({
+      success: true,
+      transactionId: 'tx-write-cte',
+      transactionPending: true,
+      data: [
+        { columns: ['affectedRows'], rows: [{ affectedRows: 3 }], statementIndex: 1 },
+      ],
+    });
+
+    let renderer!: ReactTestRenderer;
+    await act(async () => {
+      renderer = create(<QueryEditor tab={createTab({ query: sql })} />);
+    });
+
+    await act(async () => {
+      await findButton(renderer!, '运行').props.onClick();
+    });
+    await act(async () => {
+      await Promise.resolve();
+      await Promise.resolve();
+    });
+
+    expect(backendApp.DBQueryMultiTransactional).toHaveBeenCalledWith(
+      expect.anything(),
+      'main',
+      expect.stringContaining('DELETE FROM audit_logs'),
+      'query-1',
+    );
+    expect(backendApp.DBQueryMulti).not.toHaveBeenCalled();
+    expect(textContent(renderer!.root)).toContain('事务待提交');
+  });
+
   it('auto commits SQL editor DML transactions after the configured delay', async () => {
     vi.useFakeTimers();
     storeState.sqlEditorTransactionOptions = {
@@ -3521,9 +3566,9 @@ describe('QueryEditor external SQL save', () => {
     expect(source).toContain('gn-v2-query-toolbar-max-rows-select');
     expect(source).toContain('gn-v2-query-toolbar-transaction-mode-select');
     expect(source).toContain('gn-v2-query-toolbar-transaction-delay-select');
-    expect(source).toContain('这里仅选择该事务执行成功后的提交方式');
-    expect(source).toContain("label: '提交：手动'");
-    expect(source).toContain("label: '提交：自动'");
+    expect(source).toContain('这里仅选择事务执行成功后的 COMMIT 方式');
+    expect(source).toContain("label: '事务：手动提交'");
+    expect(source).toContain("label: '事务：自动提交'");
     expect(source).toContain('gn-v2-query-toolbar-action-group');
     expect(source).toContain('style={isV2Ui ? undefined : { width: 150 }}');
     expect(source).toContain('style={isV2Ui ? undefined : { width: 200 }}');
@@ -3538,11 +3583,11 @@ describe('QueryEditor external SQL save', () => {
     expect(css).toContain('display: inline-flex !important;');
     expect(css).toContain('gap: 6px;');
     expect(css).toContain('margin-left: 0 !important;');
-    expect(css).toContain('max-width: 720px;');
+    expect(css).toContain('max-width: 760px;');
     expect(css).toContain('width: 140px !important;');
     expect(css).toContain('width: 166px !important;');
     expect(css).toContain('width: 132px !important;');
-    expect(css).toContain('width: 112px !important;');
+    expect(css).toContain('width: 142px !important;');
     expect(css).toContain('width: 82px !important;');
     expect(css).toContain('width: 34px !important;');
     expect(css).toContain('@media (max-width: 900px)');
diff --git a/frontend/src/components/QueryEditor.tsx b/frontend/src/components/QueryEditor.tsx
index c441be8..16b0263 100644
--- a/frontend/src/components/QueryEditor.tsx
+++ b/frontend/src/components/QueryEditor.tsx
@@ -5328,15 +5328,15 @@ const QueryEditor: React.FC<{ tab: TabData; isActive?: boolean }> = ({ tab, isAc
                   ]}
               />
           </Tooltip>
-          <Tooltip title="SQL 编辑器执行 INSERT/UPDATE/DELETE 等 DML 时始终启用事务；这里仅选择该事务执行成功后的提交方式。">
+          <Tooltip title="SQL 编辑器执行 INSERT/UPDATE/DELETE/MERGE/REPLACE 等 DML 时会先开启受管事务；这里仅选择事务执行成功后的 COMMIT 方式。">
               <Select
                   className={isV2Ui ? 'gn-v2-query-toolbar-select gn-v2-query-toolbar-transaction-mode-select' : undefined}
-                  style={isV2Ui ? undefined : { width: 128 }}
+                  style={isV2Ui ? undefined : { width: 150 }}
                   value={sqlEditorCommitMode}
                   onChange={(mode) => setSqlEditorTransactionOptions({ commitMode: mode === 'auto' ? 'auto' : 'manual' })}
                   options={[
-                      { label: '提交：手动', value: 'manual' },
-                      { label: '提交：自动', value: 'auto' },
+                      { label: '事务：手动提交', value: 'manual' },
+                      { label: '事务：自动提交', value: 'auto' },
                   ]}
               />
           </Tooltip>
diff --git a/frontend/src/utils/sqlEditorTransaction.test.ts b/frontend/src/utils/sqlEditorTransaction.test.ts
index 5d8189c..efd76e5 100644
--- a/frontend/src/utils/sqlEditorTransaction.test.ts
+++ b/frontend/src/utils/sqlEditorTransaction.test.ts
@@ -27,6 +27,12 @@ describe('sqlEditorTransaction', () => {
     ])).toBe(false);
   });
 
+  it('uses managed transactions for data-changing CTEs even when the top-level operation is SELECT', () => {
+    const sql = 'WITH moved AS (DELETE FROM audit_logs WHERE created_at < NOW() RETURNING id) SELECT * FROM moved';
+    expect(resolveSqlEditorOperationKeyword(sql)).toBe('select');
+    expect(shouldUseSqlEditorManagedTransaction([sql])).toBe(true);
+  });
+
   it('does not wrap user-authored explicit transactions', () => {
     expect(shouldUseSqlEditorManagedTransaction([
       'BEGIN',
diff --git a/frontend/src/utils/sqlEditorTransaction.ts b/frontend/src/utils/sqlEditorTransaction.ts
index c2f5201..6cc505f 100644
--- a/frontend/src/utils/sqlEditorTransaction.ts
+++ b/frontend/src/utils/sqlEditorTransaction.ts
@@ -2,6 +2,11 @@ const SQL_EDITOR_DML_KEYWORDS = new Set(['insert', 'update', 'delete', 'replace'
 const SQL_EDITOR_READ_KEYWORDS = new Set(['select', 'with', 'show', 'describe', 'desc', 'explain', 'pragma', 'values']);
 const SQL_EDITOR_TRANSACTION_CONTROL_KEYWORDS = new Set(['begin', 'commit', 'rollback', 'savepoint', 'release']);
 
+type SqlEditorWithAnalysis = {
+    keyword: string;
+    cteHasManagedWrite: boolean;
+};
+
 const isSqlEditorKeywordChar = (char: string | undefined): boolean => !!char && /[A-Za-z0-9_]/.test(char);
 
 const skipSqlEditorTrivia = (text: string, start: number): number => {
@@ -167,8 +172,9 @@ const findTopLevelSqlEditorKeyword = (text: string, start: number, keyword: stri
     return -1;
 };
 
-const resolveSqlEditorKeywordAfterWith = (text: string, start: number): string => {
+const resolveSqlEditorWithAnalysis = (text: string, start: number): SqlEditorWithAnalysis => {
     let pos = skipSqlEditorTrivia(text, start);
+    let cteHasManagedWrite = false;
     const recursive = readSqlEditorKeyword(text, pos);
     if (recursive.keyword === 'recursive') {
         pos = recursive.end;
@@ -177,16 +183,16 @@ const resolveSqlEditorKeywordAfterWith = (text: string, start: number): string =
     while (pos < text.length) {
         pos = skipSqlEditorTrivia(text, pos);
         const identifierEnd = skipSqlEditorIdentifierToken(text, pos);
-        if (identifierEnd < 0) return '';
+        if (identifierEnd < 0) return { keyword: '', cteHasManagedWrite };
         pos = skipSqlEditorTrivia(text, identifierEnd);
         if (text[pos] === '(') {
             const columnsEnd = skipBalancedSqlEditorParens(text, pos);
-            if (columnsEnd < 0) return '';
+            if (columnsEnd < 0) return { keyword: '', cteHasManagedWrite };
             pos = skipSqlEditorTrivia(text, columnsEnd);
         }
 
         const asEnd = findTopLevelSqlEditorKeyword(text, pos, 'as');
-        if (asEnd < 0) return '';
+        if (asEnd < 0) return { keyword: '', cteHasManagedWrite };
         pos = skipSqlEditorTrivia(text, asEnd);
         const materialized = readSqlEditorKeyword(text, pos);
         if (materialized.keyword === 'not') {
@@ -199,18 +205,23 @@ const resolveSqlEditorKeywordAfterWith = (text: string, start: number): string =
         }
 
         pos = skipSqlEditorTrivia(text, pos);
-        if (text[pos] !== '(') return '';
+        if (text[pos] !== '(') return { keyword: '', cteHasManagedWrite };
+        const cteBodyStart = pos + 1;
         const cteEnd = skipBalancedSqlEditorParens(text, pos);
-        if (cteEnd < 0) return '';
+        if (cteEnd < 0) return { keyword: '', cteHasManagedWrite };
+        const cteBody = text.slice(cteBodyStart, Math.max(cteBodyStart, cteEnd - 1));
+        if (sqlEditorStatementHasManagedWrite(cteBody)) {
+            cteHasManagedWrite = true;
+        }
         pos = skipSqlEditorTrivia(text, cteEnd);
         if (text[pos] === ',') {
             pos++;
             continue;
         }
 
-        return readSqlEditorKeyword(text, pos).keyword;
+        return { keyword: readSqlEditorKeyword(text, pos).keyword, cteHasManagedWrite };
     }
-    return '';
+    return { keyword: '', cteHasManagedWrite };
 };
 
 export const resolveSqlEditorOperationKeyword = (statement: string): string => {
@@ -219,7 +230,17 @@ export const resolveSqlEditorOperationKeyword = (statement: string): string => {
     if (leading.keyword !== 'with') {
         return leading.keyword;
     }
-    return resolveSqlEditorKeywordAfterWith(text, leading.end) || leading.keyword;
+    return resolveSqlEditorWithAnalysis(text, leading.end).keyword || leading.keyword;
+};
+
+const sqlEditorStatementHasManagedWrite = (statement: string): boolean => {
+    const text = String(statement || '');
+    const leading = readSqlEditorKeyword(text, 0);
+    if (leading.keyword === 'with') {
+        const analysis = resolveSqlEditorWithAnalysis(text, leading.end);
+        return analysis.cteHasManagedWrite || SQL_EDITOR_DML_KEYWORDS.has(analysis.keyword);
+    }
+    return SQL_EDITOR_DML_KEYWORDS.has(leading.keyword);
 };
 
 const isSqlEditorTransactionControlStatement = (statement: string): boolean => {
@@ -234,12 +255,12 @@ export const shouldUseSqlEditorManagedTransaction = (statements: string[]): bool
         const trimmed = String(statement || '').trim();
         if (!trimmed) continue;
         if (isSqlEditorTransactionControlStatement(trimmed)) return false;
-        const keyword = resolveSqlEditorOperationKeyword(trimmed);
-        if (SQL_EDITOR_READ_KEYWORDS.has(keyword)) continue;
-        if (SQL_EDITOR_DML_KEYWORDS.has(keyword)) {
+        if (sqlEditorStatementHasManagedWrite(trimmed)) {
             hasManagedWrite = true;
             continue;
         }
+        const keyword = resolveSqlEditorOperationKeyword(trimmed);
+        if (SQL_EDITOR_READ_KEYWORDS.has(keyword)) continue;
         return false;
     }
     return hasManagedWrite;
diff --git a/frontend/src/v2-theme.css b/frontend/src/v2-theme.css
index ebc94b0..1401cdf 100644
--- a/frontend/src/v2-theme.css
+++ b/frontend/src/v2-theme.css
@@ -4810,7 +4810,7 @@ body[data-ui-version="v2"] .gn-v2-query-toolbar-actions {
 body[data-ui-version="v2"] .gn-v2-query-toolbar-selects {
   flex: 0 1 auto !important;
   flex-wrap: nowrap;
-  max-width: 720px;
+  max-width: 760px;
 }
 
 body[data-ui-version="v2"] .gn-v2-query-toolbar-actions {
@@ -4840,8 +4840,8 @@ body[data-ui-version="v2"] .gn-v2-query-toolbar-max-rows-select {
 }
 
 body[data-ui-version="v2"] .gn-v2-query-toolbar-transaction-mode-select {
-  width: 112px !important;
-  flex: 0 0 112px !important;
+  width: 142px !important;
+  flex: 0 0 142px !important;
 }
 
 body[data-ui-version="v2"] .gn-v2-query-toolbar-transaction-delay-select {
diff --git a/internal/app/methods_db.go b/internal/app/methods_db.go
index 033a595..294e85e 100644
--- a/internal/app/methods_db.go
+++ b/internal/app/methods_db.go
@@ -1007,6 +1007,9 @@ func (a *App) DBQueryMulti(config connection.ConnectionConfig, dbName string, qu
 
 func shouldTryQueryResultFirst(dbType string, query string) bool {
 	isSQLServer := strings.EqualFold(strings.TrimSpace(dbType), "sqlserver")
+	if keyword, withHasWrite := sqlDataOperationInfo(query); withHasWrite && keyword == "select" {
+		return true
+	}
 	keyword := leadingSQLKeyword(query)
 	switch keyword {
 	case "exec", "execute", "call":
diff --git a/internal/app/methods_db_multi_test.go b/internal/app/methods_db_multi_test.go
index 4906346..b44fa6e 100644
--- a/internal/app/methods_db_multi_test.go
+++ b/internal/app/methods_db_multi_test.go
@@ -624,6 +624,59 @@ func TestDBQueryMultiTransactionalTreatsWithDMLAsManagedWrite(t *testing.T) {
 	}
 }
 
+func TestDBQueryMultiTransactionalTreatsDataChangingCTEAsManagedWrite(t *testing.T) {
+	originalNewDatabaseFunc := newDatabaseFunc
+	t.Cleanup(func() {
+		newDatabaseFunc = originalNewDatabaseFunc
+	})
+
+	stmt := "WITH moved AS (DELETE FROM audit_logs WHERE created_at < NOW() RETURNING id) SELECT * FROM moved"
+	fakeDB := &fakeBatchWriteDB{
+		queryMap: map[string][]map[string]interface{}{
+			stmt: {{"id": 41}, {"id": 42}},
+		},
+		fieldMap: map[string][]string{
+			stmt: {"id"},
+		},
+	}
+	newDatabaseFunc = func(dbType string) (db.Database, error) {
+		return fakeDB, nil
+	}
+
+	app := NewAppWithSecretStore(secretstore.NewUnavailableStore("test"))
+	config := connection.ConnectionConfig{Type: "postgres", Host: "127.0.0.1", Port: 5432, User: "postgres"}
+
+	result := app.DBQueryMultiTransactional(config, "main", stmt, "cte-write-query")
+	if !result.Success {
+		t.Fatalf("expected transactional data-changing CTE success, got failure: %s", result.Message)
+	}
+	if result.TransactionID == "" || !result.TransactionPending {
+		t.Fatalf("expected pending transaction metadata, got id=%q pending=%v", result.TransactionID, result.TransactionPending)
+	}
+	if fakeDB.session == nil || fakeDB.session.closed {
+		t.Fatal("expected data-changing CTE transaction session to stay open")
+	}
+	wantExecs := []string{"BEGIN"}
+	if len(fakeDB.execQueries) != len(wantExecs) {
+		t.Fatalf("expected exec queries %#v, got %#v", wantExecs, fakeDB.execQueries)
+	}
+	for i, want := range wantExecs {
+		if fakeDB.execQueries[i] != want {
+			t.Fatalf("expected exec query %d = %q, got %q", i, want, fakeDB.execQueries[i])
+		}
+	}
+	if fakeDB.session.queryCalls == 0 {
+		t.Fatal("expected data-changing CTE SELECT to query returned rows inside the transaction")
+	}
+	resultSets, ok := result.Data.([]connection.ResultSetData)
+	if !ok {
+		t.Fatalf("expected []connection.ResultSetData, got %T", result.Data)
+	}
+	if len(resultSets) != 1 || len(resultSets[0].Rows) != 2 {
+		t.Fatalf("expected returned rows from data-changing CTE, got %#v", resultSets)
+	}
+}
+
 func TestDBQueryMultiTransactionalRollsBackAndClosesOnDMLFailure(t *testing.T) {
 	originalNewDatabaseFunc := newDatabaseFunc
 	t.Cleanup(func() {
diff --git a/internal/app/sql_sanitize.go b/internal/app/sql_sanitize.go
index fb102dc..d5c3467 100644
--- a/internal/app/sql_sanitize.go
+++ b/internal/app/sql_sanitize.go
@@ -53,14 +53,19 @@ func leadingSQLKeyword(query string) string {
 }
 
 func sqlDataOperationKeyword(query string) string {
+	keyword, _ := sqlDataOperationInfo(query)
+	return keyword
+}
+
+func sqlDataOperationInfo(query string) (keyword string, withHasWrite bool) {
 	keyword, keywordEnd := nextSQLKeyword(query, 0)
 	if keyword != "with" {
-		return keyword
+		return keyword, false
 	}
-	if withKeyword, ok := sqlKeywordAfterLeadingWith(query, keywordEnd); ok {
-		return withKeyword
+	if withKeyword, hasWrite, ok := sqlKeywordAfterLeadingWith(query, keywordEnd); ok {
+		return withKeyword, hasWrite
 	}
-	return keyword
+	return keyword, false
 }
 
 func nextSQLKeyword(text string, start int) (string, int) {
@@ -106,8 +111,9 @@ func skipSQLTrivia(text string, start int) int {
 	return pos
 }
 
-func sqlKeywordAfterLeadingWith(text string, start int) (string, bool) {
+func sqlKeywordAfterLeadingWith(text string, start int) (string, bool, bool) {
 	pos := skipSQLTrivia(text, start)
+	hasWriteCTE := false
 	if keyword, end := nextSQLKeyword(text, pos); keyword == "recursive" {
 		pos = end
 	}
@@ -116,20 +122,20 @@ func sqlKeywordAfterLeadingWith(text string, start int) (string, bool) {
 		pos = skipSQLTrivia(text, pos)
 		next, ok := skipSQLIdentifierToken(text, pos)
 		if !ok {
-			return "", false
+			return "", hasWriteCTE, false
 		}
 		pos = skipSQLTrivia(text, next)
 		if pos < len(text) && text[pos] == '(' {
 			next = skipBalancedSQLParens(text, pos)
 			if next < 0 {
-				return "", false
+				return "", hasWriteCTE, false
 			}
 			pos = skipSQLTrivia(text, next)
 		}
 
 		asEnd := findTopLevelSQLKeyword(text, pos, "as")
 		if asEnd < 0 {
-			return "", false
+			return "", hasWriteCTE, false
 		}
 		pos = skipSQLTrivia(text, asEnd)
 		if keyword, end := nextSQLKeyword(text, pos); keyword == "not" {
@@ -142,11 +148,19 @@ func sqlKeywordAfterLeadingWith(text string, start int) (string, bool) {
 
 		pos = skipSQLTrivia(text, pos)
 		if pos >= len(text) || text[pos] != '(' {
-			return "", false
+			return "", hasWriteCTE, false
 		}
+		cteBodyStart := pos + 1
 		next = skipBalancedSQLParens(text, pos)
 		if next < 0 {
-			return "", false
+			return "", hasWriteCTE, false
+		}
+		cteBodyEnd := next - 1
+		if cteBodyEnd >= cteBodyStart {
+			bodyKeyword, bodyHasWrite := sqlDataOperationInfo(text[cteBodyStart:cteBodyEnd])
+			if bodyHasWrite || isSQLDataWriteKeyword(bodyKeyword) {
+				hasWriteCTE = true
+			}
 		}
 		pos = skipSQLTrivia(text, next)
 		if pos < len(text) && text[pos] == ',' {
@@ -155,7 +169,7 @@ func sqlKeywordAfterLeadingWith(text string, start int) (string, bool) {
 		}
 
 		keyword, _ := nextSQLKeyword(text, pos)
-		return keyword, keyword != ""
+		return keyword, hasWriteCTE, keyword != ""
 	}
 }
 
@@ -327,7 +341,11 @@ func isReadOnlySQLQuery(dbType string, query string) bool {
 		return true
 	}
 
-	switch sqlDataOperationKeyword(query) {
+	keyword, withHasWrite := sqlDataOperationInfo(query)
+	if withHasWrite {
+		return false
+	}
+	switch keyword {
 	case "select", "with", "show", "describe", "desc", "explain", "pragma", "values":
 		return true
 	default:
@@ -340,7 +358,15 @@ func isBatchableWriteSQLStatement(dbType string, query string) bool {
 		return false
 	}
 
-	switch sqlDataOperationKeyword(query) {
+	keyword, withHasWrite := sqlDataOperationInfo(query)
+	if withHasWrite {
+		return true
+	}
+	return isSQLDataWriteKeyword(keyword)
+}
+
+func isSQLDataWriteKeyword(keyword string) bool {
+	switch keyword {
 	case "insert", "update", "delete", "replace", "merge", "upsert":
 		return true
 	default:
diff --git a/internal/app/sql_sanitize_test.go b/internal/app/sql_sanitize_test.go
index 31a77c8..3a02e5b 100644
--- a/internal/app/sql_sanitize_test.go
+++ b/internal/app/sql_sanitize_test.go
@@ -79,6 +79,11 @@ func TestIsReadOnlySQLQuery_ClassifiesWithByTopLevelOperation(t *testing.T) {
 	if isReadOnlySQLQuery("postgres", writeQuery) {
 		t.Fatal("WITH ... UPDATE should not be treated as read-only")
 	}
+
+	writeCTEQuery := "WITH moved AS (DELETE FROM audit_logs WHERE created_at < NOW() RETURNING id) SELECT * FROM moved"
+	if isReadOnlySQLQuery("postgres", writeCTEQuery) {
+		t.Fatal("data-changing CTE should not be treated as read-only")
+	}
 }
 
 func TestIsBatchableWriteSQLStatement_OnlyMatchesRealWriteStatements(t *testing.T) {
@@ -88,6 +93,9 @@ func TestIsBatchableWriteSQLStatement_OnlyMatchesRealWriteStatements(t *testing.
 	if !isBatchableWriteSQLStatement("postgres", "WITH target AS (SELECT id FROM users) DELETE FROM users WHERE id IN (SELECT id FROM target)") {
 		t.Fatal("expected WITH ... DELETE to be treated as batchable write")
 	}
+	if !isBatchableWriteSQLStatement("postgres", "WITH moved AS (DELETE FROM audit_logs WHERE created_at < NOW() RETURNING id) SELECT * FROM moved") {
+		t.Fatal("expected data-changing CTE to be treated as batchable write")
+	}
 	if isBatchableWriteSQLStatement("sqlserver", "EXEC sp_who2") {
 		t.Fatal("EXEC should not be treated as batchable write")
 	}
@@ -116,3 +124,10 @@ func TestShouldTryQueryResultFirst_TreatsSQLServerSystemCommandsAsQueryFirst(t *
 		t.Fatal("non-SQLServer system procedure name should not force query-first")
 	}
 }
+
+func TestShouldTryQueryResultFirst_TreatsDataChangingCTESelectAsQueryFirst(t *testing.T) {
+	query := "WITH moved AS (DELETE FROM audit_logs WHERE created_at < NOW() RETURNING id) SELECT * FROM moved"
+	if !shouldTryQueryResultFirst("postgres", query) {
+		t.Fatal("data-changing CTE ending in SELECT should try query-first to preserve returned rows")
+	}
+}