@marimo-team/frontend 0.16.2 → 0.16.3

package/src/components/editor/database/add-database-form.tsx

@@ -39,6 +39,7 @@ import {
   ChdbConnectionSchema,
   ClickhouseConnectionSchema,
   type DatabaseConnection,
+  DatabricksConnectionSchema,
   DataFusionConnectionSchema,
   DuckDBConnectionSchema,
   IcebergConnectionSchema,
@@ -210,6 +211,16 @@ const DATABASES = [
       preferred: "redshift",
     },
   },
+  {
+    name: "Databricks",
+    schema: DatabricksConnectionSchema,
+    color: "#c41e0c",
+    logo: "databricks",
+    connectionLibraries: {
+      libraries: ["sqlalchemy", "sqlmodel", "ibis"],
+      preferred: "sqlalchemy",
+    },
+  },
 ] satisfies ConnectionSchema[];
 
 const DATA_CATALOGS = [

package/src/components/editor/database/as-code.ts

@@ -14,7 +14,8 @@ export type ConnectionLibrary =
   | "pyiceberg"
   | "ibis"
   | "motherduck"
-  | "redshift";
+  | "redshift"
+  | "databricks";
 
 export const ConnectionDisplayNames: Record<ConnectionLibrary, string> = {
   sqlmodel: "SQLModel",
@@ -26,6 +27,7 @@ export const ConnectionDisplayNames: Record<ConnectionLibrary, string> = {
   ibis: "Ibis",
   motherduck: "MotherDuck",
   redshift: "Redshift",
+  databricks: "Databricks",
 };
 
 abstract class CodeGenerator<T extends DatabaseConnection["type"]> {
@@ -55,6 +57,9 @@ abstract class CodeGenerator<T extends DatabaseConnection["type"]> {
       case "duckdb":
         imports.add("import duckdb");
         break;
+      case "ibis":
+        imports.add("import ibis");
+        break;
     }
     return imports;
   }
@@ -138,10 +143,38 @@ class SecretContainer {
     return `{${value}}`;
   }
 
+  /**
+   * Generate a password variable for connection strings, supporting inline values,
+   * environment variable lookups, and f-string formatting.
+   * @param {string|undefined} password - Fallback password value.
+   * @param {string} passwordPlaceholder - Environment variable name.
+   * @param {boolean} inFString - If true, wrap for Python f-string.
+   * @param {string} [variableName] - Variable name (default: "password").
+   * @returns {string}
+   *
+   * @example
+   * `printPassword("token123", "DATABRICKS_ACCESS_TOKEN", true, "access_token")`
+   *
+   * Returns:
+   * ```python
+   * _access_token = os.environ.get("DATABRICKS_ACCESS_TOKEN", "token123")
+   * f"db://sample:{_access_token}@sample.com"
+   * ```
+   *
+   * @example
+   * `printPassword("token123", "DATABRICKS_ACCESS_TOKEN", false, "access_token")`
+   *
+   * Returns:
+   * ```python
+   * access_token = os.environ.get("DATABRICKS_ACCESS_TOKEN", "token123")
+   * f"db://sample:access_token@sample.com"
+   * ```
+   */
   printPassword(
     password: string | undefined,
     passwordPlaceholder: string,
     inFString: boolean,
+    variableName?: string,
   ): string {
     // Inline passwords should use printInFString, otherwise use print
     const printMethod = inFString
@@ -149,8 +182,8 @@ class SecretContainer {
       : this.print.bind(this);
 
     return isSecret(password)
-      ? printMethod("password", password)
-      : printMethod("password", passwordPlaceholder, password);
+      ? printMethod(variableName || "password", password)
+      : printMethod(variableName || "password", passwordPlaceholder, password);
   }
 
   getSecrets(): Record<string, string> {
@@ -557,7 +590,7 @@ class PyIcebergGenerator extends CodeGenerator<"iceberg"> {
 class DataFusionGenerator extends CodeGenerator<"datafusion"> {
   generateImports(): string[] {
     // To trigger installation of ibis-datafusion
-    return ["import ibis", "from datafusion import SessionContext"];
+    return ["from datafusion import SessionContext"];
   }
 
   generateConnectionCode(): string {
@@ -578,7 +611,7 @@ class DataFusionGenerator extends CodeGenerator<"datafusion"> {
 
 class PySparkGenerator extends CodeGenerator<"pyspark"> {
   generateImports(): string[] {
-    return ["import ibis", "from pyspark.sql import SparkSession"];
+    return ["from pyspark.sql import SparkSession"];
   }
 
   generateConnectionCode(): string {
@@ -668,6 +701,70 @@ ${formatUrlParams(params, (inner) => `          ${inner}`)},
   }
 }
 
+class DatabricksGenerator extends CodeGenerator<"databricks"> {
+  generateImports(): string[] {
+    return [];
+  }
+
+  generateConnectionCode(): string {
+    const useFString = this.orm !== "ibis";
+
+    const accessToken = this.secrets.printPassword(
+      this.connection.access_token,
+      "DATABRICKS_ACCESS_TOKEN",
+      useFString,
+      "access_token",
+    );
+
+    const serverHostname = this.secrets.printPassword(
+      this.connection.server_hostname,
+      "DATABRICKS_SERVER_HOSTNAME",
+      useFString,
+      "server_hostname",
+    );
+
+    const httpPath = this.secrets.printPassword(
+      this.connection.http_path,
+      "DATABRICKS_HTTP_PATH",
+      useFString,
+      "http_path",
+    );
+
+    const catalog = this.connection.catalog
+      ? this.secrets.printInFString("catalog", this.connection.catalog)
+      : undefined;
+    const schema = this.connection.schema
+      ? this.secrets.printInFString("schema", this.connection.schema)
+      : undefined;
+
+    let BASE_URL = `databricks://token:${accessToken}@${serverHostname}?http_path=${httpPath}`;
+    if (catalog) {
+      BASE_URL += `&catalog=${catalog}`;
+    }
+    if (schema) {
+      BASE_URL += `&schema=${schema}`;
+    }
+
+    if (this.orm === "ibis") {
+      const catalogParam = catalog ? `\n          catalog=${catalog},` : "";
+      const schemaParam = schema ? `\n          schema=${schema},` : "";
+
+      return dedent(`
+        engine = ibis.databricks.connect(
+          server_hostname=${serverHostname},
+          http_path=${httpPath},${catalogParam}${schemaParam}
+          access_token=${accessToken}
+        )
+      `);
+    }
+
+    return dedent(`
+      DATABASE_URL = f"${BASE_URL}"
+      engine = ${this.orm}.create_engine(DATABASE_URL)
+      `);
+  }
+}
+
 class CodeGeneratorFactory {
   public secrets = new SecretContainer();
 
@@ -706,6 +803,8 @@ class CodeGeneratorFactory {
         return new PySparkGenerator(connection, orm, this.secrets);
       case "redshift":
         return new RedshiftGenerator(connection, orm, this.secrets);
+      case "databricks":
+        return new DatabricksGenerator(connection, orm, this.secrets);
       default:
         assertNever(connection);
     }

package/src/components/editor/database/schemas.ts

@@ -16,18 +16,16 @@ function passwordField() {
     );
 }
 
-function tokenField() {
-  return z
-    .string()
-    .optional()
-    .describe(
-      FieldOptions.of({
-        label: "Token",
-        inputType: "password",
-        placeholder: "token",
-        optionRegex: ".*token.*",
-      }),
-    );
+function tokenField(label?: string, required?: boolean) {
+  const field = z.string().describe(
+    FieldOptions.of({
+      label: label || "Token",
+      inputType: "password",
+      placeholder: "token",
+      optionRegex: ".*token.*",
+    }),
+  );
+  return required ? field.nonempty() : field.optional();
 }
 
 function warehouseNameField() {
@@ -43,20 +41,22 @@ function warehouseNameField() {
     );
 }
 
-function uriField() {
-  return z
+function uriField(label?: string, required?: boolean) {
+  const field = z
     .string()
-    .optional()
-    .describe(FieldOptions.of({ label: "URI", optionRegex: ".*uri.*" }));
+    .describe(
+      FieldOptions.of({ label: label || "URI", optionRegex: ".*uri.*" }),
+    );
+  return required ? field.nonempty() : field.optional();
 }
 
-function hostField() {
+function hostField(label?: string) {
   return z
     .string()
     .nonempty()
     .describe(
       FieldOptions.of({
-        label: "Host",
+        label: label || "Host",
         placeholder: "localhost",
         optionRegex: ".*host.*",
       }),
@@ -464,6 +464,23 @@ export const RedshiftConnectionSchema = z
   })
   .describe(FieldOptions.of({ direction: "two-columns" }));
 
+export const DatabricksConnectionSchema = z
+  .object({
+    type: z.literal("databricks"),
+    access_token: tokenField("Access Token", true),
+    server_hostname: hostField("Server Hostname"),
+    http_path: uriField("HTTP Path", true),
+    catalog: z
+      .string()
+      .optional()
+      .describe(FieldOptions.of({ label: "Catalog" })),
+    schema: z
+      .string()
+      .optional()
+      .describe(FieldOptions.of({ label: "Schema" })),
+  })
+  .describe(FieldOptions.of({ direction: "two-columns" }));
+
 export const DatabaseConnectionSchema = z.discriminatedUnion("type", [
   PostgresConnectionSchema,
   MySQLConnectionSchema,
@@ -480,6 +497,7 @@ export const DatabaseConnectionSchema = z.discriminatedUnion("type", [
   DataFusionConnectionSchema,
   PySparkConnectionSchema,
   RedshiftConnectionSchema,
+  DatabricksConnectionSchema,
 ]);
 
 export type DatabaseConnection = z.infer<typeof DatabaseConnectionSchema>;

package/src/components/editor/errors/auto-fix.tsx

@@ -1,24 +1,33 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
+import { useAtomValue, useSetAtom } from "jotai";
 import { WrenchIcon } from "lucide-react";
 import { Button } from "@/components/ui/button";
 import { Tooltip } from "@/components/ui/tooltip";
+import { aiCompletionCellAtom } from "@/core/ai/state";
 import { notebookAtom, useCellActions } from "@/core/cells/cells";
 import type { CellId } from "@/core/cells/ids";
+import { aiEnabledAtom } from "@/core/config/config";
 import { getAutoFixes } from "@/core/errors/errors";
 import type { MarimoError } from "@/core/kernel/messages";
 import { store } from "@/core/state/jotai";
+import { cn } from "@/utils/cn";
 
 export const AutoFixButton = ({
   errors,
   cellId,
+  className,
 }: {
   errors: MarimoError[];
   cellId: CellId;
   className?: string;
 }) => {
   const { createNewCell } = useCellActions();
-  const autoFixes = errors.flatMap((error) => getAutoFixes(error));
+  const aiEnabled = useAtomValue(aiEnabledAtom);
+  const autoFixes = errors.flatMap((error) =>
+    getAutoFixes(error, { aiEnabled }),
+  );
+  const setAiCompletionCell = useSetAtom(aiCompletionCellAtom);
 
   if (autoFixes.length === 0) {
     return null;
@@ -33,7 +42,7 @@ export const AutoFixButton = ({
       <Button
         size="xs"
         variant="outline"
-        className="my-2 font-normal"
+        className={cn("my-2 font-normal", className)}
         onClick={() => {
           const editorView =
             store.get(notebookAtom).cellHandles[cellId].current?.editorView;
@@ -48,6 +57,7 @@ export const AutoFixButton = ({
             },
             editor: editorView,
             cellId: cellId,
+            setAiCompletionCell,
           });
           // Focus the editor
           editorView?.focus();

package/src/components/editor/errors/sql-validation-errors.tsx

@@ -2,29 +2,35 @@
 
 import { AlertCircleIcon } from "lucide-react";
 import type { CellId } from "@/core/cells/ids";
-import { useSqlValidationErrorsForCell } from "@/core/codemirror/language/languages/sql/validation-errors";
+import { useSqlValidationErrorsForCell } from "@/core/codemirror/language/languages/sql/banner-validation-errors";
 
-export const SqlValidationErrorBanner = ({ cellId }: { cellId: CellId }) => {
+export const SqlValidationErrorBanner = ({
+  cellId,
+  hide,
+}: {
+  cellId: CellId;
+  hide: boolean;
+}) => {
   const error = useSqlValidationErrorsForCell(cellId);
 
-  if (!error) {
+  if (!error || hide) {
     return;
   }
 
   return (
-    <div className="p-3 text-sm flex flex-col text-muted-foreground gap-1.5 bg-destructive/5">
+    <div className="p-3 text-sm flex flex-col text-muted-foreground gap-1.5 bg-destructive/4">
       <div className="flex items-start gap-1.5">
         <AlertCircleIcon size={13} className="mt-[3px] text-destructive" />
         <p>
           <span className="font-bold text-destructive">{error.errorType}:</span>{" "}
-          {error.errorMessage}
+          <span className="whitespace-pre-wrap">{error.errorMessage}</span>
         </p>
       </div>
 
       {error.codeblock && (
         <pre
           lang="sql"
-          className="text-xs bg-muted rounded p-2 pb-0 mx-3 overflow-x-auto font-mono whitespace-pre-wrap"
+          className="text-xs bg-muted/80 rounded p-2 pb-0 mx-3 font-medium whitespace-pre-wrap"
         >
           {error.codeblock}
         </pre>

package/src/components/editor/navigation/clipboard.ts

@@ -11,9 +11,9 @@ import { Logger } from "@/utils/Logger";
 const MARIMO_CELL_MIMETYPE = "web application/x-marimo-cell";
 
 interface ClipboardCellData {
-  cells: Array<{
+  cells: {
     code: string;
-  }>;
+  }[];
   version: "1.0";
 }
 

package/src/components/editor/output/ConsoleOutput.tsx

@@ -28,7 +28,7 @@ interface Props {
   cellId: CellId;
   cellName: string;
   className?: string;
-  consoleOutputs: Array<WithResponse<OutputMessage>>;
+  consoleOutputs: WithResponse<OutputMessage>[];
   stale: boolean;
   debuggerActive: boolean;
   onRefactorWithAI?: (opts: { prompt: string }) => void;

package/src/components/editor/output/JsonOutput.tsx

@@ -200,7 +200,7 @@ const LEAF_RENDERERS = {
 };
 
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
-const MIME_TYPES: Array<DataType<any>> = Object.entries(LEAF_RENDERERS).map(
+const MIME_TYPES: DataType<any>[] = Object.entries(LEAF_RENDERERS).map(
   ([leafType, render]) => ({
     is: (value) => typeof value === "string" && value.startsWith(leafType),
     PostComponent: PyCopyButton,

package/src/components/editor/output/MarimoErrorOutput.tsx

@@ -77,7 +77,7 @@ export const MarimoErrorOutput = ({
     alertVariant = "default";
     titleColor = "text-secondary-foreground";
   } else if (errors.some((e) => e.type === "sql-error")) {
-    titleContents = "SQL Error in statement";
+    titleContents = "SQL error";
   } else {
     // Check for exception type
     const exceptionError = errors.find((e) => e.type === "exception");
@@ -500,25 +500,30 @@ export const MarimoErrorOutput = ({
         <div key="sql-errors">
           {sqlErrors.map((error, idx) => {
             const line =
-              error.sql_line != null ? (error?.sql_line | 0) + 1 : null;
-            const col = error.sql_col != null ? (error?.sql_col | 0) + 1 : null;
+              error.sql_line == null ? null : Math.trunc(error?.sql_line) + 1;
+            const col =
+              error.sql_col == null ? null : Math.trunc(error?.sql_col) + 1;
             return (
-              <div key={`sql-error-${idx}`} className="space-y-2">
-                <p className="text-muted-foreground">{error.msg}</p>
+              <div key={`sql-error-${idx}`} className="space-y-2 mt-2">
+                <p className="text-muted-foreground font-medium">{error.msg}</p>
                 {error.hint && (
                   <div className="flex items-start gap-2">
-                    <InfoIcon className="h-4 w-4 text-muted-foreground mt-0.5 flex-shrink-0" />
-                    <pre className="whitespace-pre-wrap text-sm text-muted-foreground">
+                    <InfoIcon
+                      size={11}
+                      className="text-muted-foreground mt-1 flex-shrink-0"
+                    />
+                    <p className="whitespace-pre-wrap text-sm text-muted-foreground">
                       {error.hint}
-                    </pre>
+                    </p>
                   </div>
                 )}
                 {error.sql_statement && (
-                  <div className="bg-muted/50 p-2 rounded text-xs font-mono">
-                    <pre className="whitespace-pre-wrap">
-                      {error.sql_statement}
-                    </pre>
-                  </div>
+                  <pre
+                    lang="sql"
+                    className="text-xs bg-muted/80 rounded whitespace-pre-wrap p-3.5"
+                  >
+                    {error.sql_statement.trim()}
+                  </pre>
                 )}
                 {line !== null && col !== null && (
                   <p className="text-xs text-muted-foreground">
@@ -528,18 +533,13 @@ export const MarimoErrorOutput = ({
               </div>
             );
           })}
-          {cellId && <AutoFixButton errors={sqlErrors} cellId={cellId} />}
-          <Tip title="How to fix SQL errors">
-            <p className="pb-2">
-              SQL parsing errors often occur due to invalid syntax, missing
-              keywords, or unsupported SQL features.
-            </p>
-            <p className="py-2">
-              Check your SQL syntax and ensure you're using supported SQL
-              dialect features. The error location can help you identify the
-              problematic part of your query.
-            </p>
-          </Tip>
+          {cellId && (
+            <AutoFixButton
+              errors={sqlErrors}
+              cellId={cellId}
+              className="mt-2.5"
+            />
+          )}
         </div>,
       );
     }

package/src/components/editor/output/MarimoTracebackOutput.tsx

@@ -108,7 +108,7 @@ export const MarimoTracebackOutput = ({
             Fix with AI
           </Button>
         )}
-        {tracebackInfo && !isWasm() && (
+        {tracebackInfo && tracebackInfo.kind === "cell" && !isWasm() && (
           <Tooltip content={"Attach pdb to the exception point."}>
             <Button
               size="xs"
@@ -180,7 +180,7 @@ function lastLine(text: string): string {
 
 export const replaceTracebackFilenames = (domNode: DOMNode) => {
   const info = getTracebackInfo(domNode);
-  if (info) {
+  if (info?.kind === "cell") {
     const tooltipContent = <InsertBreakpointContent />;
     return (
       <span className="nb">
@@ -211,6 +211,21 @@ export const replaceTracebackFilenames = (domNode: DOMNode) => {
       </span>
     );
   }
+  if (info?.kind === "file") {
+    return (
+      <div
+        className="inline-block cursor-pointer text-destructive hover:underline"
+        onClick={(_) => {
+          getRequestClient().openFile({
+            path: info.filePath,
+            lineNumber: info.lineNumber,
+          });
+        }}
+      >
+        <span className="nb">"{info.filePath}"</span>
+      </div>
+    );
+  }
 };
 
 export const replaceTracebackPrefix = (domNode: DOMNode) => {

package/src/components/editor/renderers/grid-layout/types.ts

@@ -68,13 +68,13 @@ export interface GridLayout extends Omit<SerializedGridLayout, "cells"> {
   /**
    * The cells in the layout.
    */
-  cells: Array<{
+  cells: {
     i: string;
     x: number;
     y: number;
     w: number;
     h: number;
-  }>;
+  }[];
 
   scrollableCells: Set<CellId>;
 

package/src/components/editor/renderers/plugins.ts

@@ -8,7 +8,7 @@ import { VerticalLayoutPlugin } from "./vertical-layout/vertical-layout";
 
 // If more renderers are added, we may want to consider lazy loading them.
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
-export const cellRendererPlugins: Array<ICellRendererPlugin<any, any>> = [
+export const cellRendererPlugins: ICellRendererPlugin<any, any>[] = [
   GridLayoutPlugin,
   SlidesLayoutPlugin,
   VerticalLayoutPlugin,

package/src/components/editor/renderers/types.ts

@@ -17,7 +17,7 @@ export interface ICellRendererProps<L> {
   /**
    * The cells to render.
    */
-  cells: Array<CellRuntimeState & CellData>;
+  cells: (CellRuntimeState & CellData)[];
 
   /**
    * The layout configuration.

package/src/components/editor/renderers/vertical-layout/vertical-layout.tsx

@@ -198,7 +198,7 @@ const ActionButtons: React.FC<{
     actions.push(
       <DropdownMenuItem
         onSelect={onToggleShowCode}
-        id="notebook-action-show-code"
+        data-testid="notebook-action-show-code"
         key="show-code"
       >
         <Code2Icon className="mr-2" size={14} strokeWidth={1.5} />
@@ -213,7 +213,7 @@ const ActionButtons: React.FC<{
     actions.push(
       <DropdownMenuItem
         onSelect={handleDownloadAsHTML}
-        id="notebook-action-download-html"
+        data-testid="notebook-action-download-html"
         key="download-html"
       >
         <FolderDownIcon className="mr-2" size={14} strokeWidth={1.5} />
@@ -222,7 +222,7 @@ const ActionButtons: React.FC<{
       <DropdownMenuSeparator key="download-html-separator" />,
       <DropdownMenuItem
         onSelect={handleDownloadAsPNG}
-        id="notebook-action-download-png"
+        data-testid="notebook-action-download-png"
         key="download-png"
       >
         <ImageIcon className="mr-2" size={14} strokeWidth={1.5} />
@@ -239,7 +239,7 @@ const ActionButtons: React.FC<{
   // as this may be used in custom css to hide/show the actions dropdown
   return (
     <div
-      id="notebook-actions-dropdown"
+      data-testid="notebook-actions-dropdown"
       className={cn(
         "right-0 top-0 z-50 m-4 no-print flex gap-2 print:hidden",
         // If the notebook is static, we have a banner at the top, so
@@ -419,10 +419,10 @@ export const VerticalLayoutPlugin: ICellRendererPlugin<
 };
 
 export function groupCellsByColumn(
-  cells: Array<CellRuntimeState & CellData>,
-): Array<[number, Array<CellRuntimeState & CellData>]> {
+  cells: (CellRuntimeState & CellData)[],
+): [number, (CellRuntimeState & CellData)[]][] {
   // Group cells by column
-  const cellsByColumn = new Map<number, Array<CellRuntimeState & CellData>>();
+  const cellsByColumn = new Map<number, (CellRuntimeState & CellData)[]>();
   let lastSeenColumn = 0;
   cells.forEach((cell) => {
     const column = cell.config.column ?? lastSeenColumn;

package/src/components/forms/form.tsx

@@ -59,7 +59,7 @@ interface Props<T extends FieldValues> {
   form: UseFormReturn<T>;
   schema: z.ZodType;
   path?: Path<T>;
-  renderers: Array<FormRenderer<T>> | undefined;
+  renderers: FormRenderer<T>[] | undefined;
   children?: React.ReactNode;
 }
 
@@ -82,7 +82,7 @@ export function renderZodSchema<T extends FieldValues, S>(
   schema: z.ZodType<S>,
   form: UseFormReturn<T>,
   path: Path<T>,
-  renderers: Array<FormRenderer<T>>,
+  renderers: FormRenderer<T>[],
 ) {
   // Try custom renderers first
   for (const renderer of renderers) {
@@ -330,7 +330,7 @@ export function renderZodSchema<T extends FieldValues, S>(
   }
 
   if (schema instanceof z.ZodDiscriminatedUnion) {
-    const options = schema._def.options as Array<z.ZodType<unknown>>;
+    const options = schema._def.options as z.ZodType<unknown>[];
     const discriminator = schema._def.discriminator;
     const optionsMap = schema._def.optionsMap;
     return (
@@ -392,7 +392,7 @@ export function renderZodSchema<T extends FieldValues, S>(
         control={form.control}
         name={path}
         render={({ field }) => {
-          const options = schema._def.options as Array<z.ZodType<unknown>>;
+          const options = schema._def.options as z.ZodType<unknown>[];
           let value: string = field.value;
           const types = options.map((option) => {
             return getUnionLiteral(option)._def.value;
@@ -812,6 +812,6 @@ const MultiSelectFormField = ({
   );
 };
 
-function joinPath<T>(...parts: Array<string | number>): Path<T> {
+function joinPath<T>(...parts: (string | number)[]): Path<T> {
   return parts.filter((part) => part !== "").join(".") as Path<T>;
 }

package/src/components/ui/links.tsx

@@ -8,6 +8,7 @@ export const ExternalLink = ({
     | `https://console.anthropic.com/${string}`
     | `https://aistudio.google.com/${string}`
     | `https://github.com/${string}`
+    | `https://openrouter.ai/${string}`
     | `https://docs.marimo.io/${string}`
     | `https://docs.python.org/${string}`
     | `https://marimo.io/${string}`