@marimo-team/islands 0.16.5 → 0.17.1

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

@@ -6,8 +6,6 @@ import { type SQLDialect, type SQLNamespace, sql } from "@codemirror/lang-sql";
 import type { EditorState, Extension } from "@codemirror/state";
 import { Compartment } from "@codemirror/state";
 import { EditorView, keymap } from "@codemirror/view";
-import type { SyntaxNode, TreeCursor } from "@lezer/common";
-import { parser } from "@lezer/python";
 import {
   defaultSqlHoverTheme,
   NodeSqlParser,
@@ -17,9 +15,15 @@ import {
   sqlExtension,
 } from "@marimo-team/codemirror-sql";
 import { DuckDBDialect } from "@marimo-team/codemirror-sql/dialects";
-import dedent from "string-dedent";
+import { type SQLMetadata, SQLParser } from "@marimo-team/smart-cells";
+import type { CellId } from "@/core/cells/ids";
 import { cellIdState } from "@/core/codemirror/cells/state";
-import { getFeatureFlag } from "@/core/config/feature-flag";
+import type { PlaceholderType } from "@/core/codemirror/config/types";
+import type {
+  CompletionConfig,
+  DiagnosticsConfig,
+  LSPConfig,
+} from "@/core/config/config-schema";
 import {
   dataSourceConnectionsAtom,
   setLatestEngineSelected,
@@ -30,6 +34,7 @@ import {
   INTERNAL_SQL_ENGINES,
 } from "@/core/datasets/engines";
 import { ValidateSQL } from "@/core/datasets/request-registry";
+import type { HotkeyProvider } from "@/core/hotkeys/hotkeys";
 import type { ValidateSQLResult } from "@/core/kernel/messages";
 import { store } from "@/core/state/jotai";
 import { resolvedThemeAtom } from "@/theme/useTheme";
@@ -37,10 +42,6 @@ import { Logger } from "@/utils/Logger";
 import { variableCompletionSource } from "../../embedded/embedded-python";
 import { languageMetadataField } from "../../metadata";
 import type { LanguageAdapter } from "../../types";
-import { parseArgsKwargs } from "../../utils/ast";
-import { indentOneTab } from "../../utils/indentOneTab";
-import type { QuotePrefixKind } from "../../utils/quotes";
-import { MarkdownLanguageAdapter } from "../markdown";
 import {
   clearSqlValidationError,
   setSqlValidationError,
@@ -58,11 +59,7 @@ const DEFAULT_PARSER_DIALECT = "DuckDB";
 // A compartment for the SQL config, so we can update the config of codemirror
 const sqlConfigCompartment = new Compartment();
 
-export interface SQLLanguageAdapterMetadata {
-  dataframeName: string;
-  quotePrefix: QuotePrefixKind;
-  commentLines: string[];
-  showOutput: boolean;
+export interface SQLLanguageAdapterMetadata extends SQLMetadata {
   engine: ConnectionName;
 }
 
@@ -76,41 +73,30 @@ function getLatestEngine(): ConnectionName {
 export class SQLLanguageAdapter
   implements LanguageAdapter<SQLLanguageAdapterMetadata>
 {
+  private parser = new SQLParser();
   readonly type = "sql";
-  sqlLinterEnabled: boolean;
   sqlModeEnabled: boolean;
 
   constructor() {
-    try {
-      this.sqlLinterEnabled = getFeatureFlag("sql_linter");
-      this.sqlModeEnabled = getFeatureFlag("sql_mode");
-    } catch {
-      this.sqlLinterEnabled = false;
-      this.sqlModeEnabled = false;
-    }
+    this.sqlModeEnabled = true;
   }
 
   get defaultMetadata(): SQLLanguageAdapterMetadata {
     return {
-      dataframeName: "_df",
-      quotePrefix: "f",
-      commentLines: [],
-      showOutput: true,
-      engine: getLatestEngine() || this.defaultEngine,
+      ...this.parser.defaultMetadata,
+      engine: getLatestEngine() || DUCKDB_ENGINE,
     };
   }
 
   get defaultCode(): string {
-    const latestEngine = getLatestEngine();
-    if (latestEngine === this.defaultEngine) {
-      return `_df = mo.sql(f"""SELECT * FROM """)`;
+    const engine = getLatestEngine();
+    if (engine && engine !== DUCKDB_ENGINE) {
+      return `_df = mo.sql(f"""SELECT * FROM """, engine=${engine})`;
     }
-    return `_df = mo.sql(f"""SELECT * FROM """, engine=${latestEngine})`;
+    return this.parser.defaultCode;
   }
 
-  static fromQuery = (query: string) => `_df = mo.sql(f"""${query.trim()}""")`;
-
-  private readonly defaultEngine = DUCKDB_ENGINE;
+  static fromQuery = (query: string) => SQLParser.fromQuery(query);
 
   transformIn(
     pythonCode: string,
@@ -119,103 +105,38 @@ export class SQLLanguageAdapter
     queryStartOffset: number,
     metadata: SQLLanguageAdapterMetadata,
   ] {
-    pythonCode = pythonCode.trim();
-
-    // Default metadata
-    const metadata: SQLLanguageAdapterMetadata = {
-      ...this.defaultMetadata,
-      commentLines: this.extractCommentLines(pythonCode),
-    };
+    this.parser.defaultMetadata.engine = getLatestEngine() || DUCKDB_ENGINE;
+    const result = this.parser.transformIn(pythonCode);
 
-    if (!this.isSupported(pythonCode)) {
-      // Attempt to remove any markdown wrappers
-      const [transformedCode, offset] =
-        new MarkdownLanguageAdapter().transformIn(pythonCode);
-      // Just return the original code
-      return [transformedCode, offset, metadata];
-    }
+    // Handle engine selection side effect
+    const metadata = result.metadata as SQLLanguageAdapterMetadata;
 
-    // Handle empty strings
-    if (pythonCode === "") {
-      return ["", 0, metadata];
+    if (metadata.engine && metadata.engine !== DUCKDB_ENGINE) {
+      setLatestEngineSelected(metadata.engine);
     }
 
-    const sqlStatement = parseSQLStatement(pythonCode);
-    if (sqlStatement) {
-      metadata.dataframeName = sqlStatement.dfName;
-      metadata.showOutput = sqlStatement.output ?? true;
-      metadata.engine =
-        (sqlStatement.engine as ConnectionName) ?? this.defaultEngine;
-
-      if (metadata.engine !== this.defaultEngine) {
-        // User selected a new engine, set it as latest.
-        // This makes new SQL statements use the new engine by default.
-        setLatestEngineSelected(metadata.engine);
-      }
-
-      return [
-        dedent(`\n${sqlStatement.sqlString}\n`).trim(),
-        sqlStatement.startPosition,
-        metadata,
-      ];
-    }
-
-    return [pythonCode, 0, metadata];
+    return [result.code, result.offset, metadata];
   }
 
   transformOut(
     code: string,
     metadata: SQLLanguageAdapterMetadata,
   ): [string, number] {
-    const { quotePrefix, commentLines, showOutput, engine, dataframeName } =
-      metadata;
-
-    const start = `${dataframeName} = mo.sql(\n    ${quotePrefix}"""\n`;
-    const escapedCode = code.replaceAll('"""', String.raw`\"""`);
-
-    const showOutputParam = showOutput ? "" : ",\n    output=False";
-    const engineParam =
-      engine === this.defaultEngine ? "" : `,\n    engine=${engine}`;
-    const end = `\n    """${showOutputParam}${engineParam}\n)`;
-
-    return [
-      [...commentLines, start].join("\n") + indentOneTab(escapedCode) + end,
-      start.length + 1,
-    ];
+    const result = this.parser.transformOut(code, metadata);
+    return [result.code, result.offset];
   }
 
   isSupported(pythonCode: string): boolean {
-    if (pythonCode.trim() === "") {
-      return true;
-    }
-
-    // Has at least one `mo.sql` call
-    if (!pythonCode.includes("mo.sql")) {
-      return false;
-    }
-
-    // Does not have 2 `mo.sql` calls
-    if (pythonCode.split("mo.sql").length > 2) {
-      return false;
-    }
-
-    return parseSQLStatement(pythonCode) !== null;
+    return this.parser.isSupported(pythonCode);
   }
 
-  private extractCommentLines(pythonCode: string): string[] {
-    const lines = pythonCode.split("\n");
-    const commentLines = [];
-    for (const line of lines) {
-      if (line.startsWith("#")) {
-        commentLines.push(line);
-      } else {
-        break;
-      }
-    }
-    return commentLines;
-  }
-
-  getExtension(): Extension[] {
+  getExtension(
+    _cellId: CellId,
+    _completionConfig: CompletionConfig,
+    _hotkeys: HotkeyProvider,
+    _placeholderType: PlaceholderType,
+    lspConfig: LSPConfig & { diagnostics: DiagnosticsConfig },
+  ): Extension[] {
     const extensions = [
       // This can be updated with a dispatch effect
       sqlConfigCompartment.of(sql({ dialect: DEFAULT_DIALECT })),
@@ -245,7 +166,9 @@ export class SQLLanguageAdapter
       }),
     ];
 
-    if (this.sqlLinterEnabled) {
+    const sqlLinterEnabled = lspConfig?.diagnostics?.sql_linter ?? false;
+
+    if (sqlLinterEnabled) {
       const theme = store.get(resolvedThemeAtom);
       const parser = new CustomSqlParser({
         getParserOptions: (state: EditorState) => {
@@ -468,194 +391,6 @@ function connectionNameToParserDialect(
   }
 }
 
-interface SQLParseInfo {
-  dfName: string;
-  sqlString: string;
-  engine: string | undefined;
-  output: boolean | undefined;
-  startPosition: number;
-}
-
-// Finds an assignment node that is preceded only by comments.
-function findAssignment(cursor: TreeCursor): SyntaxNode | null {
-  do {
-    if (cursor.name === "AssignStatement") {
-      return cursor.node;
-    }
-
-    if (cursor.name !== "Comment") {
-      return null;
-    }
-  } while (cursor.next());
-  return null;
-}
-
-function getStringContent(node: SyntaxNode, code: string): string | null {
-  // Handle triple quoted strings
-  if (node.name === "String") {
-    const content = code.slice(node.from, node.to);
-    // Remove quotes and trim
-    if (content.startsWith('"""') || content.startsWith("'''")) {
-      return safeDedent(content.slice(3, -3));
-    }
-    // Handle single quoted strings
-    return safeDedent(content.slice(1, -1));
-  }
-  // Handle f-strings
-  if (node.name === "FormatString") {
-    const content = code.slice(node.from, node.to);
-    if (content.startsWith('f"""') || content.startsWith("f'''")) {
-      return safeDedent(content.slice(4, -3));
-    }
-    return safeDedent(content.slice(2, -1));
-  }
-  return null;
-}
-
-/**
- * Parses a SQL statement from a Python code string.
- *
- * @param code - The Python code string to parse.
- * @returns The parsed SQL statement or null if parsing fails.
- */
-function parseSQLStatement(code: string): SQLParseInfo | null {
-  try {
-    const tree = parser.parse(code);
-    const cursor = tree.cursor();
-
-    // Trees start with a Script node.
-    if (cursor.name === "Script") {
-      cursor.next();
-    }
-    const assignStmt = findAssignment(cursor);
-    if (!assignStmt) {
-      return null;
-    }
-
-    // Code after the assignment statement is not allowed.
-    if (code.slice(assignStmt.to).trim().length > 0) {
-      return null;
-    }
-
-    let dfName: string | null = null;
-    let sqlString: string | null = null;
-    let engine: string | undefined;
-    let output: boolean | undefined;
-    let startPosition = 0;
-
-    // Parse the assignment
-    const assignCursor = assignStmt.cursor();
-    assignCursor.firstChild(); // Move to first child of assignment
-
-    // First child should be the variable name
-    if (assignCursor.name === "VariableName") {
-      dfName = code.slice(assignCursor.from, assignCursor.to);
-    }
-
-    if (!dfName) {
-      return null;
-    }
-
-    // Move to the expression part (after the =)
-    while (assignCursor.next()) {
-      if (assignCursor.name === "CallExpression") {
-        // Check if it's mo.sql call
-        const callCursor = assignCursor.node.cursor();
-        let isMoSql = false;
-
-        callCursor.firstChild(); // Move to first child of call
-        if (callCursor.name === "MemberExpression") {
-          const memberText = code.slice(callCursor.from, callCursor.to);
-          isMoSql = memberText === "mo.sql";
-        }
-
-        if (!isMoSql) {
-          return null;
-        }
-
-        // Move to arguments
-        while (callCursor.next()) {
-          if (callCursor.name === "ArgList") {
-            const argListCursor = callCursor.node.cursor();
-
-            const { args, kwargs } = parseArgsKwargs(argListCursor, code);
-
-            // Parse positional args (SQL query)
-            if (args.length === 1) {
-              sqlString = getStringContent(args[0], code);
-              startPosition =
-                args[0].from +
-                getPrefixLength(code.slice(args[0].from, args[0].to));
-            }
-
-            // Parse kwargs (engine and output)
-            for (const { key, value } of kwargs) {
-              switch (key) {
-                case "engine":
-                  engine = value;
-                  break;
-                case "output":
-                  output = value === "True";
-                  break;
-              }
-            }
-
-            // Check if sql string is empty
-            if (sqlString === "") {
-              return { dfName, sqlString: "", engine, output, startPosition };
-            }
-
-            break;
-          }
-        }
-      }
-    }
-
-    if (!dfName || !sqlString) {
-      return null;
-    }
-
-    return {
-      dfName,
-      sqlString,
-      engine,
-      output,
-      startPosition,
-    };
-  } catch (error) {
-    Logger.warn("Failed to parse SQL statement", { error: error });
-    return null;
-  }
-}
-
-function getPrefixLength(code: string): number {
-  if (code === "") {
-    return 0;
-  }
-  if (code.startsWith('f"""') || code.startsWith("f'''")) {
-    return 4;
-  }
-  if (code.startsWith('"""') || code.startsWith("'''")) {
-    return 3;
-  }
-  if (code.startsWith("f'") || code.startsWith('f"')) {
-    return 2;
-  }
-  if (code.startsWith("'") || code.startsWith('"')) {
-    return 1;
-  }
-  return 0;
-}
-
-function safeDedent(code: string): string {
-  try {
-    // Dedent expects the first and last line to be empty / contain only whitespace, so we pad with \n
-    return dedent(`\n${code}\n`).trim();
-  } catch {
-    return code;
-  }
-}
-
 const SQL_VALIDATION_DEBOUNCE_MS = 300;
 
 /**
@@ -666,8 +401,10 @@ function sqlValidationExtension(): Extension {
   let lastValidationRequest: string | null = null;
 
   return EditorView.updateListener.of((update) => {
-    // Only run validation if the document has changed and editor is focused
-    if (!update.docChanged || !update.view.hasFocus) {
+    // Only run validation if the document has changed
+    // The extension only runs on keypress, so we don't need to check for focus
+    // This lets AI completions / external calls trigger validation
+    if (!update.docChanged) {
       return;
     }
 

package/src/core/codemirror/language/panel/markdown.tsx

@@ -1,5 +1,6 @@
 /* Copyright 2024 Marimo. All rights reserved. */
-import type { QuotePrefixKind } from "../utils/quotes";
+
+import type { QuotePrefixKind } from "@marimo-team/smart-cells";
 
 // Based on the current quote prefix and the checkbox state, return the new quote prefix
 export function getQuotePrefix({

package/src/core/codemirror/language/panel/panel.tsx

@@ -1,11 +1,11 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 import type { EditorView } from "@codemirror/view";
+import type { QuotePrefixKind } from "@marimo-team/smart-cells";
 import { InfoIcon, PaintRollerIcon } from "lucide-react";
 import { Button } from "@/components/ui/button";
 import { Checkbox } from "@/components/ui/checkbox";
 import { Tooltip, TooltipProvider } from "@/components/ui/tooltip";
 import { normalizeName } from "@/core/cells/names";
-import { getFeatureFlag } from "@/core/config/feature-flag";
 import { type ConnectionName, DUCKDB_ENGINE } from "@/core/datasets/engines";
 import { useAutoGrowInputProps } from "@/hooks/useAutoGrowInputProps";
 import { formatSQL } from "../../format";
@@ -21,7 +21,6 @@ import {
   updateLanguageMetadata,
 } from "../metadata";
 import type { LanguageMetadataOf } from "../types";
-import type { QuotePrefixKind } from "../utils/quotes";
 import { getQuotePrefix, MarkdownQuotePrefixTooltip } from "./markdown";
 import { SQLEngineSelect, SQLModeSelect } from "./sql";
 
@@ -71,8 +70,6 @@ export const LanguagePanelComponent: React.FC<{
       updateSQLDialectFromConnection(view, engine);
     };
 
-    const sqlModeEnabled = getFeatureFlag("sql_mode");
-
     actions = (
       <div className="flex flex-1 gap-2 items-center">
         <label className="flex gap-2 items-center">
@@ -98,9 +95,7 @@ export const LanguagePanelComponent: React.FC<{
           onChange={switchEngine}
         />
         <div className="flex items-center gap-2 ml-auto">
-          {sqlModeEnabled && metadata.engine === DUCKDB_ENGINE && (
-            <SQLModeSelect />
-          )}
+          {metadata.engine === DUCKDB_ENGINE && <SQLModeSelect />}
           <Tooltip content="Format SQL">
             <Button
               variant="text"

package/src/core/codemirror/language/utils.ts

@@ -2,6 +2,7 @@
 
 import type { EditorState } from "@codemirror/state";
 import type { EditorView } from "@codemirror/view";
+import { replaceEditorContent } from "../replace-editor-content";
 import { languageAdapterState } from "./extension";
 import { languageMetadataField } from "./metadata";
 
@@ -36,14 +37,8 @@ export function updateEditorCodeFromPython(
 ): string {
   const languageAdapter = editor.state.field(languageAdapterState);
   const [code] = languageAdapter.transformIn(pythonCode);
-  const doc = editor.state.doc;
-  // Noop if the code is the same
-  if (doc.toString() === code) {
-    return code;
-  }
-  editor.dispatch({
-    changes: { from: 0, to: doc.length, insert: code },
-  });
+  // Use replaceEditorContent which preserves cursor position when focused
+  replaceEditorContent(editor, code);
   return code;
 }
 

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

@@ -195,35 +195,36 @@ describe("createNotebookLens", () => {
 });
 
 describe("NotebookLanguageServerClient", () => {
-  let plugins: any[] = [];
   let mockClient: Mocked<ILanguageServerClient>;
   let notebookClient: NotebookLanguageServerClient;
 
   beforeEach(() => {
-    plugins = [];
     mockClient = {
-      plugins,
       ready: true,
       capabilities: {},
       initializePromise: Promise.resolve(),
+      clientCapabilities: {},
+      completionItemResolve: vi.fn(),
       initialize: vi.fn(),
       close: vi.fn(),
-      detachPlugin: vi.fn().mockImplementation((plugin) => {
-        plugins = plugins.filter((p) => p !== plugin);
-      }),
-      attachPlugin: vi.fn().mockImplementation((plugin) => {
-        plugins.push(plugin);
-      }),
+      onNotification: vi.fn(),
       textDocumentDidOpen: vi.fn(),
       textDocumentDidChange: vi.fn(),
       textDocumentHover: vi.fn(),
       textDocumentCompletion: vi.fn(),
+      textDocumentDefinition: vi.fn(),
+      textDocumentPrepareRename: vi.fn(),
+      textDocumentCodeAction: vi.fn(),
+      textDocumentSignatureHelp: vi.fn(),
       textDocumentRename: vi.fn(),
-      processNotification: vi.fn(),
-      notify: vi.fn(),
-      request: vi.fn(),
-    } as unknown as Mocked<ILanguageServerClient>;
-    notebookClient = new NotebookLanguageServerClient(mockClient, {});
+    };
+    (mockClient as any).processNotification = vi.fn();
+    (mockClient as any).notify = vi.fn();
+    notebookClient = new NotebookLanguageServerClient(mockClient, {}, () => ({
+      [Cells.cell1]: new EditorView({ doc: "# this is a comment" }),
+      [Cells.cell2]: new EditorView({ doc: "import math\nimport numpy" }),
+      [Cells.cell3]: new EditorView({ doc: "print(math.sqrt(4))" }),
+    }));
 
     // Mock the atom instead of the instance method
     vi.spyOn(store, "get").mockImplementation((atom) => {
@@ -404,6 +405,12 @@ describe("NotebookLanguageServerClient", () => {
       });
       expect(mockView3.state.doc.toString()).toBe("print(math.sqrt(4))");
 
+      (notebookClient as any).getNotebookEditors = () => ({
+        [Cells.cell1]: mockView1,
+        [Cells.cell2]: mockView2,
+        [Cells.cell3]: mockView3,
+      });
+
       // Setup rename params
       const renameParams: LSP.RenameParams = {
         textDocument: { uri: CellDocumentUri.of(Cells.cell2) },

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

@@ -1,6 +1,5 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
-import type { LanguageServerPlugin } from "@marimo-team/codemirror-languageserver";
 import type * as LSP from "vscode-languageserver-protocol";
 import { Objects } from "@/utils/objects";
 import type { ILanguageServerClient } from "./types";
@@ -24,6 +23,26 @@ export class FederatedLanguageServerClient implements ILanguageServerClient {
     this.documentUri = getLSPDocument();
   }
 
+  onNotification(
+    listener: (n: {
+      jsonrpc: "2.0";
+      id?: null | undefined;
+      method: "textDocument/publishDiagnostics";
+      params: LSP.PublishDiagnosticsParams;
+    }) => void,
+  ): () => boolean {
+    const callbacks: (() => boolean)[] = [];
+    for (const client of this.clients) {
+      callbacks.push(client.onNotification(listener));
+    }
+    return () => {
+      for (const cb of callbacks) {
+        cb();
+      }
+      return true;
+    };
+  }
+
   get clientCapabilities(): LSP.ClientCapabilities | undefined {
     const capabilities = this.clients
       .map((client) => {
@@ -153,14 +172,6 @@ export class FederatedLanguageServerClient implements ILanguageServerClient {
     return null;
   }
 
-  attachPlugin(plugin: LanguageServerPlugin): void {
-    this.clients.forEach((client) => client.attachPlugin(plugin));
-  }
-
-  detachPlugin(plugin: LanguageServerPlugin): void {
-    this.clients.forEach((client) => client.detachPlugin(plugin));
-  }
-
   // Merge completions from all clients
   async textDocumentCompletion(
     params: LSP.CompletionParams,