@marimo-team/frontend 0.16.1 → 0.16.2

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({

package/src/plugins/impl/DataTablePlugin.tsx

@@ -36,6 +36,10 @@ import {
 import { usePanelOwnership } from "@/components/data-table/hooks/use-panel-ownership";
 import { LoadingTable } from "@/components/data-table/loading-table";
 import { RowViewerPanel } from "@/components/data-table/row-viewer-panel/row-viewer";
+import {
+  type DownloadAsArgs,
+  DownloadAsSchema,
+} from "@/components/data-table/schemas";
 import {
   type BinValues,
   type ColumnHeaderStats,
@@ -182,6 +186,7 @@ interface Data<T> {
   freezeColumnsRight?: string[];
   textJustifyColumns?: Record<string, "left" | "center" | "right">;
   wrappedColumns?: string[];
+  headerTooltip?: Record<string, string>;
   totalColumns: number;
   maxColumns: number | "all";
   hasStableRowId: boolean;
@@ -190,7 +195,7 @@ interface Data<T> {
 
 // eslint-disable-next-line @typescript-eslint/consistent-type-definitions
 type DataTableFunctions = {
-  download_as: (req: { format: "csv" | "json" | "parquet" }) => Promise<string>;
+  download_as: DownloadAsArgs;
   get_column_summaries: <T>(
     opts: ColumnSummariesArgs,
   ) => Promise<ColumnSummaries<T>>;
@@ -249,6 +254,7 @@ export const DataTablePlugin = createPlugin<S>("marimo-table")
         .record(z.enum(["left", "center", "right"]))
         .optional(),
       wrappedColumns: z.array(z.string()).optional(),
+      headerTooltip: z.record(z.string()).optional(),
       fieldTypes: columnToFieldTypesSchema.nullish(),
       totalColumns: z.number(),
       maxColumns: z.union([z.number(), z.literal("all")]).default("all"),
@@ -263,9 +269,7 @@ export const DataTablePlugin = createPlugin<S>("marimo-table")
     }),
   )
   .withFunctions<DataTableFunctions>({
-    download_as: rpc
-      .input(z.object({ format: z.enum(["csv", "json", "parquet"]) }))
-      .output(z.string()),
+    download_as: DownloadAsSchema,
     get_column_summaries: rpc
       .input(z.object({ precompute: z.boolean() }))
       .output(
@@ -706,6 +710,7 @@ const DataTableComponent = ({
   freezeColumnsRight,
   textJustifyColumns,
   wrappedColumns,
+  headerTooltip,
   totalColumns,
   get_row_ids,
   cellStyles,
@@ -779,6 +784,7 @@ const DataTableComponent = ({
         fieldTypes: memoizedClampedFieldTypes,
         textJustifyColumns: memoizedTextJustifyColumns,
         wrappedColumns: memoizedWrappedColumns,
+        headerTooltip: headerTooltip,
         // Only show data types if they are explicitly set
         showDataTypes: showDataTypes,
         calculateTopKRows: calculate_top_k_rows,
@@ -791,6 +797,7 @@ const DataTableComponent = ({
       memoizedClampedFieldTypes,
       memoizedTextJustifyColumns,
       memoizedWrappedColumns,
+      headerTooltip,
       calculate_top_k_rows,
     ],
   );

package/src/plugins/impl/data-frames/DataFramePlugin.tsx

@@ -4,6 +4,10 @@ import { isEqual } from "lodash-es";
 import { Code2Icon, DatabaseIcon, FunctionSquareIcon } from "lucide-react";
 import { type JSX, memo, useEffect, useRef, useState } from "react";
 import { z } from "zod";
+import {
+  type DownloadAsArgs,
+  DownloadAsSchema,
+} from "@/components/data-table/schemas";
 import type { FieldTypesWithExternalType } from "@/components/data-table/types";
 import { ReadonlyCode } from "@/components/editor/code/readonly-python-code";
 import { Spinner } from "@/components/icons/spinner";
@@ -38,6 +42,7 @@ interface Data {
   label?: string | null;
   columns: ColumnDataTypes;
   pageSize: number;
+  showDownload: boolean;
 }
 
 // eslint-disable-next-line @typescript-eslint/consistent-type-definitions
@@ -67,6 +72,7 @@ type PluginFunctions = {
     data: TableData<T>;
     total_rows: number;
   }>;
+  download_as: DownloadAsArgs;
 };
 
 // Value is selection, but it is not currently exposed to the user
@@ -77,13 +83,14 @@ export const DataFramePlugin = createPlugin<S>("marimo-dataframe")
     z.object({
       label: z.string().nullish(),
       pageSize: z.number().default(5),
+      showDownload: z.boolean().default(true),
       columns: z
         .array(z.tuple([z.string().or(z.number()), z.string(), z.string()]))
         .transform((value) => {
           const map = new Map<ColumnId, string>();
-          value.forEach(([key, dataType]) =>
-            map.set(key as ColumnId, dataType as DataType),
-          );
+          value.forEach(([key, dataType]) => {
+            map.set(key as ColumnId, dataType as DataType);
+          });
           return map;
         }),
     }),
@@ -124,6 +131,7 @@ export const DataFramePlugin = createPlugin<S>("marimo-dataframe")
           total_rows: z.number(),
         }),
       ),
+    download_as: DownloadAsSchema,
   })
   .renderer((props) => (
     <TableProviders>
@@ -141,6 +149,8 @@ interface DataTableProps extends Data, PluginFunctions {
   value: S;
   setValue: (value: S) => void;
   host: HTMLElement;
+  showDownload: boolean;
+  download_as: DownloadAsArgs;
 }
 
 const EMPTY: Transformations = {
@@ -151,11 +161,13 @@ export const DataFrameComponent = memo(
   ({
     columns,
     pageSize,
+    showDownload,
     value,
     setValue,
     get_dataframe,
     get_column_values,
     search,
+    download_as,
     host,
   }: DataTableProps): JSX.Element => {
     const { data, error, isPending } = useAsyncData(
@@ -270,8 +282,8 @@ export const DataFrameComponent = memo(
           pagination={true}
           fieldTypes={field_types}
           rowHeaders={row_headers || Arrays.EMPTY}
-          showDownload={false}
-          download_as={Functions.THROW}
+          showDownload={showDownload}
+          download_as={download_as}
           enableSearch={false}
           showFilters={false}
           search={search}

package/src/stories/dataframe.stories.tsx

@@ -40,6 +40,8 @@ export const DataFrame: StoryObj = {
         get_dataframe={() => Promise.reject(new Error("not implemented"))}
         search={Functions.THROW}
         host={document.body}
+        showDownload={false}
+        download_as={async () => ""}
       />
     );
   },