From 812424765702f2eee96c6b0ae58b63d976921751 Mon Sep 17 00:00:00 2001
From: ordinary being <majiantong666@qq.com>
Date: Tue, 6 Oct 2026 18:58:15 +0800
Subject: [PATCH] fix(safety): scan production SQL with the connection's lexer
 rules

---
 .../src/lib/database/productionSafety.ts      | 31 +++++++++-----
 apps/desktop/src/lib/sql/sqlRisk.ts           | 17 ++++----
 tests/fixtures/production-safety-corpus.json  | 40 +++++++++++++++++++
 3 files changed, 71 insertions(+), 17 deletions(-)

diff --git a/apps/desktop/src/lib/database/productionSafety.ts b/apps/desktop/src/lib/database/productionSafety.ts
index af6f18bad0..b3dd3e35f0 100644
--- a/apps/desktop/src/lib/database/productionSafety.ts
+++ b/apps/desktop/src/lib/database/productionSafety.ts
@@ -1,6 +1,6 @@
 import type { ConnectionConfig, DatabaseType } from "@/types/database";
 import { nacosNamespaceIdentity } from "@/lib/nacos/nacosNamespaceVisibility";
-import { classifySqlRisk, isSqlRiskMutation } from "@/lib/sql/sqlRisk";
+import { classifySqlRisk, isSqlRiskMutation, usesMysqlLexerRules } from "@/lib/sql/sqlRisk";
 
 export type ProductionContextReason = "connection" | "database" | "sql_target";
 
@@ -124,7 +124,7 @@ export function isProductionMutation(sql: string): boolean {
  */
 export function assessProductionSql(sql: string, connection: ConnectionConfig | undefined, activeDatabase: string | undefined | null): ProductionSqlAssessment {
   const activeContext = productionContextForDatabase(connection, activeDatabase);
-  const targetText = sqlTargetSafetyText(sql);
+  const targetText = sqlTargetSafetyText(sql, connection?.db_type);
   const statements = splitTargetStatements(targetText.text);
   const risk = classifySqlRisk(sql, { dialect: connection?.db_type });
   const isMutation = isSqlRiskMutation(risk.risk);
@@ -258,7 +258,20 @@ function splitTargetStatements(sql: string): string[] {
     .filter(Boolean);
 }
 
-function sqlTargetSafetyText(sql: string, quotedIdentifiers = new Map<string, string>()): SqlTargetSafetyText {
+/**
+ * Removes the text no database target can be found in. Everything removed here
+ * is text this gate cannot inspect, so reading MySQL lexer rules for a
+ * connection that has none hides writes: on SQL Server `#tmp` is a temporary
+ * table and T-SQL escapes a quote by doubling it, so both
+ * `SELECT * FROM #tmp; DELETE FROM prod_app.dbo.users;` and
+ * `SELECT 'dir\'; DELETE FROM prod_app.dbo.users;` would lose the `DELETE` and
+ * the write would no longer require confirmation.
+ *
+ * Mirrors `SqlScanLexerRules` in
+ * `crates/dbx-core/src/safety/production_safety.rs`.
+ */
+function sqlTargetSafetyText(sql: string, dialect?: DatabaseType | string, quotedIdentifiers = new Map<string, string>()): SqlTargetSafetyText {
+  const mysqlLexer = usesMysqlLexerRules(dialect);
   let output = "";
   let index = 0;
 
@@ -273,7 +286,7 @@ function sqlTargetSafetyText(sql: string, quotedIdentifiers = new Map<string, st
       continue;
     }
 
-    if (char === "#") {
+    if (char === "#" && mysqlLexer) {
       index += 1;
       while (index < sql.length && sql[index] !== "\n" && sql[index] !== "\r") index += 1;
       output += " ";
@@ -286,7 +299,7 @@ function sqlTargetSafetyText(sql: string, quotedIdentifiers = new Map<string, st
       const executablePrefixLength = mysqlExecutableCommentPrefixLength(sql, index);
       if (executablePrefixLength > 0) {
         const bodyStart = skipExecutableCommentVersion(sql, index + executablePrefixLength);
-        output += ` ${sqlTargetSafetyText(sql.slice(bodyStart, close), quotedIdentifiers).text} `;
+        output += ` ${sqlTargetSafetyText(sql.slice(bodyStart, close), dialect, quotedIdentifiers).text} `;
       } else {
         output += " ";
       }
@@ -303,14 +316,14 @@ function sqlTargetSafetyText(sql: string, quotedIdentifiers = new Map<string, st
     }
 
     if (char === "'") {
-      index = readQuotedEnd(sql, index, "'", "'");
+      index = readQuotedEnd(sql, index, "'", "'", mysqlLexer);
       output += " ";
       continue;
     }
 
     if (char === '"' || char === "`" || char === "[") {
       const close = char === "[" ? "]" : char;
-      const end = readQuotedEnd(sql, index, char, close);
+      const end = readQuotedEnd(sql, index, char, close, mysqlLexer);
       const identifier = unquoteIdentifier(sql.slice(index, end), char, close).replace(/[;]/g, " ");
       const token = `__dbxq${quotedIdentifiers.size}__`;
       quotedIdentifiers.set(token.toLowerCase(), identifier);
@@ -343,10 +356,10 @@ function dollarQuoteTagAt(sql: string, index: number): string | undefined {
   return sql.slice(index).match(/^\$[A-Za-z_][A-Za-z0-9_]*\$|^\$\$/)?.[0];
 }
 
-function readQuotedEnd(sql: string, start: number, open: string, close: string): number {
+function readQuotedEnd(sql: string, start: number, open: string, close: string, backslashEscapes: boolean): number {
   let index = start + open.length;
   while (index < sql.length) {
-    if (sql[index] === "\\" && (open === "'" || open === '"')) {
+    if (backslashEscapes && sql[index] === "\\" && (open === "'" || open === '"')) {
       index += 2;
       continue;
     }
diff --git a/apps/desktop/src/lib/sql/sqlRisk.ts b/apps/desktop/src/lib/sql/sqlRisk.ts
index 82c6a0ea74..82e298fdfb 100644
--- a/apps/desktop/src/lib/sql/sqlRisk.ts
+++ b/apps/desktop/src/lib/sql/sqlRisk.ts
@@ -37,19 +37,20 @@ const SAFE_READ_PRAGMA_NAMES = new Set(["table_info", "table_xinfo", "index_list
  * string escape. Both rules are MySQL family features: on SQL Server `#tmp` is a
  * temporary table, and on PostgreSQL `standard_conforming_strings` is on by
  * default, so `'dir\'` is a complete string. Assuming the MySQL rules everywhere
- * hides SQL from the classifier — `SELECT * FROM #tmp; DELETE FROM #tmp;` (SQL
- * Server) and `SELECT 'dir\'; DELETE FROM users;` (PostgreSQL) would otherwise be
- * read as one read-only statement and auto-execute unconfirmed.
+ * hides SQL from the safety scans — `SELECT * FROM #tmp; DELETE FROM prod.users;`
+ * (SQL Server) and `SELECT 'dir\'; DELETE FROM users;` (PostgreSQL) would
+ * otherwise be read as one read-only statement and auto-execute.
  *
  * Mirrors `is_mysql_compatible_database` in `crates/dbx-sql-core/src/sql.rs`,
- * which gates `supports_hash_line_comments` the same way. An unknown dialect is
- * deliberately not treated as MySQL: keeping `#` and backslashes literal can only
- * split more statements than before, which raises the assessed risk rather than
- * lowering it.
+ * which gates `supports_hash_line_comments` the same way, and
+ * `SqlScanLexerRules` in `crates/dbx-core/src/safety/production_safety.rs`.
+ * An unknown dialect is deliberately not treated as MySQL: keeping `#` and
+ * backslashes literal can only split more statements than before, which raises
+ * the assessed risk rather than lowering it.
  */
 const MYSQL_LEXER_DATABASE_TYPES = new Set<string>(["mysql", "doris", "starrocks", "manticoresearch", "goldendb"]);
 
-function usesMysqlLexerRules(dialect?: DatabaseType | string): boolean {
+export function usesMysqlLexerRules(dialect?: DatabaseType | string): boolean {
   return typeof dialect === "string" && MYSQL_LEXER_DATABASE_TYPES.has(dialect.trim().toLowerCase());
 }
 
diff --git a/tests/fixtures/production-safety-corpus.json b/tests/fixtures/production-safety-corpus.json
index 3f1cc4c387..20d6829c18 100644
--- a/tests/fixtures/production-safety-corpus.json
+++ b/tests/fixtures/production-safety-corpus.json
@@ -308,5 +308,45 @@
     "active": false,
     "isMutation": true,
     "databases": []
+  },
+  {
+    "name": "sqlserver hash temp table keeps a later production delete in scope",
+    "dialect": "sqlserver",
+    "productionDatabases": ["prod_app"],
+    "activeDatabase": "staging",
+    "sql": "SELECT * FROM #tmp; DELETE FROM prod_app.dbo.users;",
+    "active": true,
+    "isMutation": true,
+    "databases": ["prod_app"]
+  },
+  {
+    "name": "sqlserver backslash string keeps a later production delete in scope",
+    "dialect": "sqlserver",
+    "productionDatabases": ["prod_app"],
+    "activeDatabase": "staging",
+    "sql": "SELECT 'dir\\'; DELETE FROM prod_app.dbo.users;",
+    "active": true,
+    "isMutation": true,
+    "databases": ["prod_app"]
+  },
+  {
+    "name": "mysql hash comment keeps a commented delete out of scope",
+    "dialect": "mysql",
+    "productionDatabases": ["prod_app"],
+    "activeDatabase": "staging",
+    "sql": "SELECT 1; # DELETE FROM prod_app.users",
+    "active": false,
+    "isMutation": false,
+    "databases": []
+  },
+  {
+    "name": "mysql backslash escape keeps the write inside the literal",
+    "dialect": "mysql",
+    "productionDatabases": ["prod_app"],
+    "activeDatabase": "staging",
+    "sql": "SELECT 'it\\'s; DELETE FROM prod_app.users;'",
+    "active": false,
+    "isMutation": false,
+    "databases": []
   }
 ]