@marimo-team/islands 0.16.1 → 0.16.2

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);
+  });
+}

package/src/core/codemirror/language/languages/sql/validation-errors.ts

@@ -0,0 +1,79 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import type { SupportedDialects } from "@marimo-team/codemirror-sql";
+import { atom, useAtomValue } from "jotai";
+import type { CellId } from "@/core/cells/ids";
+import { store } from "@/core/state/jotai";
+
+export interface SQLValidationError {
+  errorType: string;
+  errorMessage: string;
+  codeblock?: string; // Code block that caused the error
+}
+
+type CellToSQLErrors = Map<CellId, SQLValidationError>;
+
+export const sqlValidationErrorsAtom = atom<CellToSQLErrors>(
+  new Map<CellId, SQLValidationError>(),
+);
+
+export const useSqlValidationErrorsForCell = (cellId: CellId) => {
+  const sqlValidationErrors = useAtomValue(sqlValidationErrorsAtom);
+  return sqlValidationErrors.get(cellId);
+};
+
+export function clearSqlValidationError(cellId: CellId) {
+  const sqlValidationErrors = store.get(sqlValidationErrorsAtom);
+  const newErrors = new Map(sqlValidationErrors);
+  newErrors.delete(cellId);
+  store.set(sqlValidationErrorsAtom, newErrors);
+}
+
+export function setSqlValidationError({
+  cellId,
+  error,
+  dialect,
+}: {
+  cellId: CellId;
+  error: string;
+  dialect: SupportedDialects | null;
+}) {
+  const sqlValidationErrors = store.get(sqlValidationErrorsAtom);
+  const newErrors = new Map(sqlValidationErrors);
+
+  const errorResult: SQLValidationError =
+    dialect === "DuckDB" ? handleDuckdbError(error) : splitErrorMessage(error);
+
+  newErrors.set(cellId, errorResult);
+  store.set(sqlValidationErrorsAtom, newErrors);
+}
+
+function handleDuckdbError(error: string): SQLValidationError {
+  const { errorType, errorMessage } = splitErrorMessage(error);
+  let newErrorMessage = errorMessage;
+
+  // Extract the LINE and the rest of the message as codeblock, keep errorMessage as whatever is before
+  let codeblock: string | undefined;
+  const lineIndex = errorMessage.indexOf("LINE ");
+  if (lineIndex !== -1) {
+    codeblock = errorMessage.slice(Math.max(0, lineIndex)).trim();
+    newErrorMessage = errorMessage.slice(0, Math.max(0, lineIndex)).trim();
+  }
+
+  return {
+    errorType,
+    errorMessage: newErrorMessage,
+    codeblock,
+  };
+}
+
+function splitErrorMessage(error: string) {
+  const errorType = error.split(":")[0].trim();
+  const errorMessage = error.split(":").slice(1).join(":").trim();
+  return { errorType, errorMessage };
+}
+
+export const exportedForTesting = {
+  splitErrorMessage,
+  handleDuckdbError,
+};

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

@@ -5,7 +5,8 @@ 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 type { ConnectionName } from "@/core/datasets/engines";
+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";
 import { languageAdapterState } from "../extension";
@@ -22,7 +23,7 @@ import {
 import type { LanguageMetadataOf } from "../types";
 import type { QuotePrefixKind } from "../utils/quotes";
 import { getQuotePrefix, MarkdownQuotePrefixTooltip } from "./markdown";
-import { SQLEngineSelect } from "./sql";
+import { SQLEngineSelect, SQLModeSelect } from "./sql";
 
 const Divider = () => <div className="h-4 border-r border-border" />;
 
@@ -70,6 +71,8 @@ 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">
@@ -95,6 +98,9 @@ export const LanguagePanelComponent: React.FC<{
           onChange={switchEngine}
         />
         <div className="flex items-center gap-2 ml-auto">
+          {sqlModeEnabled && metadata.engine === DUCKDB_ENGINE && (
+            <SQLModeSelect />
+          )}
           <Tooltip content="Format SQL">
             <Button
               variant="text"

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

@@ -1,9 +1,16 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
+import type { SelectTriggerProps } from "@radix-ui/react-select";
 import { useAtomValue } from "jotai";
-import { AlertCircle, CircleHelpIcon } from "lucide-react";
+import {
+  AlertCircle,
+  CircleHelpIcon,
+  DatabaseBackup,
+  SearchCheck,
+} from "lucide-react";
 import { transformDisplayName } from "@/components/databases/display";
 import { DatabaseLogo } from "@/components/databases/icon";
+import { Button } from "@/components/ui/button";
 import {
   Select,
   SelectContent,
@@ -14,6 +21,7 @@ import {
   SelectTrigger,
   SelectValue,
 } from "@/components/ui/select";
+import { Tooltip } from "@/components/ui/tooltip";
 import {
   dataConnectionsMapAtom,
   setLatestEngineSelected,
@@ -24,6 +32,7 @@ import {
 } from "@/core/datasets/engines";
 import type { DataSourceConnection } from "@/core/kernel/messages";
 import { useNonce } from "@/hooks/useNonce";
+import { type SQLMode, useSQLMode } from "../languages/sql/sql-mode";
 
 interface SelectProps {
   selectedEngine: ConnectionName;
@@ -77,7 +86,7 @@ export const SQLEngineSelect: React.FC<SelectProps> = ({
       <SelectItem key={connection.name} value={connection.name}>
         <div className="flex items-center gap-1">
           <DatabaseLogo className="h-3 w-3" name={connection.dialect} />
-          <span className="truncate">
+          <span className="truncate ml-0.5">
             {transformDisplayName(connection.display_name)}
           </span>
         </div>
@@ -88,9 +97,9 @@ export const SQLEngineSelect: React.FC<SelectProps> = ({
   return (
     <div className="flex flex-row gap-1 items-center">
       <Select value={selectedEngine} onValueChange={handleSelectEngine}>
-        <SelectTrigger className="text-xs border-border shadow-none! ring-0! h-4.5 px-1.5">
+        <SQLSelectTrigger>
           <SelectValue placeholder="Select an engine" />
-        </SelectTrigger>
+        </SQLSelectTrigger>
         <SelectContent>
           <SelectGroup>
             <SelectLabel>Database connections</SelectLabel>
@@ -130,3 +139,71 @@ export const SQLEngineSelect: React.FC<SelectProps> = ({
 const HELP_KEY = "__help__";
 const HELP_URL =
   "http://docs.marimo.io/guides/working_with_data/sql/#connecting-to-a-custom-database";
+
+export const SQLModeSelect: React.FC = () => {
+  const { sqlMode, setSQLMode } = useSQLMode();
+
+  const handleToggleMode = () => {
+    setSQLMode(sqlMode === "validate" ? "default" : "validate");
+  };
+
+  const getModeIcon = (mode: SQLMode) => {
+    return mode === "validate" ? (
+      <SearchCheck className="h-3 w-3" />
+    ) : (
+      <DatabaseBackup className="h-3 w-3" />
+    );
+  };
+
+  const getTooltipContent = (mode: SQLMode) => {
+    return mode === "validate" ? (
+      <div className="text-xs">
+        <div className="font-semibold mb-1 flex flex-row items-center gap-1">
+          <SearchCheck className="h-3 w-3" />
+          Validate Mode
+        </div>
+        <p>Queries are validated as you write them</p>
+      </div>
+    ) : (
+      <div className="text-xs">
+        <div className="font-semibold mb-1 flex flex-row items-center gap-1">
+          <DatabaseBackup className="h-3 w-3" />
+          Default Mode
+        </div>
+        <p>Standard editing</p>
+      </div>
+    );
+  };
+
+  return (
+    <div className="flex flex-row gap-1 items-center">
+      <Tooltip delayDuration={300} content={getTooltipContent(sqlMode)}>
+        <Button
+          variant="ghost"
+          size="sm"
+          onClick={handleToggleMode}
+          className="h-5 px-1.5 text-xs border-border shadow-none hover:bg-accent"
+        >
+          {getModeIcon(sqlMode)}
+          <span className="ml-1">
+            {sqlMode === "validate" ? "Validate" : "Default"}
+          </span>
+        </Button>
+      </Tooltip>
+    </div>
+  );
+};
+
+const SQLSelectTrigger: React.FC<SelectTriggerProps> = ({
+  children,
+  ...props
+}) => {
+  return (
+    <SelectTrigger
+      className="text-xs border-border shadow-none! ring-0! h-5 px-1.5 hover:bg-accent transition-colors"
+      {...props}
+    >
+      {children}
+    </SelectTrigger>
+  );
+};

package/src/core/config/feature-flag.tsx

@@ -14,6 +14,7 @@ export interface ExperimentalFeatures {
   mcp_docs: boolean;
   sql_linter: boolean;
   external_agents: boolean;
+  sql_mode: boolean;
   // Add new feature flags here
 }
 
@@ -26,13 +27,14 @@ const defaultValues: ExperimentalFeatures = {
   mcp_docs: false,
   sql_linter: false,
   external_agents: import.meta.env.DEV,
+  sql_mode: false,
 };
 
 export function getFeatureFlag<T extends keyof ExperimentalFeatures>(
   feature: T,
 ): ExperimentalFeatures[T] {
   return (
-    (getResolvedMarimoConfig().experimental?.[
+    (getResolvedMarimoConfig()?.experimental?.[
       feature
     ] as ExperimentalFeatures[T]) ?? defaultValues[feature]
   );

package/src/core/datasets/request-registry.ts

@@ -1,10 +1,15 @@
 /* Copyright 2024 Marimo. All rights reserved. */
-import type { SQLTableListPreview, SQLTablePreview } from "../kernel/messages";
+import type {
+  SQLTableListPreview,
+  SQLTablePreview,
+  ValidateSQLResult,
+} from "../kernel/messages";
 import { DeferredRequestRegistry } from "../network/DeferredRequestRegistry";
 import { getRequestClient } from "../network/requests";
 import type {
   PreviewSQLTableListRequest,
   PreviewSQLTableRequest,
+  ValidateSQLRequest,
 } from "../network/types";
 
 // We make a request to the backend to preview the table, passing in Engine, DB, Schema, and Table
@@ -32,3 +37,14 @@ export const PreviewSQLTableList = new DeferredRequestRegistry<
     ...req,
   });
 });
+
+export const ValidateSQL = new DeferredRequestRegistry<
+  Omit<ValidateSQLRequest, "requestId">,
+  ValidateSQLResult
+>("validate-sql", async (requestId, req) => {
+  const client = getRequestClient();
+  await client.validateSQL({
+    requestId: requestId,
+    ...req,
+  });
+});

package/src/core/islands/bridge.ts

@@ -153,6 +153,7 @@ export class IslandsPyodideBridge implements RunRequests, EditRequests {
   previewSQLTable = throwNotImplemented;
   previewSQLTableList = throwNotImplemented;
   previewDataSourceConnection = throwNotImplemented;
+  validateSQL = throwNotImplemented;
   openFile = throwNotImplemented;
   sendListFiles = throwNotImplemented;
   sendSearchFiles = throwNotImplemented;

package/src/core/islands/main.ts

@@ -119,6 +119,7 @@ export async function initialize() {
       case "sql-table-list-preview":
       case "datasets":
       case "data-source-connections":
+      case "validate-sql-result":
       case "secret-keys-result":
       case "startup-logs":
         // Unsupported

package/src/core/kernel/messages.ts

@@ -38,6 +38,7 @@ export type DataColumnPreview = OperationMessageData<"data-column-preview">;
 export type SQLTablePreview = OperationMessageData<"sql-table-preview">;
 export type SQLTableListPreview =
   OperationMessageData<"sql-table-list-preview">;
+export type ValidateSQLResult = OperationMessageData<"validate-sql-result">;
 export type SecretKeysResult = OperationMessageData<"secret-keys-result">;
 export type StartupLogs = OperationMessageData<"startup-logs">;
 export type CellMessage = OperationMessageData<"cell-op">;

package/src/core/network/requests-network.ts

@@ -207,6 +207,13 @@ export function createNetworkRequests(): EditRequests & RunRequests {
         })
         .then(handleResponseReturnNull);
     },
+    validateSQL: (request) => {
+      return getClient()
+        .POST("/api/sql/validate", {
+          body: request,
+        })
+        .then(handleResponseReturnNull);
+    },
     openFile: async (request) => {
       await waitForConnectionOpen();
       await getClient()

package/src/core/network/requests-static.ts

@@ -57,6 +57,7 @@ export function createStaticRequests(): EditRequests & RunRequests {
     previewSQLTable: throwNotInEditMode,
     previewSQLTableList: throwNotInEditMode,
     previewDataSourceConnection: throwNotInEditMode,
+    validateSQL: throwNotInEditMode,
     openFile: throwNotInEditMode,
     getUsageStats: throwNotInEditMode,
     sendListFiles: throwNotInEditMode,

package/src/core/network/requests-toasting.ts

@@ -35,6 +35,7 @@ export function createErrorToastingRequests(
     previewSQLTable: "Failed to fetch SQL table",
     previewSQLTableList: "Failed to fetch SQL table list",
     previewDataSourceConnection: "Failed to preview data source connection",
+    validateSQL: "Failed to validate SQL",
     openFile: "Failed to open file",
     getUsageStats: "", // No toast
     sendListFiles: "Failed to list files",

package/src/core/network/types.ts

@@ -62,6 +62,7 @@ export type PreviewSQLTableRequest = schemas["PreviewSQLTableRequest"];
 export type PreviewSQLTableListRequest = schemas["PreviewSQLTableListRequest"];
 export type PreviewDataSourceConnectionRequest =
   schemas["PreviewDataSourceConnectionRequest"];
+export type ValidateSQLRequest = schemas["ValidateSQLRequest"];
 export type PdbRequest = schemas["PdbRequest"];
 export type ReadCodeResponse = schemas["ReadCodeResponse"];
 export type RecentFilesResponse = schemas["RecentFilesResponse"];
@@ -140,6 +141,7 @@ export interface EditRequests {
   previewDataSourceConnection: (
     request: PreviewDataSourceConnectionRequest,
   ) => Promise<null>;
+  validateSQL: (request: ValidateSQLRequest) => Promise<null>;
   openFile: (request: { path: string }) => Promise<null>;
   getUsageStats: () => Promise<UsageResponse>;
   // Debugger

package/src/core/wasm/bridge.ts

@@ -496,6 +496,11 @@ export class PyodideBridge implements RunRequests, EditRequests {
       return null;
     };
 
+  validateSQL: EditRequests["validateSQL"] = async (request) => {
+    await this.putControlRequest(request);
+    return null;
+  };
+
   sendModelValue: RunRequests["sendModelValue"] = async (request) => {
     await this.putControlRequest(request);
     return null;

package/src/core/websocket/useMarimoWebSocket.tsx

@@ -32,6 +32,7 @@ import type { ConnectionName } from "../datasets/engines";
 import {
   PreviewSQLTable,
   PreviewSQLTableList,
+  ValidateSQL,
 } from "../datasets/request-registry";
 import { useDatasetsActions } from "../datasets/state";
 import { UI_ELEMENT_REGISTRY } from "../dom/uiregistry";
@@ -238,6 +239,9 @@ export function useMarimoWebSocket(opts: {
       case "sql-table-list-preview":
         PreviewSQLTableList.resolve(msg.data.request_id as RequestId, msg.data);
         return;
+      case "validate-sql-result":
+        ValidateSQL.resolve(msg.data.request_id as RequestId, msg.data);
+        return;
       case "secret-keys-result":
         SECRETS_REGISTRY.resolve(msg.data.request_id as RequestId, msg.data);
         return;

package/src/plugins/core/registerReactComponent.tsx

@@ -24,7 +24,7 @@ import ReactDOM, { type Root } from "react-dom/client";
 import useEvent from "react-use-event-hook";
 import type { ZodSchema } from "zod";
 import { notebookAtom } from "@/core/cells/cells.ts";
-import { findCellId } from "@/core/cells/ids.ts";
+import { HTMLCellId } from "@/core/cells/ids.ts";
 import { isUninstantiated } from "@/core/cells/utils";
 import { createInputEvent, MarimoValueUpdateEvent } from "@/core/dom/events";
 import { getUIElementObjectId } from "@/core/dom/ui-element";
@@ -182,24 +182,28 @@ function PluginSlotInternal<T>(
         const objectId = getUIElementObjectId(hostElement);
         invariant(objectId, "Object ID should exist");
 
-        const cellId = findCellId(hostElement);
-        invariant(cellId, "Cell ID should exist");
-
-        const notebookState = store.get(notebookAtom);
-        const cellRuntime = notebookState.cellRuntime[cellId];
-        const cellData = notebookState.cellData[cellId];
-
-        const cellNotInitialized = isUninstantiated({
-          executionTime:
-            cellRuntime.runElapsedTimeMs ?? cellData.lastExecutionTime,
-          status: cellRuntime.status,
-          errored: cellRuntime.errored,
-          interrupted: cellRuntime.interrupted,
-          stopped: cellRuntime.stopped,
-        });
-
-        if (cellNotInitialized) {
-          throw new CellNotInitializedError();
+        const htmlId = HTMLCellId.findElementThroughShadowDOMs(hostElement)?.id;
+        const cellId = htmlId ? HTMLCellId.parse(htmlId) : null;
+        if (cellId) {
+          // If the cell is not initialized, throw an error
+          const notebookState = store.get(notebookAtom);
+          const cellRuntime = notebookState.cellRuntime[cellId];
+          const cellData = notebookState.cellData[cellId];
+
+          const cellNotInitialized = isUninstantiated({
+            executionTime:
+              cellRuntime.runElapsedTimeMs ?? cellData.lastExecutionTime,
+            status: cellRuntime.status,
+            errored: cellRuntime.errored,
+            interrupted: cellRuntime.interrupted,
+            stopped: cellRuntime.stopped,
+          });
+
+          if (cellNotInitialized) {
+            throw new CellNotInitializedError();
+          }
+        } else {
+          Logger.warn(`Cell ID ${cellId} cannot be found`);
         }
 
         const response = await FUNCTIONS_REGISTRY.request({