@marimo-team/frontend 0.23.16-dev9 → 0.23.16

package/src/core/codemirror/language/languages/python.ts

@@ -38,7 +38,11 @@ import { createLspMarkdownRenderer } from "../../lsp/markdown-renderer";
 import { NotebookLanguageServerClient } from "../../lsp/notebook-lsp";
 import { createTransport } from "../../lsp/transports";
 import { CellDocumentUri, type ILanguageServerClient } from "../../lsp/types";
-import { getLspRootUri, getLspWorkspaceFolders } from "../../lsp/utils";
+import {
+  getLspRootUri,
+  getLspWorkspaceFolders,
+  isKnownCellDocumentUri,
+} from "../../lsp/utils";
 import {
   clickablePlaceholderExtension,
   smartPlaceholderExtension,
@@ -338,6 +342,21 @@ export class PythonLanguageAdapter implements LanguageAdapter<{}> {
           clients.length === 1
             ? (clients[0] as NotebookLanguageServerClient)
             : new FederatedLanguageServerClient(clients);
+        const openExternalLocation = (result: { uri: string }) => {
+          Logger.debug("openExternalLocation", result);
+          // Neither the merged document nor a cell is a file on disk. The
+          // plugin handles same-cell locations itself; revealing a location in
+          // another cell isn't supported yet.
+          if (
+            client.documentUri === result.uri ||
+            isKnownCellDocumentUri(result.uri)
+          ) {
+            return;
+          }
+          getRequestClient().openFile({
+            path: result.uri.replace("file://", ""),
+          });
+        };
 
         return [
           languageServerWithClient({
@@ -346,6 +365,7 @@ export class PythonLanguageAdapter implements LanguageAdapter<{}> {
             allowHTMLContent: true,
             markdownRenderer: createLspMarkdownRenderer(),
             useSnippetOnCompletion: true,
+            clientSideFiltering: true,
             hoverConfig: hoverOptions,
             completionConfig: autocompleteOptions,
             // Default to false
@@ -362,16 +382,8 @@ export class PythonLanguageAdapter implements LanguageAdapter<{}> {
               rename: hotkeys.getHotkey("cell.renameSymbol").key,
             },
             completionMatchBefore,
-            onGoToDefinition: (result) => {
-              Logger.debug("onGoToDefinition", result);
-              if (client.documentUri === result.uri) {
-                // Local definition
-                return;
-              }
-              getRequestClient().openFile({
-                path: result.uri.replace("file://", ""),
-              });
-            },
+            onGoToDefinition: openExternalLocation,
+            onShowLocation: openExternalLocation,
           }),
           documentUri.of(CellDocumentUri.of(cellId)),
         ];

package/src/core/codemirror/language/languages/sql/completion-store.ts

@@ -109,7 +109,11 @@ class SQLCompletionStore {
     return {
       shouldAddLocalTables: true,
       schema: builder.build(),
-      defaultSchema: default_schema ?? undefined,
+      defaultSchema:
+        default_schema ??
+        (defaultDb?.schemas.length === 1
+          ? defaultDb.schemas[0]?.name
+          : undefined),
     };
   }
 

package/src/core/codemirror/language/languages/sql/renderers.tsx

@@ -52,7 +52,7 @@ const SOURCE_TYPE_COLORS = {
   catalog: "bg-(--purple-4) text-(--purple-11)",
 } as const;
 
-const CONTAINER_STYLES = "p-3 min-w-[250px] flex flex-col divide-y";
+export const CONTAINER_STYLES = "p-3 min-w-[250px] flex flex-col divide-y";
 
 const columnsText = new PluralWord("column", "columns");
 const rowsText = new PluralWord("row", "rows");
@@ -61,7 +61,7 @@ const tablesText = new PluralWord("table", "tables");
 const databasesText = new PluralWord("database", "databases");
 
 // Helper components and functions
-const SectionHeader: React.FC<{
+export const SectionHeader: React.FC<{
   icon: React.ReactNode;
   title: string;
   badge?: React.ReactNode;
@@ -73,7 +73,7 @@ const SectionHeader: React.FC<{
   </div>
 );
 
-const MetadataRow: React.FC<{
+export const MetadataRow: React.FC<{
   label: string;
   value: React.ReactNode;
 }> = ({ label, value }) => (

package/src/core/codemirror/language/languages/sql/sql.ts

@@ -12,12 +12,17 @@ import {
   type ViewUpdate,
 } from "@codemirror/view";
 import {
+  aliasColumnCompletionSource,
+  createCteCompletionSource,
   defaultSqlHoverTheme,
   NodeSqlParser,
   type NodeSqlParserResult,
+  QueryContextAnalyzer,
   type SupportedDialects as ParserDialects,
   type SqlParseError,
   sqlExtension,
+  SqlStructureAnalyzer,
+  unqualifiedColumnCompletionSource,
 } from "@marimo-team/codemirror-sql";
 import { DuckDBDialect } from "@marimo-team/codemirror-sql/dialects";
 import { type SQLMetadata, SQLParser } from "@marimo-team/smart-cells";
@@ -61,7 +66,7 @@ import { getSQLMode, type SQLMode } from "./sql-mode";
 import { isKnownDialect } from "./utils";
 
 const DEFAULT_DIALECT = DuckDBDialect;
-const DEFAULT_PARSER_DIALECT = "DuckDB";
+const DEFAULT_PARSER_DIALECT: ParserDialects = "DuckDB";
 
 // A compartment for the SQL config, so we can update the config of codemirror
 const sqlConfigCompartment = new Compartment();
@@ -142,6 +147,7 @@ export class SQLLanguageAdapter implements LanguageAdapter<SQLLanguageAdapterMet
     _placeholderType: PlaceholderType,
     lspConfig: LSPConfig & { diagnostics: DiagnosticsConfig },
   ): Extension[] {
+    const analysis = createSQLAnalysis();
     const extensions = [
       // This can be updated with a dispatch effect
       sqlConfigCompartment.of(sql({ dialect: DEFAULT_DIALECT })),
@@ -167,6 +173,10 @@ export class SQLLanguageAdapter implements LanguageAdapter<SQLLanguageAdapterMet
           variableCompletionSource,
           // Completions for dialect keywords
           customKeywordCompletionSource(),
+          // Completions based on the current query
+          analysis.cteCompletionSource,
+          analysis.aliasCompletionSource,
+          analysis.columnCompletionSource,
         ],
       }),
     ];
@@ -185,10 +195,23 @@ export class SQLLanguageAdapter implements LanguageAdapter<SQLLanguageAdapterMet
 
       extensions.push(
         sqlExtension({
+          schema: getSchema,
           enableLinting: true,
           linterConfig: {
             delay: 250, // Delay before running validation
             parser: parser,
+            // CustomSqlParser performs backend DuckDB validation in validateSql.
+            perStatement: false,
+          },
+          enableSemanticLinting: true,
+          semanticLinterConfig: {
+            parser: analysis.parser,
+            structureAnalyzer: analysis.structureAnalyzer,
+            severity: {
+              unknownTable: "off",
+              unknownColumn: "warning",
+              ambiguousColumn: "warning",
+            },
           },
           enableGutterMarkers: true,
           gutterConfig: {
@@ -198,14 +221,21 @@ export class SQLLanguageAdapter implements LanguageAdapter<SQLLanguageAdapterMet
             parser: parser,
           },
           hoverConfig: {
-            schema: getSchema, // Use the same schema as autocomplete
             hoverTime: 300, // 300ms hover delay
             enableKeywords: true, // Show keyword information
             enableTables: true, // Show table information
             enableColumns: true, // Show column information
-            parser: parser,
+            parser: analysis.parser,
+            contextAnalyzer: analysis.contextAnalyzer,
             theme: defaultSqlHoverTheme(theme),
           },
+          enableNavigation: true,
+          navigationConfig: {
+            contextAnalyzer: analysis.contextAnalyzer,
+            keymap: false,
+            parser: analysis.parser,
+            structureAnalyzer: analysis.structureAnalyzer,
+          },
         }),
         EditorView.updateListener.of((update) => {
           if (update.focusChanged) {
@@ -223,8 +253,127 @@ export class SQLLanguageAdapter implements LanguageAdapter<SQLLanguageAdapterMet
   }
 }
 
+function createSQLAnalysis() {
+  const parser = new NodeSqlParser({
+    getParserOptions: (state: EditorState) => {
+      return {
+        database: getAnalysisDialect(state),
+      };
+    },
+  });
+  const contextAnalyzer = new DialectAwareQueryContextAnalyzer(parser);
+  const structureAnalyzer = new DialectAwareSqlStructureAnalyzer(parser);
+
+  return {
+    parser,
+    contextAnalyzer,
+    structureAnalyzer,
+    cteCompletionSource: createCteCompletionSource({
+      parser,
+      contextAnalyzer,
+      structureAnalyzer,
+    }),
+    aliasCompletionSource: aliasColumnCompletionSource({
+      schema: getSchema,
+      parser,
+      contextAnalyzer,
+    }),
+    columnCompletionSource: unqualifiedColumnCompletionSource({
+      schema: getSchema,
+      parser,
+      contextAnalyzer,
+    }),
+  };
+}
+
+function getAnalysisDialect(state: EditorState): ParserDialects {
+  return guessParserDialect(state) ?? DEFAULT_PARSER_DIALECT;
+}
+
+class DialectAwareQueryContextAnalyzer extends QueryContextAnalyzer {
+  private readonly analyzers = new Map<ParserDialects, QueryContextAnalyzer>();
+  private readonly analysisParser: NodeSqlParser;
+
+  constructor(analysisParser: NodeSqlParser) {
+    super(analysisParser);
+    this.analysisParser = analysisParser;
+  }
+
+  override getContext(
+    sql: string,
+    opts: { state: EditorState },
+  ): ReturnType<QueryContextAnalyzer["getContext"]> {
+    return this.getAnalyzer(opts.state).getContext(sql, opts);
+  }
+
+  override clearCache(): void {
+    for (const analyzer of this.analyzers.values()) {
+      analyzer.clearCache();
+    }
+  }
+
+  private getAnalyzer(state: EditorState): QueryContextAnalyzer {
+    const dialect = getAnalysisDialect(state);
+    let analyzer = this.analyzers.get(dialect);
+    if (!analyzer) {
+      analyzer = new QueryContextAnalyzer(this.analysisParser);
+      this.analyzers.set(dialect, analyzer);
+    }
+    return analyzer;
+  }
+}
+
+class DialectAwareSqlStructureAnalyzer extends SqlStructureAnalyzer {
+  private readonly analyzers = new Map<ParserDialects, SqlStructureAnalyzer>();
+  private readonly analysisParser: NodeSqlParser;
+
+  constructor(analysisParser: NodeSqlParser) {
+    super(analysisParser);
+    this.analysisParser = analysisParser;
+  }
+
+  override analyzeDocument(
+    state: EditorState,
+  ): ReturnType<SqlStructureAnalyzer["analyzeDocument"]> {
+    return this.getAnalyzer(state).analyzeDocument(state);
+  }
+
+  override getStatementAtPosition(
+    state: EditorState,
+    position: number,
+  ): ReturnType<SqlStructureAnalyzer["getStatementAtPosition"]> {
+    return this.getAnalyzer(state).getStatementAtPosition(state, position);
+  }
+
+  override getStatementsInRange(
+    state: EditorState,
+    from: number,
+    to: number,
+  ): ReturnType<SqlStructureAnalyzer["getStatementsInRange"]> {
+    return this.getAnalyzer(state).getStatementsInRange(state, from, to);
+  }
+
+  override clearCache(): void {
+    for (const analyzer of this.analyzers.values()) {
+      analyzer.clearCache();
+    }
+  }
+
+  private getAnalyzer(state: EditorState): SqlStructureAnalyzer {
+    const dialect = getAnalysisDialect(state);
+    let analyzer = this.analyzers.get(dialect);
+    if (!analyzer) {
+      analyzer = new SqlStructureAnalyzer(this.analysisParser);
+      this.analyzers.set(dialect, analyzer);
+    }
+    return analyzer;
+  }
+}
+
 class CustomSqlParser extends NodeSqlParser {
   private validationTimeout: number | null = null;
+  private pendingValidationResolve: ((errors: SqlParseError[]) => void) | null =
+    null;
   private readonly VALIDATION_DELAY_MS = 300; // Wait 300ms after user stops typing
   private isFocused = false; // Only validate if the editor is focused
 
@@ -232,37 +381,51 @@ class CustomSqlParser extends NodeSqlParser {
     this.isFocused = focused;
   }
 
+  private resolvePendingValidation(errors: SqlParseError[]): void {
+    this.pendingValidationResolve?.(errors);
+    this.pendingValidationResolve = null;
+  }
+
   private async validateWithDelay(
     sql: string,
-    engine: string,
+    engine: ConnectionName,
     dialect: ParserDialects | null,
   ): Promise<SqlParseError[]> {
-    // Clear any existing delay call
+    // Clear any existing delay call, resolving its promise so a superseded
+    // request doesn't hang forever.
     if (this.validationTimeout) {
       window.clearTimeout(this.validationTimeout);
+      this.resolvePendingValidation([]);
     }
 
     // Set up a new request to be called after the delay
     return new Promise((resolve) => {
+      this.pendingValidationResolve = resolve;
       this.validationTimeout = window.setTimeout(async () => {
+        this.validationTimeout = null;
+
         // Only validate if the editor is still focused
         if (!this.isFocused) {
-          resolve([]);
+          this.resolvePendingValidation([]);
           return;
         }
 
         try {
-          const sqlMode = getSQLMode();
+          // For validate mode, we run EXPLAIN queries on the engine, which can be
+          // expensive for remote databases. So, we only run for internal engines.
+          const sqlMode = INTERNAL_SQL_ENGINES.has(engine)
+            ? getSQLMode()
+            : "default";
           const result = await validateSQL(sql, engine, dialect, sqlMode);
           if (result.error) {
             Logger.error("Failed to validate SQL", { error: result.error });
-            resolve([]);
+            this.resolvePendingValidation([]);
             return;
           }
-          resolve(result.parse_result?.errors ?? []);
+          this.resolvePendingValidation(result.parse_result?.errors ?? []);
         } catch (error) {
           Logger.error("Failed to validate SQL", { error });
-          resolve([]);
+          this.resolvePendingValidation([]);
         }
       }, this.VALIDATION_DELAY_MS);
     });
@@ -273,19 +436,23 @@ class CustomSqlParser extends NodeSqlParser {
     opts: { state: EditorState },
   ): Promise<SqlParseError[]> {
     const metadata = getSQLMetadata(opts.state);
+    const dialect = connectionNameToParserDialect(metadata.engine);
 
     // Only validate if the editor is focused
     if (!this.isFocused) {
       return [];
     }
 
-    // Only perform custom validation for DuckDB
-    if (!INTERNAL_SQL_ENGINES.has(metadata.engine)) {
+    // Only perform custom validation for DuckDB as we have a custom validation endpoint for it.
+    if (!isDuckDBConnection(metadata.engine, dialect)) {
       return super.validateSql(sql, opts);
     }
 
-    const dialect = guessParserDialect(opts.state);
-    return this.validateWithDelay(sql, metadata.engine, dialect);
+    return this.validateWithDelay(
+      sql,
+      metadata.engine,
+      dialect ?? DEFAULT_PARSER_DIALECT,
+    );
   }
 
   override async parse(
@@ -294,9 +461,10 @@ class CustomSqlParser extends NodeSqlParser {
   ): Promise<NodeSqlParserResult> {
     const metadata = getSQLMetadata(opts.state);
     const engine = metadata.engine;
+    const dialect = connectionNameToParserDialect(engine);
 
     // For now, always return success for DuckDB
-    if (engine === DUCKDB_ENGINE) {
+    if (isDuckDBConnection(engine, dialect)) {
       return { success: true, errors: [] };
     }
 
@@ -304,6 +472,13 @@ class CustomSqlParser extends NodeSqlParser {
   }
 }
 
+function isDuckDBConnection(
+  engine: ConnectionName,
+  dialect: ParserDialects | null,
+): boolean {
+  return engine === DUCKDB_ENGINE || dialect === "DuckDB";
+}
+
 /**
  * Update the SQL dialect in the editor view.
  */
@@ -344,7 +519,61 @@ function getSchema(view: EditorView): SQLNamespace {
     return {};
   }
 
-  return config.schema;
+  const schema = config.schema;
+  if (!isNamespaceMap(schema)) {
+    return schema;
+  }
+
+  const defaultSchema = config.defaultSchema;
+  const defaultSchemaNamespace = defaultSchema
+    ? schema[defaultSchema]
+    : undefined;
+
+  // The completion schema contains both default-schema and fully qualified
+  // paths. Promote default-schema tables so unqualified references resolve to
+  // one exact match instead of appearing ambiguous.
+  if (
+    defaultSchemaNamespace &&
+    !Array.isArray(defaultSchemaNamespace) &&
+    typeof defaultSchemaNamespace === "object" &&
+    "children" in defaultSchemaNamespace &&
+    isNamespaceMap(defaultSchemaNamespace.children)
+  ) {
+    return {
+      ...schema,
+      ...Object.fromEntries(
+        Object.entries(defaultSchemaNamespace.children).filter(
+          ([name]) => !Object.hasOwn(schema, name),
+        ),
+      ),
+    };
+  }
+
+  return schema;
+}
+
+function isNamespaceMap(
+  namespace: SQLNamespace,
+): namespace is Record<string, SQLNamespace> {
+  if (Array.isArray(namespace)) {
+    return false;
+  }
+
+  if (!("self" in namespace) || !("children" in namespace)) {
+    return true;
+  }
+
+  // A self/children wrapper stores a Completion in `self`. A table or schema
+  // literally named "self" stores another namespace there.
+  const self = namespace.self;
+  const isSelfChildrenWrapper =
+    typeof self === "object" &&
+    self !== null &&
+    !Array.isArray(self) &&
+    "label" in self &&
+    typeof self.label === "string";
+
+  return !isSelfChildrenWrapper;
 }
 
 function guessParserDialect(state: EditorState): ParserDialects | null {
@@ -529,3 +758,9 @@ async function validateSQL(
   }
   return result;
 }
+
+export const exportedForTesting = {
+  createSQLAnalysis,
+  CustomSqlParser,
+  getSchema,
+};

package/src/core/codemirror/lsp/__tests__/federated-lsp.test.ts

@@ -0,0 +1,161 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { describe, expect, it, type Mocked, vi } from "vitest";
+import type * as LSP from "vscode-languageserver-protocol";
+import { FederatedLanguageServerClient } from "../federated-lsp";
+import type { ILanguageServerClient } from "../types";
+
+function createClient(
+  supportedMethods: string[] = [],
+): Mocked<ILanguageServerClient> {
+  const client: Mocked<ILanguageServerClient> = {
+    ready: true,
+    capabilities: {},
+    initializePromise: Promise.resolve(),
+    clientCapabilities: {},
+    initialize: vi.fn(),
+    close: vi.fn(),
+    hasCapability: vi.fn((method) => supportedMethods.includes(method)),
+    textDocumentDidOpen: vi.fn().mockResolvedValue(false),
+    textDocumentDidChange: vi.fn(),
+    textDocumentDidClose: vi.fn(),
+    textDocumentWillSave: vi.fn(),
+    textDocumentWillSaveWaitUntil: vi.fn().mockResolvedValue(null),
+    textDocumentDidSave: vi.fn(),
+    textDocumentHover: vi.fn().mockResolvedValue({ contents: [] }),
+    textDocumentCompletion: vi.fn().mockResolvedValue(null),
+    completionItemResolve: vi.fn(),
+    textDocumentDefinition: vi.fn().mockResolvedValue(null),
+    textDocumentCodeAction: vi.fn().mockResolvedValue(null),
+    codeActionResolve: vi.fn(),
+    textDocumentRename: vi.fn().mockResolvedValue(null),
+    textDocumentPrepareRename: vi.fn().mockResolvedValue(null),
+    textDocumentSignatureHelp: vi.fn().mockResolvedValue(null),
+    onNotification: vi.fn().mockReturnValue(() => true),
+  };
+  return client;
+}
+
+describe("FederatedLanguageServerClient", () => {
+  it("routes requests using dynamic method capabilities", async () => {
+    const staticOnlyClient = createClient();
+    const dynamicClient = createClient(["textDocument/definition"]);
+    const definition: LSP.Location = {
+      uri: "file:///definition.py",
+      range: {
+        start: { line: 1, character: 2 },
+        end: { line: 1, character: 5 },
+      },
+    };
+    dynamicClient.textDocumentDefinition.mockResolvedValue(definition);
+    const client = new FederatedLanguageServerClient([
+      staticOnlyClient,
+      dynamicClient,
+    ]);
+    const params: LSP.DefinitionParams = {
+      textDocument: { uri: "file:///notebook.py" },
+      position: { line: 0, character: 0 },
+    };
+
+    await expect(client.textDocumentDefinition(params)).resolves.toEqual(
+      definition,
+    );
+    expect(client.hasCapability("textDocument/definition")).toBe(true);
+    expect(staticOnlyClient.textDocumentDefinition).not.toHaveBeenCalled();
+    expect(dynamicClient.textDocumentDefinition).toHaveBeenCalledWith(params);
+  });
+
+  it("prefers the first client's willSaveWaitUntil edits regardless of latency", async () => {
+    const slowEdits: LSP.TextEdit[] = [
+      {
+        range: {
+          start: { line: 0, character: 0 },
+          end: { line: 0, character: 1 },
+        },
+        newText: "first",
+      },
+    ];
+    const first = createClient();
+    const second = createClient();
+    // Answers last, but still wins
+    first.textDocumentWillSaveWaitUntil.mockImplementation(
+      () => new Promise((resolve) => setTimeout(() => resolve(slowEdits), 5)),
+    );
+    second.textDocumentWillSaveWaitUntil.mockResolvedValue([
+      {
+        range: {
+          start: { line: 1, character: 0 },
+          end: { line: 1, character: 1 },
+        },
+        newText: "second",
+      },
+    ]);
+    const client = new FederatedLanguageServerClient([first, second]);
+
+    await expect(
+      client.textDocumentWillSaveWaitUntil({
+        textDocument: { uri: "file:///notebook.py" },
+        reason: 1,
+      }),
+    ).resolves.toEqual(slowEdits);
+    expect(second.textDocumentWillSaveWaitUntil).toHaveBeenCalled();
+  });
+
+  it("skips clients that reject willSaveWaitUntil", async () => {
+    const failing = createClient();
+    const working = createClient();
+    const edits: LSP.TextEdit[] = [
+      {
+        range: {
+          start: { line: 0, character: 0 },
+          end: { line: 0, character: 1 },
+        },
+        newText: "ok",
+      },
+    ];
+    failing.textDocumentWillSaveWaitUntil.mockRejectedValue(
+      new Error("server exploded"),
+    );
+    working.textDocumentWillSaveWaitUntil.mockResolvedValue(edits);
+    const client = new FederatedLanguageServerClient([failing, working]);
+
+    await expect(
+      client.textDocumentWillSaveWaitUntil({
+        textDocument: { uri: "file:///notebook.py" },
+        reason: 1,
+      }),
+    ).resolves.toEqual(edits);
+  });
+
+  it("fans out document lifecycle and save notifications", async () => {
+    const first = createClient();
+    const second = createClient();
+    second.textDocumentDidOpen.mockResolvedValue(true);
+    const client = new FederatedLanguageServerClient([first, second]);
+    const openParams: LSP.DidOpenTextDocumentParams = {
+      textDocument: {
+        uri: "file:///notebook.py",
+        languageId: "python",
+        version: 1,
+        text: "value = 1",
+      },
+    };
+    const closeParams: LSP.DidCloseTextDocumentParams = {
+      textDocument: { uri: "file:///notebook.py" },
+    };
+    const saveParams: LSP.DidSaveTextDocumentParams = {
+      textDocument: { uri: "file:///notebook.py" },
+      text: "value = 1",
+    };
+
+    await expect(client.textDocumentDidOpen(openParams)).resolves.toBe(true);
+    await client.textDocumentDidClose(closeParams);
+    await client.textDocumentDidSave(saveParams);
+
+    for (const child of [first, second]) {
+      expect(child.textDocumentDidOpen).toHaveBeenCalledWith(openParams);
+      expect(child.textDocumentDidClose).toHaveBeenCalledWith(closeParams);
+      expect(child.textDocumentDidSave).toHaveBeenCalledWith(saveParams);
+    }
+  });
+});