From 135d4e15db2eb31c92560b7635614232bc486a3a Mon Sep 17 00:00:00 2001
From: miracle <cmiracle@163.com>
Date: Thu, 9 Jul 2026 00:34:51 +0800
Subject: [PATCH] feat(sql): add support for @set variables in scripts and copy
 button

---
 .../components/editor/SqlParameterDialog.vue  |  18 +-
 .../src/composables/useSqlExecution.ts        |   9 +-
 .../lib/__tests__/sql/sqlVariables.spec.ts    |  61 ++++
 apps/desktop/src/lib/sql/sqlVariables.ts      | 331 ++++++++++++++++++
 4 files changed, 414 insertions(+), 5 deletions(-)
 create mode 100644 apps/desktop/src/lib/__tests__/sql/sqlVariables.spec.ts
 create mode 100644 apps/desktop/src/lib/sql/sqlVariables.ts

diff --git a/apps/desktop/src/components/editor/SqlParameterDialog.vue b/apps/desktop/src/components/editor/SqlParameterDialog.vue
index c07ba33693..e272800558 100644
--- a/apps/desktop/src/components/editor/SqlParameterDialog.vue
+++ b/apps/desktop/src/components/editor/SqlParameterDialog.vue
@@ -1,6 +1,6 @@
 <script setup lang="ts">
 import { computed, nextTick, ref, watch } from "vue";
-import { Braces } from "@lucide/vue";
+import { Braces, Copy } from "@lucide/vue";
 import { useI18n } from "vue-i18n";
 import { Button } from "@/components/ui/button";
 import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog";
@@ -11,10 +11,13 @@ import TruncatedTextTooltip from "@/components/ui/TruncatedTextTooltip.vue";
 import { loadSqlParameterHistory, rememberSqlParameterValues } from "@/lib/sql/sqlParameterHistory";
 import { substituteSqlParameters, type SqlParameterDescriptor, type SqlParameterInput, type SqlParameterSyntax, type SqlParameterValueKind } from "@/lib/sql/sqlParameters";
 import { useSqlHighlighter } from "@/composables/useSqlHighlighter";
+import { useToast } from "@/composables/useToast";
+import { copyToClipboard } from "@/lib/common/clipboard";
 import type { DatabaseType } from "@/types/database";
 
 const { t } = useI18n();
 const { highlight } = useSqlHighlighter();
+const { toast } = useToast();
 
 const open = defineModel<boolean>("open", { default: false });
 
@@ -105,6 +108,15 @@ function execute() {
   open.value = false;
   emit("execute", resolvedSql.value);
 }
+
+async function copyResolvedSql() {
+  try {
+    await copyToClipboard(resolvedSql.value);
+    toast(t("grid.copied"));
+  } catch (e: any) {
+    toast(t("grid.copyFailed", { message: e?.message || String(e) }), 5000);
+  }
+}
 </script>
 
 <template>
@@ -183,6 +195,10 @@ function execute() {
 
       <DialogFooter>
         <Button variant="outline" @click="open = false">{{ t("dangerDialog.cancel") }}</Button>
+        <Button variant="outline" @click="copyResolvedSql">
+          <Copy class="mr-1.5 h-4 w-4" />
+          {{ t("grid.copy") }}
+        </Button>
         <Button @click="execute">{{ t("sqlParameters.execute") }}</Button>
       </DialogFooter>
     </DialogContent>
diff --git a/apps/desktop/src/composables/useSqlExecution.ts b/apps/desktop/src/composables/useSqlExecution.ts
index a9dbdb43b1..c17bb6ea96 100644
--- a/apps/desktop/src/composables/useSqlExecution.ts
+++ b/apps/desktop/src/composables/useSqlExecution.ts
@@ -12,6 +12,7 @@ import { sqlMetadataRefreshTarget } from "@/lib/sql/sqlMetadataRefresh";
 import { classifyRedisCommandSafety, firstRedisCommandToken } from "@/lib/redis/redisCommandSafety";
 import { isSqlExecutionSnapshot, resolveExecutableSql, type SqlExecutionOverride, type SqlExecutionSnapshot } from "@/lib/sql/sqlExecutionTarget";
 import { extractSqlParameterDescriptors, type SqlParameterDescriptor } from "@/lib/sql/sqlParameters";
+import { expandSqlVariables } from "@/lib/sql/sqlVariables";
 import type { ConnectionConfig, DatabaseType, QueryTab } from "@/types/database";
 
 const DANGER_RE = /^\s*(DROP|DELETE|TRUNCATE|ALTER|UPDATE|MERGE|REPLACE)\b/i;
@@ -64,10 +65,10 @@ export function useSqlExecution(deps: {
   const sqlParameterDatabaseType = ref<DatabaseType | undefined>();
 
   async function resolvedExecutableSql(source?: SqlExecutionOverride): Promise<string> {
-    if (typeof source === "string") return source;
-    if (deps.resolveExecutableSql) return await deps.resolveExecutableSql(source);
-    if (isSqlExecutionSnapshot(source)) return resolveExecutableSql(source.fullSql, source.selectedSql, { cursorPos: source.cursorPos });
-    return deps.executableSql.value;
+    if (typeof source === "string") return expandSqlVariables(source).sql;
+    if (deps.resolveExecutableSql) return expandSqlVariables(await deps.resolveExecutableSql(source)).sql;
+    if (isSqlExecutionSnapshot(source)) return expandSqlVariables(resolveExecutableSql(source.fullSql, source.selectedSql, { cursorPos: source.cursorPos })).sql;
+    return expandSqlVariables(deps.executableSql.value).sql;
   }
 
   async function tryExecute(sqlOverride?: SqlExecutionOverride) {
diff --git a/apps/desktop/src/lib/__tests__/sql/sqlVariables.spec.ts b/apps/desktop/src/lib/__tests__/sql/sqlVariables.spec.ts
new file mode 100644
index 0000000000..b97ca1b518
--- /dev/null
+++ b/apps/desktop/src/lib/__tests__/sql/sqlVariables.spec.ts
@@ -0,0 +1,61 @@
+import { describe, expect, it } from "vitest";
+import { expandSqlVariables } from "@/lib/sql/sqlVariables";
+
+describe("expandSqlVariables", () => {
+  it("returns the SQL unchanged when there is no @set declaration", () => {
+    const sql = "select * from t where id = @id";
+    expect(expandSqlVariables(sql)).toEqual({ sql, expanded: false });
+  });
+
+  it("inlines a declared IN-list verbatim across the script", () => {
+    const sql = ["@set client_id = (606,322,634);", "select * from invoices where client_id in @client_id"].join("\n");
+    const { sql: result, expanded } = expandSqlVariables(sql);
+    expect(expanded).toBe(true);
+    expect(result).toBe("select * from invoices where client_id in (606,322,634)");
+  });
+
+  it("inlines a quoted string value verbatim", () => {
+    const sql = ["@set date_start = '2026-07-04 00:00:00';", "select * from t where created_at < @date_start"].join("\n");
+    expect(expandSqlVariables(sql).sql).toBe("select * from t where created_at < '2026-07-04 00:00:00'");
+  });
+
+  it("expands the same variable in multiple places", () => {
+    const sql = ["@set tenant = 42;", "select @tenant, count(*) from t where tenant_id = @tenant"].join("\n");
+    expect(expandSqlVariables(sql).sql).toBe("select 42, count(*) from t where tenant_id = 42");
+  });
+
+  it("supports several declarations", () => {
+    const sql = ["@set a = 1;", "@set b = 'x';", "select @a, @b"].join("\n");
+    expect(expandSqlVariables(sql).sql).toBe("select 1, 'x'");
+  });
+
+  it("leaves undeclared @name references untouched", () => {
+    const sql = ["@set a = 1;", "select @a, @b, @@version"].join("\n");
+    expect(expandSqlVariables(sql).sql).toBe("select 1, @b, @@version");
+  });
+
+  it("does not expand references inside strings, comments, or quoted identifiers", () => {
+    const sql = ["@set a = 1;", "select '@a' as s, \"@a\" as q, `@a` as b -- @a"].join("\n");
+    expect(expandSqlVariables(sql).sql).toBe("select '@a' as s, \"@a\" as q, `@a` as b -- @a");
+  });
+
+  it("does not treat @set inside a string as a declaration", () => {
+    const sql = "select '@set a = 1;' as note, @a";
+    expect(expandSqlVariables(sql)).toEqual({ sql, expanded: false });
+  });
+
+  it("keeps a value's own quotes and parentheses intact", () => {
+    const sql = ["@set filter = (status = 'drafted' and deleted_at is null);", "select * from t where @filter"].join("\n");
+    expect(expandSqlVariables(sql).sql).toBe("select * from t where (status = 'drafted' and deleted_at is null)");
+  });
+
+  it("matches @set case-insensitively", () => {
+    const sql = ["@SET a = 7;", "select @a"].join("\n");
+    expect(expandSqlVariables(sql).sql).toBe("select 7");
+  });
+
+  it("does not treat @settings as an @set declaration", () => {
+    const sql = "select @settings from t";
+    expect(expandSqlVariables(sql)).toEqual({ sql, expanded: false });
+  });
+});
diff --git a/apps/desktop/src/lib/sql/sqlVariables.ts b/apps/desktop/src/lib/sql/sqlVariables.ts
new file mode 100644
index 0000000000..b193158ad8
--- /dev/null
+++ b/apps/desktop/src/lib/sql/sqlVariables.ts
@@ -0,0 +1,331 @@
+// Client-side SQL variable expansion.
+//
+// Lets users declare reusable values inside a SQL script with `@set name = value;`
+// and reference them elsewhere as `@name`. Expansion happens entirely on the client:
+// declarations are stripped and every reference is replaced with the declared value
+// verbatim (raw). Because the SQL sent to the server contains only plain literals,
+// this works uniformly across PostgreSQL, MySQL, SQL Server and every other backend
+// regardless of whether they have native variable support.
+//
+// This is intentionally separate from the placeholder parameter system
+// (`sqlParameters.ts`). A `@name` reference is only expanded when a matching
+// `@set` declaration exists in the same script; any other `@name` (SQL Server
+// native variables, `@@version`, dialog parameters) is left untouched.
+
+const VARIABLE_NAME_START_RE = /[\p{L}_]/u;
+const VARIABLE_NAME_CHAR_RE = /[\p{L}\p{N}_]/u;
+
+export interface SqlVariableExpansion {
+  sql: string;
+  expanded: boolean;
+}
+
+interface DeclarationSpan {
+  name: string;
+  value: string;
+  start: number;
+  end: number;
+}
+
+/**
+ * Expand `@set name = value;` declarations within a SQL script.
+ *
+ * Returns the SQL with declarations removed and `@name` references replaced by
+ * their declared values. When no declaration is found the input is returned
+ * unchanged with `expanded: false`.
+ */
+export function expandSqlVariables(sql: string): SqlVariableExpansion {
+  const declarations = collectDeclarations(sql);
+  if (!declarations.length) return { sql, expanded: false };
+
+  const values = new Map<string, string>();
+  for (const declaration of declarations) {
+    values.set(declaration.name.toLowerCase(), declaration.value);
+  }
+
+  // Remove declaration spans first (right-to-left keeps earlier offsets valid),
+  // then replace references in the remaining text.
+  let result = sql;
+  for (let i = declarations.length - 1; i >= 0; i -= 1) {
+    const declaration = declarations[i];
+    result = stripDeclaration(result, declaration.start, declaration.end);
+  }
+
+  result = replaceReferences(result, values);
+  return { sql: result, expanded: true };
+}
+
+function collectDeclarations(sql: string): DeclarationSpan[] {
+  const declarations: DeclarationSpan[] = [];
+  let i = 0;
+  let dollarQuoteEnd = "";
+
+  while (i < sql.length) {
+    if (dollarQuoteEnd) {
+      const end = sql.indexOf(dollarQuoteEnd, i);
+      if (end === -1) break;
+      i = end + dollarQuoteEnd.length;
+      dollarQuoteEnd = "";
+      continue;
+    }
+
+    const ch = sql[i];
+    const next = sql[i + 1];
+
+    if (ch === "'" || ch === '"' || ch === "`") {
+      i = skipQuoted(sql, i, ch);
+      continue;
+    }
+    if (ch === "[") {
+      i = skipBracketIdentifier(sql, i);
+      continue;
+    }
+    if (ch === "-" && next === "-") {
+      i = skipLine(sql, i + 2);
+      continue;
+    }
+    if (ch === "/" && next === "*") {
+      i = skipBlockComment(sql, i + 2);
+      continue;
+    }
+    if (ch === "$") {
+      const marker = readDollarQuoteMarker(sql, i);
+      if (marker) {
+        dollarQuoteEnd = marker;
+        i += marker.length;
+        continue;
+      }
+    }
+    if (ch === "@" && matchesWord(sql, i + 1, "set") && isStatementStart(sql, i)) {
+      const declaration = readDeclaration(sql, i);
+      if (declaration) {
+        declarations.push(declaration);
+        i = declaration.end;
+        continue;
+      }
+    }
+    i += 1;
+  }
+
+  return declarations;
+}
+
+// Parse `@set name = value` starting at `start` (the `@`). The value runs until
+// the terminating `;` or end of input, honouring nested quotes, comments and
+// parentheses so that `IN (...)` lists and quoted strings survive intact.
+function readDeclaration(sql: string, start: number): DeclarationSpan | null {
+  let i = start + 1 + "set".length;
+  i = skipInlineWhitespace(sql, i);
+
+  const name = readVariableName(sql, i);
+  if (!name) return null;
+  i += name.length;
+
+  i = skipInlineWhitespace(sql, i);
+  if (sql[i] !== "=") return null;
+  i += 1;
+  i = skipInlineWhitespace(sql, i);
+
+  const valueStart = i;
+  const valueEnd = readValueEnd(sql, i);
+  const value = sql.slice(valueStart, valueEnd).trim();
+  if (!value) return null;
+
+  // Consume the terminating semicolon so the declaration disappears cleanly.
+  let end = valueEnd;
+  if (sql[end] === ";") end += 1;
+
+  return { name, value, start, end };
+}
+
+function readValueEnd(sql: string, start: number): number {
+  let i = start;
+  let depth = 0;
+  while (i < sql.length) {
+    const ch = sql[i];
+    const next = sql[i + 1];
+    if (ch === "'" || ch === '"' || ch === "`") {
+      i = skipQuoted(sql, i, ch);
+      continue;
+    }
+    if (ch === "[") {
+      i = skipBracketIdentifier(sql, i);
+      continue;
+    }
+    if (ch === "-" && next === "-") return i;
+    if (ch === "/" && next === "*") return i;
+    if (ch === "$") {
+      const marker = readDollarQuoteMarker(sql, i);
+      if (marker) {
+        const end = sql.indexOf(marker, i + marker.length);
+        i = end === -1 ? sql.length : end + marker.length;
+        continue;
+      }
+    }
+    if (ch === "(") depth += 1;
+    else if (ch === ")") depth = Math.max(0, depth - 1);
+    else if ((ch === ";" || ch === "\n") && depth === 0) return i;
+    i += 1;
+  }
+  return sql.length;
+}
+
+function replaceReferences(sql: string, values: Map<string, string>): string {
+  let result = "";
+  let i = 0;
+  let dollarQuoteEnd = "";
+
+  while (i < sql.length) {
+    if (dollarQuoteEnd) {
+      const end = sql.indexOf(dollarQuoteEnd, i);
+      const stop = end === -1 ? sql.length : end + dollarQuoteEnd.length;
+      result += sql.slice(i, stop);
+      i = stop;
+      dollarQuoteEnd = "";
+      continue;
+    }
+
+    const ch = sql[i];
+    const next = sql[i + 1];
+
+    if (ch === "'" || ch === '"' || ch === "`") {
+      const end = skipQuoted(sql, i, ch);
+      result += sql.slice(i, end);
+      i = end;
+      continue;
+    }
+    if (ch === "[") {
+      const end = skipBracketIdentifier(sql, i);
+      result += sql.slice(i, end);
+      i = end;
+      continue;
+    }
+    if (ch === "-" && next === "-") {
+      const end = skipLine(sql, i + 2);
+      result += sql.slice(i, end);
+      i = end;
+      continue;
+    }
+    if (ch === "/" && next === "*") {
+      const end = skipBlockComment(sql, i + 2);
+      result += sql.slice(i, end);
+      i = end;
+      continue;
+    }
+    if (ch === "$") {
+      const marker = readDollarQuoteMarker(sql, i);
+      if (marker) {
+        result += marker;
+        i += marker.length;
+        dollarQuoteEnd = marker;
+        continue;
+      }
+    }
+    if (ch === "@" && next !== "@" && sql[i - 1] !== "@") {
+      const name = readVariableName(sql, i + 1);
+      if (name) {
+        const value = values.get(name.toLowerCase());
+        if (value !== undefined) {
+          result += value;
+          i += 1 + name.length;
+          continue;
+        }
+      }
+    }
+    result += ch;
+    i += 1;
+  }
+
+  return result;
+}
+
+// Remove a declaration span, collapsing the leftover blank line it leaves behind
+// so the expanded SQL stays tidy.
+function stripDeclaration(sql: string, start: number, end: number): string {
+  let lineStart = start;
+  while (lineStart > 0 && sql[lineStart - 1] !== "\n") {
+    if (sql[lineStart - 1] !== " " && sql[lineStart - 1] !== "\t" && sql[lineStart - 1] !== "\r") break;
+    lineStart -= 1;
+  }
+  const trimmableStart = lineStart === 0 || sql[lineStart - 1] === "\n" ? lineStart : start;
+
+  let lineEnd = end;
+  while (lineEnd < sql.length && (sql[lineEnd] === " " || sql[lineEnd] === "\t" || sql[lineEnd] === "\r")) lineEnd += 1;
+  if (trimmableStart === lineStart && sql[lineEnd] === "\n") lineEnd += 1;
+
+  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;
+  return !VARIABLE_NAME_CHAR_RE.test(sql[start + word.length] ?? "");
+}
+
+function readVariableName(sql: string, start: number): string {
+  if (!VARIABLE_NAME_START_RE.test(sql[start] ?? "")) return "";
+  let i = start + 1;
+  while (i < sql.length && VARIABLE_NAME_CHAR_RE.test(sql[i])) i += 1;
+  return sql.slice(start, i);
+}
+
+function skipInlineWhitespace(sql: string, start: number): number {
+  let i = start;
+  while (i < sql.length && (sql[i] === " " || sql[i] === "\t" || sql[i] === "\r")) i += 1;
+  return i;
+}
+
+function skipQuoted(sql: string, start: number, quote: string): number {
+  let i = start + 1;
+  while (i < sql.length) {
+    if (sql[i] === "\\" && quote === "'" && i + 1 < sql.length) {
+      i += 2;
+      continue;
+    }
+    if (sql[i] === quote) {
+      if (sql[i + 1] === quote) {
+        i += 2;
+        continue;
+      }
+      return i + 1;
+    }
+    i += 1;
+  }
+  return sql.length;
+}
+
+function skipBracketIdentifier(sql: string, start: number): number {
+  let i = start + 1;
+  while (i < sql.length) {
+    if (sql[i] === "]") {
+      if (sql[i + 1] === "]") {
+        i += 2;
+        continue;
+      }
+      return i + 1;
+    }
+    i += 1;
+  }
+  return sql.length;
+}
+
+function skipLine(sql: string, start: number): number {
+  const nextNewline = sql.indexOf("\n", start);
+  return nextNewline === -1 ? sql.length : nextNewline + 1;
+}
+
+function skipBlockComment(sql: string, start: number): number {
+  const end = sql.indexOf("*/", start);
+  return end === -1 ? sql.length : end + 2;
+}
+
+function readDollarQuoteMarker(sql: string, start: number): string {
+  const match = sql.slice(start).match(/^\$[A-Za-z_][A-Za-z0-9_]*\$|^\$\$/);
+  return match?.[0] ?? "";
+}