From bbb57b32c8800655e25c2b6f153038f2659232c4 Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Wed, 19 Aug 2026 16:43:20 +0000 Subject: [PATCH] fix(sql): expand @set variables in selected statements Closes #6719 --- .../__tests__/useSqlExecution.spec.ts | 28 +++++++++++ .../src/composables/useSqlExecution.ts | 5 +- .../lib/__tests__/sql/sqlVariables.spec.ts | 16 +++++++ apps/desktop/src/lib/sql/sqlVariables.ts | 47 +++++++++++++------ 4 files changed, 79 insertions(+), 17 deletions(-) diff --git a/apps/desktop/src/composables/__tests__/useSqlExecution.spec.ts b/apps/desktop/src/composables/__tests__/useSqlExecution.spec.ts index 62d4c0a270..c4f2cb6b20 100644 --- a/apps/desktop/src/composables/__tests__/useSqlExecution.spec.ts +++ b/apps/desktop/src/composables/__tests__/useSqlExecution.spec.ts @@ -208,6 +208,34 @@ describe("useSqlExecution", () => { expect(executeCurrentSql).toHaveBeenCalledWith(selectedSql, { sourceOffset: selectionFrom }); }); + 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 selectionFrom = fullSql.indexOf("SELECT"); + const selectionTo = selectionFrom + selectedSql.length; + const activeTab = ref({ ...queryTab("app"), sql: fullSql }); + const activeConnection = ref(connection("postgres")); + const activeOutputView = ref<"result" | "summary" | "explain" | "chart">("result"); + const queryStore = useQueryStore(); + const executeCurrentSql = vi.spyOn(queryStore, "executeCurrentSql").mockImplementation(async () => { + if (activeTab.value) activeTab.value.result = { columns: ["post_id"], rows: [["224"]], affected_rows: 0, execution_time_ms: 1 }; + }); + vi.spyOn(useHistoryStore(), "add").mockResolvedValue(undefined); + + const execution = useSqlExecution({ + activeTab: computed(() => activeTab.value), + activeConnection: computed(() => activeConnection.value), + executableSql: computed(() => selectedSql), + resolveExecutableSql: async () => selectedSql, + activeOutputView, + }); + + await execution.tryExecute({ fullSql, selectedSql, cursorPos: selectionFrom, selectionFrom, selectionTo }); + + expect(execution.showSqlParameterDialog.value).toBe(false); + expect(executeCurrentSql).toHaveBeenCalledWith("SELECT * FROM patrol WHERE post_id = '224';", {}); + }); + it("opens the execution summary for a multi-statement batch", async () => { const sql = "SELECT 1;\nSELECT 2;"; const activeTab = ref({ ...queryTab("app"), sql }); diff --git a/apps/desktop/src/composables/useSqlExecution.ts b/apps/desktop/src/composables/useSqlExecution.ts index 8e064fa5ce..d494a3d3d3 100644 --- a/apps/desktop/src/composables/useSqlExecution.ts +++ b/apps/desktop/src/composables/useSqlExecution.ts @@ -144,11 +144,12 @@ export function useSqlExecution(deps: { async function resolvedExecutableSql(source?: SqlExecutionOverride): Promise<{ sql: string; sourceOffset?: number; editorViewportRequestId?: number }> { const atSetEnabled = resolveSqlVariableSyntaxToggles(settingsStore.editorSettings.sqlVariableSyntaxOverrides, deps.activeConnection.value?.db_type, settingsStore.editorSettings.sqlVariableSubstitutionEnabled).atSet; - const expand = (sql: string) => (atSetEnabled ? expandSqlVariables(sql).sql : sql); + const expand = (sql: string, declarationSql?: string) => (atSetEnabled ? expandSqlVariables(sql, { declarationSql }).sql : sql); if (typeof source === "string") return { sql: expand(source) }; const resolved = deps.resolveExecutableSql ? await deps.resolveExecutableSql(source) : isSqlExecutionSnapshot(source) ? resolveExecutableSql(source.fullSql, source.selectedSql, { cursorPos: source.cursorPos }) : deps.executableSql.value; - const sql = expand(resolved); + const declarationSql = isSqlExecutionSnapshot(source) ? source.fullSql.slice(0, source.selectionTo) : resolved; + const sql = expand(resolved, declarationSql); const editorViewportRequestId = isSqlExecutionSnapshot(source) ? source.editorViewportRequestId : undefined; if (!isSqlExecutionSnapshot(source) || !source.selectedSql.trim() || sql !== resolved) return { sql, editorViewportRequestId }; diff --git a/apps/desktop/src/lib/__tests__/sql/sqlVariables.spec.ts b/apps/desktop/src/lib/__tests__/sql/sqlVariables.spec.ts index b97ca1b518..d77d90ef2f 100644 --- a/apps/desktop/src/lib/__tests__/sql/sqlVariables.spec.ts +++ b/apps/desktop/src/lib/__tests__/sql/sqlVariables.spec.ts @@ -19,6 +19,22 @@ describe("expandSqlVariables", () => { expect(expandSqlVariables(sql).sql).toBe("select * from t where created_at < '2026-07-04 00:00:00'"); }); + it("inlines a declared shell-style reference", () => { + const sql = ["@set postid = '224';", "select * from t where post_id = ${postid}"].join("\n"); + expect(expandSqlVariables(sql).sql).toBe("select * from t where post_id = '224'"); + }); + + 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"); + expect(expandSqlVariables(selectedSql, { declarationSql }).sql).toBe("select * from t where post_id = '224'"); + }); + + it("leaves undeclared shell-style references for parameter prompting", () => { + const sql = ["@set postid = '224';", "select ${postid}, ${missing}"].join("\n"); + expect(expandSqlVariables(sql).sql).toBe("select '224', ${missing}"); + }); + 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"); diff --git a/apps/desktop/src/lib/sql/sqlVariables.ts b/apps/desktop/src/lib/sql/sqlVariables.ts index b193158ad8..0e5fc9bd83 100644 --- a/apps/desktop/src/lib/sql/sqlVariables.ts +++ b/apps/desktop/src/lib/sql/sqlVariables.ts @@ -1,16 +1,15 @@ // 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: +// and reference them elsewhere as `${name}` or the legacy `@name` form. 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. +// This runs before the placeholder parameter system (`sqlParameters.ts`). A +// reference is only expanded when a matching `@set` declaration exists; any +// unresolved `${name}` or `@name` remains available to the parameter dialog. const VARIABLE_NAME_START_RE = /[\p{L}_]/u; const VARIABLE_NAME_CHAR_RE = /[\p{L}\p{N}_]/u; @@ -20,6 +19,10 @@ export interface SqlVariableExpansion { expanded: boolean; } +export interface SqlVariableExpansionOptions { + declarationSql?: string; +} + interface DeclarationSpan { name: string; value: string; @@ -30,12 +33,13 @@ interface DeclarationSpan { /** * 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`. + * Returns the SQL with declarations removed and matching `${name}`/`@name` + * references replaced by their declared values. `declarationSql` can provide + * the document prefix for a separately executed statement. */ -export function expandSqlVariables(sql: string): SqlVariableExpansion { - const declarations = collectDeclarations(sql); +export function expandSqlVariables(sql: string, options: SqlVariableExpansionOptions = {}): SqlVariableExpansion { + const declarationSql = options.declarationSql ?? sql; + const declarations = collectDeclarations(declarationSql); if (!declarations.length) return { sql, expanded: false }; const values = new Map(); @@ -43,16 +47,17 @@ export function expandSqlVariables(sql: string): SqlVariableExpansion { 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. + // Only spans inside the execution target are removed. Declarations supplied + // by the surrounding document are context and must not affect target offsets. + const targetDeclarations = declarationSql === sql ? declarations : collectDeclarations(sql); let result = sql; - for (let i = declarations.length - 1; i >= 0; i -= 1) { - const declaration = declarations[i]; + for (let i = targetDeclarations.length - 1; i >= 0; i -= 1) { + const declaration = targetDeclarations[i]; result = stripDeclaration(result, declaration.start, declaration.end); } result = replaceReferences(result, values); - return { sql: result, expanded: true }; + return { sql: result, expanded: result !== sql }; } function collectDeclarations(sql: string): DeclarationSpan[] { @@ -220,6 +225,18 @@ function replaceReferences(sql: string, values: Map): string { dollarQuoteEnd = marker; continue; } + if (next === "{") { + const name = readVariableName(sql, i + 2); + const end = i + 2 + name.length; + if (name && sql[end] === "}") { + const value = values.get(name.toLowerCase()); + if (value !== undefined) { + result += value; + i = end + 1; + continue; + } + } + } } if (ch === "@" && next !== "@" && sql[i - 1] !== "@") { const name = readVariableName(sql, i + 1);