marimo-team
diff --git a/‎README.md‎
Lines changed: 32 additions & 1 deletion b/‎README.md‎
Lines changed: 32 additions & 1 deletion
diff --git a/‎demo/data.ts‎
Lines changed: 13 additions & 2 deletions b/‎demo/data.ts‎
Lines changed: 13 additions & 2 deletions
diff --git a/‎demo/index.html‎
Lines changed: 23 additions & 1 deletion b/‎demo/index.html‎
Lines changed: 23 additions & 1 deletion
diff --git a/‎demo/index.ts‎
Lines changed: 9 additions & 2 deletions b/‎demo/index.ts‎
Lines changed: 9 additions & 2 deletions
diff --git a/‎src/__tests__/index.test.ts‎
Lines changed: 9 additions & 0 deletions b/‎src/__tests__/index.test.ts‎
Lines changed: 9 additions & 0 deletions
diff --git a/‎src/index.ts‎
Lines changed: 22 additions & 1 deletion b/‎src/index.ts‎
Lines changed: 22 additions & 1 deletion
diff --git a/‎src/sql/__tests__/cte-completion-source.test.ts‎
Lines changed: 141 additions & 7 deletions b/‎src/sql/__tests__/cte-completion-source.test.ts‎
Lines changed: 141 additions & 7 deletions
@@ -8,8 +8,9 @@ A CodeMirror extension for SQL linting and visual gutter indicators. Built by an
 - 🧠 **Schema-aware linting** - Warns about unknown tables, unknown columns, and ambiguous column references based on your schema
 - 🎨 **Visual gutter** - Color-coded statement indicators and error highlighting
 - 💡 **Hover tooltips** - Schema info, keywords, and column details on hover
-- 🔮 **CTE autocomplete** - Auto-complete support for CTEs
+- 🔮 **CTE autocomplete** - Statement-scoped completion of CTE names and their output columns
 - 🏷️ **Alias resolution** - Hover and completion understand table aliases (`SELECT u.name FROM users u`)
+- 🧭 **Navigation** - Go-to-definition, reference highlighting, and rename for CTEs and aliases
 - 🎯 **Query-aware resolution** - Context-sensitive schema and column suggestions
 - 🔍 **Additional dialects** - DuckDB, BigQuery, Dremio
 - 🛠️ **Custom renderers** - Customizable tooltip rendering for tables, columns, and keywords
@@ -110,6 +111,36 @@ diagnostics carry `source: "sql-schema"`; syntax diagnostics use
 `source: "sql-parser"`. If the schema is provided as a function, it is called
 on every lint pass and should be cheap/memoized.
 
+### Navigation: go-to-definition, highlights, rename
+
+`sqlExtension` includes navigation for statement-local identifiers (CTE names,
+table aliases, select aliases) by default: the references of the identifier
+under the cursor are highlighted, and Mod-click (Cmd/Ctrl-click) on a
+resolvable identifier jumps to its definition. Keybindings (`F12`/`Mod-b` for
+go-to-definition, `F2` for rename) are opt-in:
+
+```ts
+import { renameSqlIdentifier, sqlExtension } from "@marimo-team/codemirror-sql";
+
+sqlExtension({
+  enableNavigation: true, // default
+  navigationConfig: {
+    keymap: true, // enable F12 / Mod-b / F2
+    // Supply your own rename UI (defaults to window.prompt)
+    prompt: (currentName) => window.prompt(`Rename '${currentName}' to:`, currentName),
+  },
+});
+
+// Rename programmatically: rewrites the definition and all references
+// in a single undo step
+await renameSqlIdentifier(view, { prompt: () => "new_name" });
+```
+
+The pieces are also exported individually: `sqlHighlightReferences`,
+`sqlGotoDefinition`, `sqlNavigationKeymap`, `gotoSqlDefinition`, and
+`findReferences`. Rename returns `false` when the identifier can't be
+confidently resolved — it never falls back to text search-and-replace.
+
 ## Additional Dialects
 
 This extension adds support for additional dialects:
 
@@ -4,9 +4,20 @@ import type { Schema } from "./custom-renderers";
 export const defaultSqlDoc = `-- Welcome to the SQL Editor Demo!
 -- Try editing the queries below to see real-time validation
 
-WITH cte_name AS (
-  SELECT * FROM users
+-- Put the cursor on a CTE name or alias to highlight its references;
+-- Cmd/Ctrl-click (or F12) jumps to the definition, F2 renames.
+WITH recent_orders AS (
+  SELECT customer_id, total_amount FROM orders WHERE order_date >= '2024-01-01'
+),
+top_customers AS (
+  SELECT customer_id, SUM(total_amount) AS total_spent
+  FROM recent_orders
+  GROUP BY customer_id
 )
+SELECT c.first_name, t.total_spent AS amount
+FROM customers c
+JOIN top_customers t ON t.customer_id = c.id
+ORDER BY amount DESC;
 
 -- Valid queries (no errors):
 SELECT id, name, email
 
@@ -28,10 +28,15 @@ <h1 class="text-4xl font-bold mb-4">
       <div class="grid gap-8">
         <section>
           <h2 class="text-2xl font-bold mb-4">SQL Editor with Diagnostics</h2>
-          <p class="mb-4 text-gray-600">
+          <p class="mb-2 text-gray-600">
             Try typing invalid SQL syntax to see real-time error highlighting and messages.
             Valid tables are: <span class="font-bold">users, posts, orders, customers, categories</span>
           </p>
+          <p class="mb-4 text-gray-600">
+            Put the cursor on a CTE name or alias to highlight its references,
+            <span class="font-bold">Cmd/Ctrl-click</span> (or <span class="font-bold">F12</span>) to jump to its
+            definition, and press <span class="font-bold">F2</span> to rename it everywhere in the statement.
+          </p>
 
           <div class="flex justify-end items-center mb-2">
             <select id="database-select" class="border rounded-md p-1.5">
@@ -64,6 +69,9 @@ <h3 class="font-semibold text-lg">Valid Queries</h3>
               <button class="example-btn w-full p-3 text-left bg-green-50 hover:bg-green-100 border border-green-200 rounded">
                 <code class="text-sm">INSERT INTO products (name, price) VALUES ('Laptop', 999.99)</code>
               </button>
+              <button class="example-btn w-full p-3 text-left bg-green-50 hover:bg-green-100 border border-green-200 rounded">
+                <code class="text-sm">WITH active_users AS (SELECT id, name FROM users WHERE active = true) SELECT a.name FROM active_users a ORDER BY a.name</code>
+              </button>
               <button class="example-btn w-full p-3 text-left bg-blue-50 hover:bg-blue-100 border border-blue-200 rounded">
                 <code class="text-sm">from nyc.rideshare select * limit 100</code>
               </button>
@@ -102,6 +110,20 @@ <h3 class="font-semibold mb-2">Multiple SQL Dialects</h3>
               <h3 class="font-semibold mb-2">DuckDB Support</h3>
               <p class="text-sm text-gray-700">Special support for DuckDB-specific syntax like "from table select * limit 100".</p>
             </div>
+            <div class="p-4 bg-blue-50 rounded-lg">
+              <h3 class="font-semibold mb-2">CTE &amp; Alias Navigation</h3>
+              <p class="text-sm text-gray-700">
+                Go-to-definition (Cmd/Ctrl-click or F12), reference highlighting, and rename (F2) for CTE names,
+                table aliases, and select aliases.
+              </p>
+            </div>
+            <div class="p-4 bg-blue-50 rounded-lg">
+              <h3 class="font-semibold mb-2">CTE Autocomplete</h3>
+              <p class="text-sm text-gray-700">
+                Statement-scoped completion of CTE names and their output columns, including after
+                <code>my_cte.</code>.
+              </p>
+            </div>
             <div class="p-4 bg-blue-50 rounded-lg">
               <h3 class="font-semibold mb-2">TypeScript Support</h3>
               <p class="text-sm text-gray-700">Full TypeScript support with comprehensive type definitions.</p>
 
@@ -5,7 +5,7 @@ import { keymap } from "@codemirror/view";
 import { basicSetup, EditorView } from "codemirror";
 import {
   aliasColumnCompletionSource,
-  cteCompletionSource,
+  createCteCompletionSource,
   DefaultSqlTooltipRenders,
   defaultSqlHoverTheme,
   NodeSqlParser,
@@ -174,9 +174,16 @@ function initializeEditor() {
         theme: defaultSqlHoverTheme("light"),
         parser,
       },
+      // Reference highlighting and go-to-definition for CTEs/aliases
+      enableNavigation: true,
+      navigationConfig: {
+        keymap: true, // F12/Mod-b go-to-definition, F2 rename
+        parser,
+      },
     }),
     defaultDialect.language.data.of({
-      autocomplete: cteCompletionSource,
+      // Statement-scoped CTE names and their output columns
+      autocomplete: createCteCompletionSource({ parser }),
     }),
     defaultDialect.language.data.of({
       // Complete `u.` -> columns of `users` in `SELECT ... FROM users u`
 
@@ -12,15 +12,24 @@ describe("index.ts exports", () => {
         "DefaultSqlTooltipRenders",
         "NodeSqlParser",
         "QueryContextAnalyzer",
+        "SqlReferenceResolver",
         "SqlStructureAnalyzer",
         "aliasColumnCompletionSource",
         "analyzeQueryContext",
+        "createCteCompletionSource",
         "cteCompletionSource",
         "defaultSqlHoverTheme",
+        "findReferences",
+        "gotoSqlDefinition",
+        "renameSqlIdentifier",
         "resolveSqlSchema",
         "sqlExtension",
+        "sqlGotoDefinition",
+        "sqlHighlightReferences",
         "sqlHover",
         "sqlLinter",
+        "sqlNavigation",
+        "sqlNavigationKeymap",
         "sqlSchemaFacet",
         "sqlSemanticLinter",
         "sqlStructureGutter",
 
@@ -2,7 +2,11 @@ export {
   type AliasCompletionConfig,
   aliasColumnCompletionSource,
 } from "./sql/alias-completion-source.js";
-export { cteCompletionSource } from "./sql/cte-completion-source.js";
+export {
+  createCteCompletionSource,
+  type CteCompletionConfig,
+  cteCompletionSource,
+} from "./sql/cte-completion-source.js";
 export { type SqlLinterConfig, sqlLinter } from "./sql/diagnostics.js";
 export { type SqlExtensionConfig, sqlExtension } from "./sql/extension.js";
 export type {
@@ -12,6 +16,15 @@ export type {
   SqlKeywordInfo,
 } from "./sql/hover.js";
 export { DefaultSqlTooltipRenders, defaultSqlHoverTheme, sqlHover } from "./sql/hover.js";
+export {
+  gotoSqlDefinition,
+  renameSqlIdentifier,
+  type SqlNavigationConfig,
+  sqlGotoDefinition,
+  sqlHighlightReferences,
+  sqlNavigation,
+  sqlNavigationKeymap,
+} from "./sql/navigation-extension.js";
 export {
   NodeSqlParser,
   type NodeSqlParserOptions,
@@ -26,6 +39,14 @@ export {
   type QueryContextCte,
   type QueryContextTable,
 } from "./sql/query-context.js";
+export {
+  findReferences,
+  type SqlIdentifierKind,
+  type SqlRange,
+  type SqlReferenceConfig,
+  SqlReferenceResolver,
+  type SqlReferenceResult,
+} from "./sql/references.js";
 export { resolveSqlSchema, type SqlSchemaSource, sqlSchemaFacet } from "./sql/schema-facet.js";
 export {
   type SemanticSeverity,
 
@@ -23,9 +23,9 @@ function createMockContext(doc: string, pos: number, explicit = false): Completi
     matchBefore: (pattern: RegExp) => {
       const before = doc.slice(0, pos);
 
-      // For \w* pattern, find the word at the end
-      if (pattern.source === "\\w*") {
-        const wordMatch = before.match(/(\w*)$/);
+      // For word patterns, find the word at the end
+      if (pattern.source === "\\w*" || pattern.source === "[\\w$]*") {
+        const wordMatch = before.match(/([\w$]*)$/);
         if (!wordMatch) return null;
 
         const text = wordMatch[1] || "";
@@ -182,7 +182,7 @@ describe("cteCompletionSource", () => {
       expect(result?.options[0].label).toBe("filtered_users");
     });
 
-    it("should handle multiple WITH clauses in different statements", async () => {
+    it("scopes CTEs to the statement containing the cursor", async () => {
       const sql = `WITH first_cte AS (SELECT 1)
       SELECT * FROM first_cte;
 
@@ -193,10 +193,22 @@ describe("cteCompletionSource", () => {
       const result = await getCompletionResult(context);
 
       expect(result).toBeTruthy();
-      expect(result?.options).toHaveLength(2);
+      // first_cte belongs to the previous statement and must not leak in
+      expect(result?.options.map((opt) => opt.label)).toEqual(["second_cte"]);
+    });
 
-      const labels = result?.options.map((opt) => opt.label).sort();
-      expect(labels).toEqual(["first_cte", "second_cte"]);
+    it("offers the first statement's CTE when the cursor is in it", async () => {
+      const sql = `WITH first_cte AS (SELECT 1)
+      SELECT * FROM first_;
+
+      WITH second_cte AS (SELECT 2)
+      SELECT * FROM second_cte`;
+
+      const pos = sql.indexOf("first_;") + "first_".length;
+      const context = createMockContext(sql, pos, true);
+      const result = await getCompletionResult(context);
+
+      expect(result?.options.map((opt) => opt.label)).toEqual(["first_cte"]);
     });
   });
 
@@ -316,5 +328,127 @@ describe("cteCompletionSource", () => {
       expect(labels).toContain("a");
       expect(labels).toContain("b");
     });
+
+    it("detects 3+ comma-separated CTEs with nested parens mid-edit", async () => {
+      const sql =
+        "WITH a AS (SELECT max(x) FROM (SELECT 1) s), b AS (SELECT count(*) FROM t), c AS (SELECT 3) SELECT * FROM ";
+
+      const context = createMockContext(sql, sql.length, true);
+      const result = await getCompletionResult(context);
+
+      const labels = result?.options.map((option) => option.label) ?? [];
+      expect(labels).toEqual(["a", "b", "c"]);
+    });
+  });
+
+  describe("quoted CTE names", () => {
+    it("detects quoted CTE names", async () => {
+      const sql = `WITH "My Stats" AS (SELECT 1) SELECT * FROM `;
+
+      const context = createMockContext(sql, sql.length, true);
+      const result = await getCompletionResult(context);
+
+      expect(result?.options.map((opt) => opt.label)).toEqual(["My Stats"]);
+      expect(result?.options[0].apply).toBe('"My Stats"');
+    });
+
+    it("does not add an apply override for bare names", async () => {
+      const sql = "WITH stats AS (SELECT 1) SELECT * FROM ";
+
+      const context = createMockContext(sql, sql.length, true);
+      const result = await getCompletionResult(context);
+
+      expect(result?.options[0].apply).toBeUndefined();
+    });
+
+    it("completes names containing $ as one token", async () => {
+      const sql = "WITH t$1 AS (SELECT 1) SELECT * FROM t$";
+
+      const context = createMockContext(sql, sql.length, true);
+      const result = await getCompletionResult(context);
+
+      expect(result?.from).toBe(sql.lastIndexOf("t$"));
+      expect(result?.options[0].label).toBe("t$1");
+    });
+
+    it("quotes non-bare column labels on apply", async () => {
+      const sql = 'WITH t("My Col", b) AS (SELECT 1, 2) SELECT t. FROM t';
+      const pos = sql.indexOf("t. FROM") + 2;
+
+      const context = createMockContext(sql, pos, false);
+      const result = await getCompletionResult(context);
+
+      const myCol = result?.options.find((opt) => opt.label === "My Col");
+      expect(myCol?.apply).toBe('"My Col"');
+      expect(result?.options.find((opt) => opt.label === "b")?.apply).toBeUndefined();
+    });
+  });
+
+  describe("CTE column completion", () => {
+    it("completes columns after <cte>. from a declared column list", async () => {
+      const sql = "WITH t(a, b) AS (SELECT 1, 2) SELECT t. FROM t";
+      const pos = sql.indexOf("t. FROM") + 2;
+
+      const context = createMockContext(sql, pos, false);
+      const result = await getCompletionResult(context);
+
+      expect(result).toBeTruthy();
+      expect(result?.from).toBe(pos);
+      expect(result?.options.map((opt) => opt.label)).toEqual(["a", "b"]);
+      expect(result?.options[0].type).toBe("property");
+      expect(result?.options[0].detail).toBe("column of t");
+    });
+
+    it("completes columns after <cte>. inferred from the select list", async () => {
+      const sql = "WITH t AS (SELECT id, name AS n FROM users) SELECT t. FROM t";
+      const pos = sql.indexOf("t. FROM") + 2;
+
+      const context = createMockContext(sql, pos, false);
+      const result = await getCompletionResult(context);
+
+      expect(result?.options.map((opt) => opt.label)).toEqual(["id", "n"]);
+    });
+
+    it("offers no columns for a SELECT * CTE", async () => {
+      const sql = "WITH t AS (SELECT * FROM users) SELECT t. FROM t";
+      const pos = sql.indexOf("t. FROM") + 2;
+
+      const context = createMockContext(sql, pos, false);
+      const result = await getCompletionResult(context);
+
+      expect(result).toBeNull();
+    });
+
+    it("does not complete columns for a multi-segment qualifier", async () => {
+      const sql = "WITH t(a) AS (SELECT 1) SELECT db.t. FROM t";
+      const pos = sql.indexOf("db.t. FROM") + "db.t.".length;
+
+      const context = createMockContext(sql, pos, false);
+      const result = await getCompletionResult(context);
+
+      expect(result).toBeNull();
+    });
+
+    it("offers CTE columns unqualified in a column position", async () => {
+      const sql = "WITH t(a, b) AS (SELECT 1, 2) SELECT  FROM t";
+      const pos = sql.indexOf("SELECT  FROM") + "SELECT ".length;
+
+      const context = createMockContext(sql, pos, true);
+      const result = await getCompletionResult(context);
+
+      const labels = result?.options.map((opt) => opt.label) ?? [];
+      expect(labels).toContain("t");
+      expect(labels).toContain("a");
+      expect(labels).toContain("b");
+    });
+
+    it("does not offer columns right after FROM", async () => {
+      const sql = "WITH t(a, b) AS (SELECT 1, 2) SELECT a FROM ";
+
+      const context = createMockContext(sql, sql.length, true);
+      const result = await getCompletionResult(context);
+
+      expect(result?.options.map((opt) => opt.label)).toEqual(["t"]);
+    });
   });
 });