From d48167ef6650264208ef06e215f7f4512b1f92c4 Mon Sep 17 00:00:00 2001
From: Eve <92513238+Eve-z1@users.noreply.github.com>
Date: Thu, 10 Sep 2026 12:44:23 +0800
Subject: [PATCH] fix(editor): include SQL Server variable prefix in word
 selection

---
 apps/desktop/src/components/editor/QueryEditor.vue |  3 ++-
 .../editor/queryEditorLineComment.spec.ts          | 14 +++++++++++++-
 .../src/lib/editor/queryEditorLineComment.ts       |  4 ++++
 3 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/apps/desktop/src/components/editor/QueryEditor.vue b/apps/desktop/src/components/editor/QueryEditor.vue
index dbab1dfa8a..2b9f93e4dc 100644
--- a/apps/desktop/src/components/editor/QueryEditor.vue
+++ b/apps/desktop/src/components/editor/QueryEditor.vue
@@ -140,7 +140,7 @@ import { sqlSemanticTableNameSpansForSyntaxTree } from "@/lib/editor/codemirrorS
 import { startsQueryEditorRectangularSelection, usesQueryEditorObjectNavigationModifier } from "@/lib/editor/queryEditorPointerSelection";
 import { LARGE_PASTE_HISTORY_USER_EVENT, normalizeQueryEditorPasteText, recoverableNativePasteSuffix, shouldRecoverLargeTauriPaste } from "@/lib/editor/queryEditorLargePaste";
 import { computePasteCaretResyncTarget } from "@/lib/editor/queryEditorPasteCaretResync";
-import { queryEditorCommentTokens, queryEditorLineCommentToken } from "@/lib/editor/queryEditorLineComment";
+import { queryEditorCommentTokens, queryEditorLineCommentToken, queryEditorWordLanguageData } from "@/lib/editor/queryEditorLineComment";
 import { createShellLineCommentHighlight } from "@/lib/editor/codemirrorShellLineCommentHighlight";
 import { extendQueryEditorSelection, runQueryEditorAltExtendSelection } from "@/lib/editor/queryEditorExtendSelection";
 import { addNextQueryEditorSelectionOccurrence, selectAllQueryEditorSelectionOccurrences } from "@/lib/editor/queryEditorOccurrenceSelection";
@@ -6035,6 +6035,7 @@ onMounted(async () => {
     // Non-SQL editors (MongoDB shell) keep the SQL grammar for highlighting, so override the
     // comment marker that toggleLineComment reads from language data.
     Prec.highest(EditorState.languageData.of(() => [{ commentTokens: queryEditorCommentTokens(props.databaseType) }])),
+    Prec.highest(EditorState.languageData.of(() => queryEditorWordLanguageData(props.databaseType))),
     // The SQL grammar does not tokenize `//`, so those comments are highlighted by hand.
     queryEditorLineCommentToken(props.databaseType) === "//" ? shellLineCommentHighlightPlugin : [],
   ];
diff --git a/apps/desktop/src/lib/__tests__/editor/queryEditorLineComment.spec.ts b/apps/desktop/src/lib/__tests__/editor/queryEditorLineComment.spec.ts
index 329e922ff5..af57bc884c 100644
--- a/apps/desktop/src/lib/__tests__/editor/queryEditorLineComment.spec.ts
+++ b/apps/desktop/src/lib/__tests__/editor/queryEditorLineComment.spec.ts
@@ -3,7 +3,7 @@ import { toggleLineComment } from "@codemirror/commands";
 import { sql } from "@codemirror/lang-sql";
 import { EditorState, Prec, type Transaction } from "@codemirror/state";
 import { describe, expect, it, vi } from "vitest";
-import { queryEditorCommentTokens, queryEditorLineCommentToken } from "@/lib/editor/queryEditorLineComment";
+import { queryEditorCommentTokens, queryEditorLineCommentToken, queryEditorWordLanguageData } from "@/lib/editor/queryEditorLineComment";
 
 const queryEditorSource = readFileSync(new URL("../../../components/editor/QueryEditor.vue", import.meta.url), "utf8");
 const editorThemesSource = readFileSync(new URL("../../editor/editorThemes.ts", import.meta.url), "utf8");
@@ -47,6 +47,18 @@ describe("queryEditorLineCommentToken", () => {
   });
 });
 
+describe("QueryEditor word selection", () => {
+  it("includes the @ prefix in SQL Server variable words", () => {
+    const sqlServer = EditorState.create({ doc: "@name", extensions: [EditorState.languageData.of(() => queryEditorWordLanguageData("sqlserver"))] });
+    const mysql = EditorState.create({ doc: "@name", extensions: [EditorState.languageData.of(() => queryEditorWordLanguageData("mysql"))] });
+    const systemVariable = EditorState.create({ doc: "@@ROWCOUNT", extensions: [EditorState.languageData.of(() => queryEditorWordLanguageData("sqlserver"))] });
+
+    for (let position = 0; position <= 5; position += 1) expect(sqlServer.wordAt(position)).toMatchObject({ from: 0, to: 5 });
+    expect(mysql.wordAt(2)).toMatchObject({ from: 1, to: 5 });
+    expect(systemVariable.wordAt(1)).toMatchObject({ from: 0, to: 10 });
+  });
+});
+
 describe("QueryEditor line comment", () => {
   it("overrides the language comment tokens in the SQL language compartment", () => {
     expect(queryEditorSource).toContain("Prec.highest(EditorState.languageData.of(() => [{ commentTokens: queryEditorCommentTokens(props.databaseType) }]))");
diff --git a/apps/desktop/src/lib/editor/queryEditorLineComment.ts b/apps/desktop/src/lib/editor/queryEditorLineComment.ts
index 39edfd160d..7408895b91 100644
--- a/apps/desktop/src/lib/editor/queryEditorLineComment.ts
+++ b/apps/desktop/src/lib/editor/queryEditorLineComment.ts
@@ -18,3 +18,7 @@ export function queryEditorCommentTokens(dbType?: DatabaseType) {
     block: { open: "/*", close: "*/" },
   };
 }
+
+export function queryEditorWordLanguageData(dbType?: DatabaseType) {
+  return dbType === "sqlserver" ? [{ wordChars: "@" }] : [];
+}