From 6e189723bd892a4456a34d02dd2c721e529a16e8 Mon Sep 17 00:00:00 2001
From: t8y2 <1156263951@qq.com>
Date: Sun, 23 Aug 2026 09:43:12 +0000
Subject: [PATCH] fix(sql): recognize @set declarations after comments

Closes #6974
---
 .../__tests__/useSqlExecution.spec.ts         |  2 +-
 .../lib/__tests__/sql/sqlVariables.spec.ts    | 15 ++++++++++++
 apps/desktop/src/lib/sql/sqlVariables.ts      | 23 +++++++++++++------
 3 files changed, 32 insertions(+), 8 deletions(-)

diff --git a/apps/desktop/src/composables/__tests__/useSqlExecution.spec.ts b/apps/desktop/src/composables/__tests__/useSqlExecution.spec.ts
index eb4e5c809a..3a6034aa4d 100644
--- a/apps/desktop/src/composables/__tests__/useSqlExecution.spec.ts
+++ b/apps/desktop/src/composables/__tests__/useSqlExecution.spec.ts
@@ -210,7 +210,7 @@ describe("useSqlExecution", () => {
 
   it("expands preceding @set values in a selected shell-style statement", async () => {
     const selectedSql = "SELECT * FROM patrol WHERE post_id = ${postid};";
-    const fullSql = ["@set postid = '224';", selectedSql, "@set postid = 'future';"].join("\n");
+    const fullSql = ["-- saved query", "@set postid='224';", selectedSql, "@set postid = 'future';"].join("\n");
     const selectionFrom = fullSql.indexOf("SELECT");
     const selectionTo = selectionFrom + selectedSql.length;
     const activeTab = ref<QueryTab | undefined>({ ...queryTab("app"), sql: fullSql });
diff --git a/apps/desktop/src/lib/__tests__/sql/sqlVariables.spec.ts b/apps/desktop/src/lib/__tests__/sql/sqlVariables.spec.ts
index e1cbf3c399..468cf7b900 100644
--- a/apps/desktop/src/lib/__tests__/sql/sqlVariables.spec.ts
+++ b/apps/desktop/src/lib/__tests__/sql/sqlVariables.spec.ts
@@ -24,6 +24,21 @@ describe("expandSqlVariables", () => {
     expect(expandSqlVariables(sql).sql).toBe("select * from t where post_id = '224'");
   });
 
+  it("recognizes a declaration after a leading line comment", () => {
+    const sql = ["-- saved query", "@set devid=7062;", "select * from t where device_id = ${devid}"].join("\n");
+    expect(expandSqlVariables(sql).sql).toBe("-- saved query\nselect * from t where device_id = 7062");
+  });
+
+  it("recognizes declarations after comments at a statement boundary", () => {
+    const sql = ["select 1;", "/* next statement */", "@set devid = 7062;", "select ${devid}"].join("\n");
+    expect(expandSqlVariables(sql).sql).toBe("select 1;\n/* next statement */\nselect 7062");
+  });
+
+  it("does not recognize a declaration after an unterminated statement", () => {
+    const sql = ["select 1", "-- same statement", "@set devid = 7062;", "select ${devid}"].join("\n");
+    expect(expandSqlVariables(sql)).toEqual({ sql, expanded: false });
+  });
+
   it("uses declarations from the surrounding document for a selected statement", () => {
     const selectedSql = "select * from t where post_id = ${postid}";
     const declarationSql = ["@set postid = '224';", selectedSql].join("\n");
diff --git a/apps/desktop/src/lib/sql/sqlVariables.ts b/apps/desktop/src/lib/sql/sqlVariables.ts
index 2fb0d30a0b..0f6d446ba1 100644
--- a/apps/desktop/src/lib/sql/sqlVariables.ts
+++ b/apps/desktop/src/lib/sql/sqlVariables.ts
@@ -67,6 +67,7 @@ function collectDeclarations(sql: string, databaseType?: DatabaseType): Declarat
   const declarations: DeclarationSpan[] = [];
   let i = 0;
   let dollarQuoteEnd = "";
+  let statementStart = true;
 
   while (i < sql.length) {
     if (dollarQuoteEnd) {
@@ -80,11 +81,17 @@ function collectDeclarations(sql: string, databaseType?: DatabaseType): Declarat
     const ch = sql[i];
     const next = sql[i + 1];
 
+    if (/\s/.test(ch)) {
+      i += 1;
+      continue;
+    }
     if (ch === "'" || ch === '"' || ch === "`") {
+      statementStart = false;
       i = skipQuoted(sql, i, ch);
       continue;
     }
     if (ch === "[") {
+      statementStart = false;
       i = skipBracketIdentifier(sql, i, databaseType);
       continue;
     }
@@ -99,19 +106,27 @@ function collectDeclarations(sql: string, databaseType?: DatabaseType): Declarat
     if (ch === "$") {
       const marker = readDollarQuoteMarker(sql, i);
       if (marker) {
+        statementStart = false;
         dollarQuoteEnd = marker;
         i += marker.length;
         continue;
       }
     }
-    if (ch === "@" && matchesWord(sql, i + 1, "set") && isStatementStart(sql, i)) {
+    if (ch === "@" && statementStart && matchesWord(sql, i + 1, "set")) {
       const declaration = readDeclaration(sql, i, databaseType);
       if (declaration) {
         declarations.push(declaration);
         i = declaration.end;
+        statementStart = true;
         continue;
       }
     }
+    if (ch === ";") {
+      statementStart = true;
+      i += 1;
+      continue;
+    }
+    statementStart = false;
     i += 1;
   }
 
@@ -283,12 +298,6 @@ function stripDeclaration(sql: string, start: number, end: number): string {
   return sql.slice(0, trimmableStart) + sql.slice(lineEnd);
 }
 
-function isStatementStart(sql: string, start: number): boolean {
-  let i = start - 1;
-  while (i >= 0 && /\s/.test(sql[i])) i -= 1;
-  return i < 0 || sql[i] === ";";
-}
-
 function matchesWord(sql: string, start: number, word: string): boolean {
   const value = sql.slice(start, start + word.length);
   if (value.toLowerCase() !== word) return false;