@marimo-team/frontend 0.16.0 → 0.16.2

package/src/core/cells/__tests__/add-missing-import.test.ts

@@ -1,18 +1,35 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
 import { createStore } from "jotai";
-import { describe, expect, it, vi } from "vitest";
+import { beforeEach, describe, expect, it, vi } from "vitest";
+import { MockNotebook } from "@/__mocks__/notebook";
+import { MockRequestClient } from "@/__mocks__/requests";
+import { store } from "@/core/state/jotai";
 import { variablesAtom } from "@/core/variables/state";
 import type { Variables } from "@/core/variables/types";
-import { CollapsibleTree, MultiColumn } from "@/utils/id-tree";
-import { maybeAddMissingImport } from "../add-missing-import";
-import { type NotebookState, notebookAtom } from "../cells";
+import {
+  maybeAddAltairImport,
+  maybeAddMarimoImport,
+  maybeAddMissingImport,
+} from "../add-missing-import";
+import { notebookAtom } from "../cells";
 import type { CellId } from "../ids";
-import type { CellData } from "../types";
 
-const CELL_IDS = new MultiColumn([CollapsibleTree.from(["1" as CellId])]);
+// Mock the getRequestClient function
+const mockRequestClient = MockRequestClient.create();
+vi.mock("@/core/network/requests", () => ({
+  getRequestClient: () => mockRequestClient,
+}));
+
+const Cell1 = "1" as CellId;
+const Cell2 = "2" as CellId;
 
 describe("maybeAddMissingImport", () => {
+  beforeEach(() => {
+    store.set(variablesAtom, {} as Variables);
+    store.set(notebookAtom, MockNotebook.notebookState({ cellData: {} }));
+  });
+
   it("should not add an import if the variable is already in the variables state", () => {
     const appStore = createStore();
     appStore.set(variablesAtom, { mo: {} } as Variables);
@@ -36,14 +53,14 @@ describe("maybeAddMissingImport", () => {
     (code) => {
       const appStore = createStore();
       appStore.set(variablesAtom, {} as Variables);
-      appStore.set(notebookAtom, {
-        cellData: {
-          ["1" as CellId]: {
-            code: code,
-          } as CellData,
-        },
-        cellIds: CELL_IDS,
-      } as NotebookState);
+      appStore.set(
+        notebookAtom,
+        MockNotebook.notebookState({
+          cellData: {
+            [Cell1]: { code: code },
+          },
+        }),
+      );
       const onAddImport = vi.fn();
       maybeAddMissingImport({
         moduleName: "marimo",
@@ -58,14 +75,14 @@ describe("maybeAddMissingImport", () => {
   it("should add an import if the variable is not in the variables state and the import statement does not exist in the notebook", () => {
     const appStore = createStore();
     appStore.set(variablesAtom, {} as Variables);
-    appStore.set(notebookAtom, {
-      cellData: {
-        ["1" as CellId]: {
-          code: "mo.md('hello')",
-        } as CellData,
-      },
-      cellIds: CELL_IDS,
-    } as NotebookState);
+    appStore.set(
+      notebookAtom,
+      MockNotebook.notebookState({
+        cellData: {
+          [Cell1]: { code: "mo.md('hello')" },
+        },
+      }),
+    );
     const onAddImport = vi.fn();
     maybeAddMissingImport({
       moduleName: "marimo",
@@ -75,4 +92,32 @@ describe("maybeAddMissingImport", () => {
     });
     expect(onAddImport).toHaveBeenCalledWith("import marimo as mo");
   });
+
+  it("should not create a new cell if import already exists due to skipIfCodeExists", () => {
+    const addImports = [maybeAddMarimoImport, maybeAddAltairImport];
+    for (const addImport of addImports) {
+      store.set(variablesAtom, {} as Variables);
+      store.set(
+        notebookAtom,
+        MockNotebook.notebookState({
+          cellData: {
+            [Cell1]: { code: "import marimo as mo" },
+            [Cell2]: { code: "import altair as alt" },
+          },
+        }),
+      );
+
+      const createNewCell = vi.fn();
+      const result = addImport({
+        autoInstantiate: false,
+        createNewCell,
+        fromCellId: Cell1,
+        before: false,
+      });
+
+      // Should not create a new cell since the import already exists
+      expect(createNewCell).not.toHaveBeenCalled();
+      expect(result).toBeNull();
+    }
+  });
 });

package/src/core/cells/add-missing-import.ts

@@ -47,27 +47,40 @@ export function maybeAddMissingImport({
   return true;
 }
 
+/**
+ * Adds a marimo import to the notebook if not already present.
+ * @param autoInstantiate Whether to automatically run the cell.
+ * @param createNewCell The function to create a new cell.
+ * @param fromCellId The cell to add the import to.
+ * @param before Whether to add the import before or after the cell.
+ *
+ * Returns the ID of the new cell if added, otherwise null.
+ */
 export function maybeAddMarimoImport({
   autoInstantiate,
   createNewCell,
   fromCellId,
+  before,
 }: {
   autoInstantiate: boolean;
   createNewCell: CellActions["createNewCell"];
   fromCellId?: CellId | null;
-}): boolean {
+  before?: boolean;
+}): CellId | null {
   const client = getRequestClient();
-  return maybeAddMissingImport({
+  let newCellId: CellId | null = null;
+  const added = maybeAddMissingImport({
     moduleName: "marimo",
     variableName: "mo",
     onAddImport: (importStatement) => {
-      const newCellId = CellId.create();
+      newCellId = CellId.create();
       createNewCell({
         cellId: fromCellId ?? "__end__",
-        before: false,
+        before: before ?? false,
         code: importStatement,
         lastCodeRun: autoInstantiate ? importStatement : undefined,
         newCellId: newCellId,
+        skipIfCodeExists: true,
         autoFocus: false,
       });
       if (autoInstantiate) {
@@ -78,6 +91,7 @@ export function maybeAddMarimoImport({
       }
     },
   });
+  return added ? newCellId : null;
 }
 
 export function maybeAddAltairImport({
@@ -88,19 +102,21 @@ export function maybeAddAltairImport({
   autoInstantiate: boolean;
   createNewCell: CellActions["createNewCell"];
   fromCellId?: CellId | null;
-}): boolean {
+}): CellId | null {
   const client = getRequestClient();
-  return maybeAddMissingImport({
+  let newCellId: CellId | null = null;
+  const added = maybeAddMissingImport({
     moduleName: "altair",
     variableName: "alt",
     onAddImport: (importStatement) => {
-      const newCellId = CellId.create();
+      newCellId = CellId.create();
       createNewCell({
         cellId: fromCellId ?? "__end__",
         before: false,
         code: importStatement,
         lastCodeRun: autoInstantiate ? importStatement : undefined,
         newCellId: newCellId,
+        skipIfCodeExists: true,
         autoFocus: false,
       });
       if (autoInstantiate) {
@@ -111,4 +127,5 @@ export function maybeAddAltairImport({
       }
     },
   });
+  return added ? newCellId : null;
 }

package/src/core/cells/cells.ts

@@ -145,6 +145,31 @@ export function initialNotebookState(): NotebookState {
   });
 }
 
+export interface CreateNewCellAction {
+  /** The target cell ID to create a new cell relative to. Can be:
+   * - A CellId string for an existing cell
+   * - "__end__" to append at the end of the first column
+   * - {type: "__end__", columnId} to append at the end of a specific column
+   */
+  cellId: CellId | "__end__" | { type: "__end__"; columnId: CellColumnId };
+  /** Whether to insert before (true) or after (false) the target cell */
+  before: boolean;
+  /** Initial code content for the new cell */
+  code?: string;
+  /** The last executed code for the new cell */
+  lastCodeRun?: string;
+  /** Timestamp of the last execution */
+  lastExecutionTime?: number;
+  /** Optional custom ID for the new cell. Auto-generated if not provided */
+  newCellId?: CellId;
+  /** Whether to focus the new cell after creation */
+  autoFocus?: boolean;
+  /** If true, skip creation if code already exists */
+  skipIfCodeExists?: boolean;
+  /** Hide the code in the new cell. This will be initially shown until the cell is blurred for the first time. */
+  hideCode?: boolean;
+}
+
 /**
  * Actions and reducer for the notebook state.
  */
@@ -154,33 +179,7 @@ const {
   useActions,
   valueAtom: notebookAtom,
 } = createReducerAndAtoms(initialNotebookState, {
-  createNewCell: (
-    state,
-    action: {
-      /** The target cell ID to create a new cell relative to. Can be:
-       * - A CellId string for an existing cell
-       * - "__end__" to append at the end of the first column
-       * - {type: "__end__", columnId} to append at the end of a specific column
-       */
-      cellId: CellId | "__end__" | { type: "__end__"; columnId: CellColumnId };
-      /** Whether to insert before (true) or after (false) the target cell */
-      before: boolean;
-      /** Initial code content for the new cell */
-      code?: string;
-      /** The last executed code for the new cell */
-      lastCodeRun?: string;
-      /** Timestamp of the last execution */
-      lastExecutionTime?: number;
-      /** Optional custom ID for the new cell. Auto-generated if not provided */
-      newCellId?: CellId;
-      /** Whether to focus the new cell after creation */
-      autoFocus?: boolean;
-      /** If true, skip creation if code already exists */
-      skipIfCodeExists?: boolean;
-      /** Hide the code in the new cell. This will be initially shown until the cell is blurred for the first time. */
-      hideCode?: boolean;
-    },
-  ) => {
+  createNewCell: (state, action: CreateNewCellAction) => {
     const {
       cellId,
       before,

package/src/core/cells/logs.ts

@@ -102,7 +102,7 @@ export function formatLogTimestamp(timestamp: number): string {
   try {
     // parse from UTC
     const date = fromUnixTime(timestamp);
-    return date.toLocaleTimeString("en-US", {
+    return date.toLocaleTimeString(undefined, {
       hour12: true,
       hour: "numeric",
       minute: "numeric",

package/src/core/codemirror/find-replace/search-highlight.ts

@@ -91,8 +91,10 @@ const HighlightMargin = 250;
 export const searchHighlighter = ViewPlugin.fromClass(
   class {
     decorations: DecorationSet;
+    readonly view: EditorView;
 
-    constructor(readonly view: EditorView) {
+    constructor(view: EditorView) {
+      this.view = view;
       this.decorations = this.highlight(view.state.field(searchState));
     }
 

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

@@ -1,20 +1,26 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
+import { once } from "@/utils/once";
 import { MarkdownLanguageAdapter } from "./languages/markdown";
 import { PythonLanguageAdapter } from "./languages/python";
 import { SQLLanguageAdapter } from "./languages/sql/sql";
 import type { LanguageAdapter, LanguageAdapterType } from "./types";
 
+// Create cached instances
+const createPythonAdapter = once(() => new PythonLanguageAdapter());
+const createMarkdownAdapter = once(() => new MarkdownLanguageAdapter());
+const createSqlAdapter = once(() => new SQLLanguageAdapter());
+
 export const LanguageAdapters: Record<LanguageAdapterType, LanguageAdapter> = {
   // Getters to prevent circular dependencies
   get python() {
-    return new PythonLanguageAdapter();
+    return createPythonAdapter();
   },
   get markdown() {
-    return new MarkdownLanguageAdapter();
+    return createMarkdownAdapter();
   },
   get sql() {
-    return new SQLLanguageAdapter();
+    return createSqlAdapter();
   },
 };
 

package/src/core/codemirror/language/__tests__/extension.test.ts

@@ -14,6 +14,7 @@ import {
   languageAdapterState,
   switchLanguage,
 } from "../extension";
+import { exportedForTesting as sqlValidationErrorsForTesting } from "../languages/sql/validation-errors";
 import { languageMetadataField } from "../metadata";
 
 let view: EditorView | null = null;
@@ -258,3 +259,26 @@ describe("switchLanguage", () => {
     });
   });
 });
+
+describe("sqlValidationErrors", () => {
+  const { splitErrorMessage } = sqlValidationErrorsForTesting;
+
+  describe("split error message", () => {
+    it("should split the error message into error type and error message", () => {
+      const error = "SyntaxError: SELECT * FROM df";
+      const { errorType, errorMessage } = splitErrorMessage(error);
+      expect(errorType).toBe("SyntaxError");
+      expect(errorMessage).toBe("SELECT * FROM df");
+    });
+
+    it("should handle multiple colons", () => {
+      const error =
+        "SyntaxError: SELECT * FROM df:SyntaxError: SELECT * FROM df";
+      const { errorType, errorMessage } = splitErrorMessage(error);
+      expect(errorType).toBe("SyntaxError");
+      expect(errorMessage).toBe(
+        "SELECT * FROM df:SyntaxError: SELECT * FROM df",
+      );
+    });
+  });
+});

package/src/core/codemirror/language/__tests__/sql-validation.test.ts

@@ -0,0 +1,133 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import { describe, expect, it } from "vitest";
+import { exportedForTesting } from "../languages/sql/validation-errors";
+
+describe("Error Message Splitting", () => {
+  it("should handle error message splitting correctly", () => {
+    const { splitErrorMessage } = exportedForTesting;
+
+    const result1 = splitErrorMessage("Syntax error: unexpected token");
+    expect(result1.errorType).toBe("Syntax error");
+    expect(result1.errorMessage).toBe("unexpected token");
+
+    const result2 = splitErrorMessage("Multiple: colons: in error");
+    expect(result2.errorType).toBe("Multiple");
+    expect(result2.errorMessage).toBe("colons: in error");
+
+    const result3 = splitErrorMessage("No colon error");
+    expect(result3.errorType).toBe("No colon error");
+    expect(result3.errorMessage).toBe("");
+  });
+});
+
+describe("DuckDB Error Handling", () => {
+  it("should extract codeblock from error with LINE information", () => {
+    const { handleDuckdbError } = exportedForTesting;
+
+    const error =
+      'Binder Error: Referenced column "attacks" not found in FROM clause! Candidate bindings: "Attack", "Total" LINE 1:... from pokemon WHERE \'type_2\' = 32 and attack = 32 and not attacks = \'hi\' ^';
+
+    const result = handleDuckdbError(error);
+
+    expect(result.errorType).toBe("Binder Error");
+    expect(result.errorMessage).toBe(
+      'Referenced column "attacks" not found in FROM clause! Candidate bindings: "Attack", "Total"',
+    );
+    expect(result.codeblock).toBe(
+      "LINE 1:... from pokemon WHERE 'type_2' = 32 and attack = 32 and not attacks = 'hi' ^",
+    );
+  });
+
+  it("should handle error without LINE information", () => {
+    const { handleDuckdbError } = exportedForTesting;
+
+    const error = "Syntax Error: Invalid syntax near WHERE";
+
+    const result = handleDuckdbError(error);
+
+    expect(result.errorType).toBe("Syntax Error");
+    expect(result.errorMessage).toBe("Invalid syntax near WHERE");
+    expect(result.codeblock).toBeUndefined();
+  });
+
+  it("should handle error with LINE at the beginning", () => {
+    const { handleDuckdbError } = exportedForTesting;
+
+    const error = "LINE 1: SELECT * FROM table WHERE invalid_column = 1 ^";
+
+    const result = handleDuckdbError(error);
+
+    expect(result.errorType).toBe("LINE 1");
+    expect(result.errorMessage).toBe(
+      "SELECT * FROM table WHERE invalid_column = 1 ^",
+    );
+    expect(result.codeblock).toBeUndefined();
+  });
+
+  it("should handle error with multiple LINE occurrences", () => {
+    const { handleDuckdbError } = exportedForTesting;
+
+    const error =
+      "Error: Something went wrong LINE 1: SELECT * FROM table WHERE invalid_column = 1 ^";
+
+    const result = handleDuckdbError(error);
+
+    expect(result.errorType).toBe("Error");
+    expect(result.errorMessage).toBe("Something went wrong");
+    expect(result.codeblock).toBe(
+      "LINE 1: SELECT * FROM table WHERE invalid_column = 1 ^",
+    );
+  });
+
+  it("should handle complex error with nested quotes", () => {
+    const { handleDuckdbError } = exportedForTesting;
+
+    const error =
+      "Binder Error: Column \"name\" not found! LINE 1: SELECT * FROM users WHERE name = 'John' AND age > 25 ^";
+
+    const result = handleDuckdbError(error);
+
+    expect(result.errorType).toBe("Binder Error");
+    expect(result.errorMessage).toBe('Column "name" not found!');
+    expect(result.codeblock).toBe(
+      "LINE 1: SELECT * FROM users WHERE name = 'John' AND age > 25 ^",
+    );
+  });
+
+  it("should handle error with LINE but no caret", () => {
+    const { handleDuckdbError } = exportedForTesting;
+
+    const error = "Error: Invalid query LINE 1: SELECT * FROM table";
+
+    const result = handleDuckdbError(error);
+
+    expect(result.errorType).toBe("Error");
+    expect(result.errorMessage).toBe("Invalid query");
+    expect(result.codeblock).toBe("LINE 1: SELECT * FROM table");
+  });
+
+  it("should trim whitespace from codeblock", () => {
+    const { handleDuckdbError } = exportedForTesting;
+
+    const error = "Error: Something wrong   LINE 1: SELECT * FROM table   ^   ";
+
+    const result = handleDuckdbError(error);
+
+    expect(result.errorType).toBe("Error");
+    expect(result.errorMessage).toBe("Something wrong");
+    expect(result.codeblock).toBe("LINE 1: SELECT * FROM table   ^");
+  });
+
+  it("should handle empty error message", () => {
+    const { handleDuckdbError } = exportedForTesting;
+
+    const error = "";
+
+    const result = handleDuckdbError(error);
+
+    expect(result.errorType).toBe("");
+    expect(result.errorMessage).toBe("");
+    expect(result.codeblock).toBeUndefined();
+  });
+});

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

@@ -0,0 +1,20 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import { useAtom } from "jotai";
+import { atomWithStorage } from "jotai/utils";
+import { store } from "@/core/state/jotai";
+
+const BASE_KEY = "marimo:notebook-sql-mode";
+
+export type SQLMode = "validate" | "default";
+
+const sqlModeAtom = atomWithStorage<SQLMode>(BASE_KEY, "default");
+
+export function useSQLMode() {
+  const [sqlMode, setSQLMode] = useAtom(sqlModeAtom);
+  return { sqlMode, setSQLMode };
+}
+
+export function getSQLMode() {
+  return store.get(sqlModeAtom);
+}

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

@@ -5,7 +5,7 @@ import { insertTab } from "@codemirror/commands";
 import { type SQLDialect, type SQLNamespace, sql } from "@codemirror/lang-sql";
 import type { EditorState, Extension } from "@codemirror/state";
 import { Compartment } from "@codemirror/state";
-import { type EditorView, keymap } from "@codemirror/view";
+import { EditorView, keymap } from "@codemirror/view";
 import type { SyntaxNode, TreeCursor } from "@lezer/common";
 import { parser } from "@lezer/python";
 import {
@@ -16,12 +16,19 @@ import {
 } from "@marimo-team/codemirror-sql";
 import { DuckDBDialect } from "@marimo-team/codemirror-sql/dialects";
 import dedent from "string-dedent";
+import { cellIdState } from "@/core/codemirror/cells/state";
 import { getFeatureFlag } from "@/core/config/feature-flag";
 import {
   dataSourceConnectionsAtom,
   setLatestEngineSelected,
 } from "@/core/datasets/data-source-connections";
-import { type ConnectionName, DUCKDB_ENGINE } from "@/core/datasets/engines";
+import {
+  type ConnectionName,
+  DUCKDB_ENGINE,
+  INTERNAL_SQL_ENGINES,
+} from "@/core/datasets/engines";
+import { ValidateSQL } from "@/core/datasets/request-registry";
+import type { ValidateSQLResult } from "@/core/kernel/messages";
 import { store } from "@/core/state/jotai";
 import { resolvedThemeAtom } from "@/theme/useTheme";
 import { Logger } from "@/utils/Logger";
@@ -37,6 +44,11 @@ import {
   tablesCompletionSource,
 } from "./completion-sources";
 import { SCHEMA_CACHE } from "./completion-store";
+import { getSQLMode } from "./sql-mode";
+import {
+  clearSqlValidationError,
+  setSqlValidationError,
+} from "./validation-errors";
 
 const DEFAULT_DIALECT = DuckDBDialect;
 const DEFAULT_PARSER_DIALECT = "DuckDB";
@@ -64,12 +76,15 @@ export class SQLLanguageAdapter
 {
   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;
     }
   }
 
@@ -265,6 +280,10 @@ export class SQLLanguageAdapter
       );
     }
 
+    if (this.sqlModeEnabled) {
+      extensions.push(sqlValidationExtension());
+    }
+
     return extensions;
   }
 }
@@ -315,9 +334,14 @@ function getSchema(view: EditorView): SQLNamespace {
 function guessParserDialect(state: EditorState): ParserDialects | null {
   const metadata = getSQLMetadata(state);
   const connectionName = metadata.engine;
+  return connectionNameToParserDialect(connectionName);
+}
+
+function connectionNameToParserDialect(
+  connectionName: ConnectionName,
+): ParserDialects | null {
   const dialect =
     SCHEMA_CACHE.getInternalDialect(connectionName)?.toLowerCase();
-
   switch (dialect) {
     case "postgresql":
     case "postgres":
@@ -543,3 +567,66 @@ function safeDedent(code: string): string {
     return code;
   }
 }
+
+function sqlValidationExtension(): Extension {
+  let debounceTimeout: NodeJS.Timeout | null = null;
+  let lastValidationRequest: string | null = null;
+
+  return EditorView.updateListener.of((update) => {
+    const sqlMode = getSQLMode();
+    if (sqlMode !== "validate") {
+      return;
+    }
+
+    const metadata = getSQLMetadata(update.state);
+    const connectionName = metadata.engine;
+    if (!INTERNAL_SQL_ENGINES.has(connectionName)) {
+      // Currently only internal engines are supported
+      return;
+    }
+
+    if (!update.docChanged) {
+      return;
+    }
+
+    const doc = update.state.doc;
+    const sqlContent = doc.toString();
+
+    // Clear existing timeout
+    if (debounceTimeout) {
+      clearTimeout(debounceTimeout);
+    }
+
+    // Debounce the validation call
+    debounceTimeout = setTimeout(async () => {
+      // Skip if content hasn't changed since last validation
+      if (lastValidationRequest === sqlContent) {
+        return;
+      }
+
+      lastValidationRequest = sqlContent;
+      const cellId = update.view.state.facet(cellIdState);
+
+      if (sqlContent === "") {
+        clearSqlValidationError(cellId);
+        return;
+      }
+
+      try {
+        const result: ValidateSQLResult = await ValidateSQL.request({
+          engine: connectionName,
+          query: sqlContent,
+        });
+
+        if (result.error) {
+          const dialect = connectionNameToParserDialect(connectionName);
+          setSqlValidationError({ cellId, error: result.error, dialect });
+        } else {
+          clearSqlValidationError(cellId);
+        }
+      } catch (error) {
+        Logger.warn("Failed to validate SQL", { error });
+      }
+    }, 300);
+  });
+}