@marimo-team/frontend 0.16.1 → 0.16.3

package/src/components/editor/Cell.tsx

@@ -80,6 +80,7 @@ import { useDeleteCellCallback } from "./cell/useDeleteCell";
 import { useRunCell } from "./cell/useRunCells";
 import { HideCodeButton } from "./code/readonly-python-code";
 import { cellDomProps } from "./common";
+import { SqlValidationErrorBanner } from "./errors/sql-validation-errors";
 import { useCellNavigationProps } from "./navigation/navigation";
 import {
   useTemporarilyShownCode,
@@ -450,6 +451,7 @@ const EditableCellComponent = ({
   });
   const canCollapse = canCollapseOutline(cellRuntime.outline);
   const hasOutput = !isOutputEmpty(cellRuntime.output);
+  const isStaleCell = outputIsStale(cellRuntime, cellData.edited);
   const hasConsoleOutput = cellRuntime.consoleOutputs.length > 0;
   const cellOutput = userConfig.display.cell_output;
 
@@ -510,7 +512,7 @@ const EditableCellComponent = ({
         className="output-area"
         cellId={cellId}
         output={cellRuntime.output}
-        stale={outputIsStale(cellRuntime, cellData.edited)}
+        stale={isStaleCell}
         loading={outputIsLoading(cellRuntime.status)}
       />
       {isMarkdownCodeHidden &&
@@ -653,6 +655,10 @@ const EditableCellComponent = ({
                 )}
               </div>
             </div>
+            <SqlValidationErrorBanner
+              cellId={cellId}
+              hide={cellRuntime.errored && !isStaleCell}
+            />
             {cellOutput === "below" && outputArea}
             {cellRuntime.serialization && (
               <div className="py-1 px-2 flex items-center justify-end gap-2 last:rounded-b">

package/src/components/editor/ai/transport/chat-transport.tsx

@@ -16,7 +16,7 @@ export class StreamingChunkTransport<
   private onChunkReceived: (chunk: UIMessageChunk) => void;
 
   constructor(
-    options: HttpChatTransportInitOptions<UI_MESSAGE> = {},
+    options: HttpChatTransportInitOptions<UI_MESSAGE>,
     onChunkReceived: (chunk: UIMessageChunk) => void,
   ) {
     super(options);

package/src/components/editor/chrome/panels/outline/useActiveOutline.tsx

@@ -18,7 +18,7 @@ function getRootScrollableElement() {
  * React hook to find the active header in the outline
  */
 export function useActiveOutline(
-  headerElements: Array<readonly [HTMLElement, string]>,
+  headerElements: (readonly [HTMLElement, string])[],
 ) {
   const [activeHeaderId, setActiveHeaderId] = useState<string | undefined>(
     undefined,

package/src/components/editor/chrome/panels/packages-panel.tsx

@@ -300,7 +300,7 @@ const InstallPackageForm: React.FC<{
 
 const PackagesList: React.FC<{
   onSuccess: () => void;
-  packages: Array<{ name: string; version: string }>;
+  packages: { name: string; version: string }[];
 }> = ({ onSuccess, packages }) => {
   if (packages.length === 0) {
     return (

package/src/components/editor/columns/storage.ts

@@ -7,7 +7,7 @@ import { NotebookScopedLocalStorage } from "@/utils/localStorage";
 const BASE_KEY = "marimo:notebook-col-sizes";
 
 interface ColumnSizes {
-  widths: Array<number | "contentWidth">;
+  widths: (number | "contentWidth")[];
 }
 
 function initialState(): ColumnSizes {

package/src/components/editor/database/__tests__/__snapshots__/as-code.test.ts.snap

@@ -28,6 +28,42 @@ engine = clickhouse_connect.get_client(
 )"
 `;
 
+exports[`generateDatabaseCode > basic connections > databricks 1`] = `
+"import os
+import sqlalchemy
+
+_access_token = os.environ.get("DATABRICKS_ACCESS_TOKEN", "my_access_token")
+_server_hostname = os.environ.get("DATABRICKS_SERVER_HOSTNAME", "localhost")
+_http_path = os.environ.get("DATABRICKS_HTTP_PATH", "http://localhost:8080")
+DATABASE_URL = f"databricks://token:{_access_token}@{_server_hostname}?http_path={_http_path}"
+engine = sqlalchemy.create_engine(DATABASE_URL)"
+`;
+
+exports[`generateDatabaseCode > basic connections > databricks with catalog and schema 1`] = `
+"import os
+import sqlalchemy
+
+_access_token = os.environ.get("DATABRICKS_ACCESS_TOKEN", "my_access_token")
+_server_hostname = os.environ.get("DATABRICKS_SERVER_HOSTNAME", "localhost")
+_http_path = os.environ.get("DATABRICKS_HTTP_PATH", "http://localhost:8080")
+DATABASE_URL = f"databricks://token:{_access_token}@{_server_hostname}?http_path={_http_path}&catalog=my_catalog&schema=my_schema"
+engine = sqlalchemy.create_engine(DATABASE_URL)"
+`;
+
+exports[`generateDatabaseCode > basic connections > databricks with ibis 1`] = `
+"import ibis
+import os
+
+_access_token = os.environ.get("DATABRICKS_ACCESS_TOKEN", "my_access_token")
+_server_hostname = os.environ.get("DATABRICKS_SERVER_HOSTNAME", "localhost")
+_http_path = os.environ.get("DATABRICKS_HTTP_PATH", "http://localhost:8080")
+engine = ibis.databricks.connect(
+  server_hostname=_server_hostname,
+  http_path=_http_path,
+  access_token=_access_token
+)"
+`;
+
 exports[`generateDatabaseCode > basic connections > datafusion 1`] = `
 "from datafusion import SessionContext
 import ibis

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

@@ -192,8 +192,24 @@ describe("generateDatabaseCode", () => {
     },
   };
 
+  const databricksConnection: DatabaseConnection = {
+    type: "databricks",
+    access_token: "my_access_token",
+    server_hostname: "localhost",
+    http_path: "http://localhost:8080",
+  };
+
+  const databricksWithCatalogSchema: DatabaseConnection = {
+    type: "databricks",
+    access_token: "my_access_token",
+    server_hostname: "localhost",
+    http_path: "http://localhost:8080",
+    catalog: "my_catalog",
+    schema: "my_schema",
+  };
+
   describe("basic connections", () => {
-    const testCases: Array<[string, DatabaseConnection, ConnectionLibrary]> = [
+    const testCases: [string, DatabaseConnection, ConnectionLibrary][] = [
       ["postgres with SQLModel", basePostgres, "sqlmodel"],
       ["postgres with SQLAlchemy", basePostgres, "sqlalchemy"],
       ["mysql with SQLModel", baseMysql, "sqlmodel"],
@@ -218,9 +234,16 @@ describe("generateDatabaseCode", () => {
       ["pyspark with session", pysparkConnSession, "ibis"],
       ["redshift with DB credentials", redshiftDBConnection, "redshift"],
       ["redshift with IAM credentials", redshiftIAMConnection, "redshift"],
+      ["databricks", databricksConnection, "sqlalchemy"],
+      [
+        "databricks with catalog and schema",
+        databricksWithCatalogSchema,
+        "sqlalchemy",
+      ],
+      ["databricks with ibis", databricksConnection, "ibis"],
     ];
 
-    it.each(testCases)("%s", (name, connection, orm) => {
+    it.each(testCases)("%s", (_name, connection, orm) => {
       expect(generateDatabaseCode(connection, orm)).toMatchSnapshot();
     });
   });
@@ -292,7 +315,7 @@ describe("generateDatabaseCode", () => {
         },
         "duckdb",
       ],
-    ])("%s", (name, connection, orm) => {
+    ])("%s", (_name, connection, orm) => {
       expect(
         generateDatabaseCode(connection, orm as ConnectionLibrary),
       ).toMatchSnapshot();
@@ -300,7 +323,7 @@ describe("generateDatabaseCode", () => {
   });
 
   describe("edge cases", () => {
-    const testCases: Array<[string, DatabaseConnection, string]> = [
+    const testCases: [string, DatabaseConnection, string][] = [
       [
         "ENV with special chars SQLModel",
         {
@@ -504,7 +527,7 @@ describe("generateDatabaseCode", () => {
       ],
     ];
 
-    it.each(testCases)("%s", (name, connection, orm) => {
+    it.each(testCases)("%s", (_name, connection, orm) => {
       expect(
         generateDatabaseCode(connection, orm as ConnectionLibrary),
       ).toMatchSnapshot();
@@ -548,7 +571,7 @@ describe("generateDatabaseCode", () => {
           credentials_json: '{"type": "service_account", "project_id": "test"',
         },
       ],
-    ])("%s", (name, connection) => {
+    ])("%s", (_name, connection) => {
       expect(generateDatabaseCode(connection, "sqlmodel")).toMatchSnapshot();
       expect(generateDatabaseCode(connection, "sqlalchemy")).toMatchSnapshot();
     });
@@ -591,7 +614,7 @@ describe("generateDatabaseCode", () => {
             "sqlmodel",
           ),
       ],
-    ])("%s", (name, fn) => {
+    ])("%s", (_name, fn) => {
       expect(fn).toThrow();
     });
   });

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

@@ -0,0 +1,40 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import { AlertCircleIcon } from "lucide-react";
+import type { CellId } from "@/core/cells/ids";
+import { useSqlValidationErrorsForCell } from "@/core/codemirror/language/languages/sql/banner-validation-errors";
+
+export const SqlValidationErrorBanner = ({
+  cellId,
+  hide,
+}: {
+  cellId: CellId;
+  hide: boolean;
+}) => {
+  const error = useSqlValidationErrorsForCell(cellId);
+
+  if (!error || hide) {
+    return;
+  }
+
+  return (
+    <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>{" "}
+          <span className="whitespace-pre-wrap">{error.errorMessage}</span>
+        </p>
+      </div>
+
+      {error.codeblock && (
+        <pre
+          lang="sql"
+          className="text-xs bg-muted/80 rounded p-2 pb-0 mx-3 font-medium whitespace-pre-wrap"
+        >
+          {error.codeblock}
+        </pre>
+      )}
+    </div>
+  );
+};

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

@@ -14,6 +14,7 @@ import type { OutputMessage } from "@/core/kernel/messages";
 import { useSelectAllContent } from "@/hooks/useSelectAllContent";
 import { cn } from "@/utils/cn";
 import { copyToClipboard } from "@/utils/copy";
+import { ansiToPlainText, parseHtmlContent } from "@/utils/dom";
 import { invariant } from "@/utils/invariant";
 import { Strings } from "@/utils/strings";
 import { NameCellContentEditable } from "../actions/name-cell-input";
@@ -27,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;
@@ -124,7 +125,18 @@ const ConsoleOutputInternal = (props: Props): React.ReactNode => {
                 onClick={() => {
                   const text = reversedOutputs
                     .filter((output) => output.channel !== "pdb")
-                    .map((output) => Strings.asString(output.data))
+                    .map((output) => {
+                      // If starts with `<`, then assume it's HTML
+                      if (
+                        typeof output.data === "string" &&
+                        output.data.startsWith("<")
+                      ) {
+                        return parseHtmlContent(output.data);
+                      }
+
+                      // Otherwise, convert the ANSI to HTML, then parse as HTML
+                      return ansiToPlainText(Strings.asString(output.data));
+                    })
                     .join("\n");
                   void copyToClipboard(text);
                 }}

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,