@marimo-team/frontend 0.20.3-dev94 → 0.20.3-dev96

package/src/components/editor/connections/components.tsx

@@ -0,0 +1,177 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { zodResolver } from "@hookform/resolvers/zod";
+import React from "react";
+import { type DefaultValues, type FieldValues, useForm } from "react-hook-form";
+import type { z } from "zod";
+import { type FormRenderer, ZodForm } from "@/components/forms/form";
+import { getDefaults } from "@/components/forms/form-utils";
+import { Button } from "@/components/ui/button";
+import { FormErrorsBanner } from "@/components/ui/form";
+import {
+  Select,
+  SelectContent,
+  SelectGroup,
+  SelectItem,
+  SelectTrigger,
+  SelectValue,
+} from "@/components/ui/select";
+import { useCellActions } from "@/core/cells/cells";
+import { useLastFocusedCellId } from "@/core/cells/focus";
+import { ENV_RENDERER, SecretsProvider } from "./form-renderers";
+
+const RENDERERS: FormRenderer[] = [ENV_RENDERER];
+
+/**
+ * Grid layout for provider/database selector buttons.
+ */
+export const SelectorGrid: React.FC<{ children: React.ReactNode }> = ({
+  children,
+}) => <div className="grid grid-cols-2 md:grid-cols-3 gap-4">{children}</div>;
+
+/**
+ * A colored button tile for selecting a provider/database.
+ */
+export const SelectorButton: React.FC<{
+  name: string;
+  color: string;
+  icon: React.ReactNode;
+  onSelect: () => void;
+}> = ({ name, color, icon, onSelect }) => (
+  <button
+    type="button"
+    className="py-3 flex flex-col items-center justify-center gap-1 transition-all hover:scale-105 hover:brightness-110 rounded shadow-sm-solid hover:shadow-md-solid"
+    style={{ backgroundColor: color }}
+    onClick={onSelect}
+  >
+    {icon}
+    <span className="text-white font-medium text-lg">{name}</span>
+  </button>
+);
+
+/**
+ * Footer with Back/Add buttons and a library picker.
+ */
+export const ConnectionFormFooter = <L extends string>({
+  onBack,
+  isValid,
+  libraries,
+  preferredLibrary,
+  onLibraryChange,
+  displayNames,
+  libraryLabel = "Preferred library",
+}: {
+  onBack: () => void;
+  isValid: boolean;
+  libraries: L[];
+  preferredLibrary: L;
+  onLibraryChange: (library: L) => void;
+  displayNames: Record<L, string>;
+  libraryLabel?: string;
+}) => (
+  <div className="flex gap-2 justify-between">
+    <div className="flex gap-2">
+      <Button type="button" variant="outline" onClick={onBack}>
+        Back
+      </Button>
+      <Button type="submit" disabled={!isValid}>
+        Add
+      </Button>
+    </div>
+    <div>
+      <Select value={preferredLibrary} onValueChange={onLibraryChange}>
+        <div className="flex flex-col gap-1 items-end">
+          <SelectTrigger>
+            <SelectValue placeholder="Select a library" />
+          </SelectTrigger>
+          <span className="text-xs text-muted-foreground">{libraryLabel}</span>
+        </div>
+        <SelectContent>
+          <SelectGroup>
+            {libraries.map((library) => (
+              <SelectItem key={library} value={library}>
+                {displayNames[library] ?? library}
+              </SelectItem>
+            ))}
+          </SelectGroup>
+        </SelectContent>
+      </Select>
+    </div>
+  </div>
+);
+
+/**
+ * Returns a callback that inserts code into a new cell after the last focused cell.
+ */
+export function useInsertCode() {
+  const { createNewCell } = useCellActions();
+  const lastFocusedCellId = useLastFocusedCellId();
+
+  return (code: string) => {
+    createNewCell({
+      code,
+      before: false,
+      cellId: lastFocusedCellId ?? "__end__",
+      skipIfCodeExists: true,
+    });
+  };
+}
+
+/**
+ * Generic connection form: Zod-driven form with secrets support, a library
+ * picker, and Back/Add buttons. Used by both database and storage forms.
+ */
+export const ConnectionForm = <T extends FieldValues, L extends string>({
+  schema,
+  libraries,
+  preferredLibrary: initialPreferred,
+  displayNames,
+  libraryLabel,
+  generateCode,
+  onSubmit,
+  onBack,
+}: {
+  schema: z.ZodType<T, FieldValues>;
+  libraries: L[];
+  preferredLibrary: L;
+  displayNames: Record<L, string>;
+  libraryLabel?: string;
+  generateCode: (values: T, library: L) => string;
+  onSubmit: () => void;
+  onBack: () => void;
+}) => {
+  const defaults = getDefaults(schema);
+  const form = useForm<T>({
+    defaultValues: defaults as DefaultValues<T>,
+    resolver: zodResolver(schema),
+    reValidateMode: "onChange",
+  });
+
+  const [preferredLibrary, setPreferredLibrary] =
+    React.useState<L>(initialPreferred);
+  const insertCode = useInsertCode();
+
+  const handleSubmit = (values: T) => {
+    insertCode(generateCode(values, preferredLibrary));
+    onSubmit();
+  };
+
+  return (
+    <form onSubmit={form.handleSubmit(handleSubmit)} className="space-y-4">
+      <SecretsProvider>
+        <ZodForm schema={schema} form={form} renderers={RENDERERS}>
+          <FormErrorsBanner />
+        </ZodForm>
+      </SecretsProvider>
+      <ConnectionFormFooter
+        onBack={onBack}
+        isValid={form.formState.isValid}
+        libraries={libraries}
+        preferredLibrary={preferredLibrary}
+        onLibraryChange={setPreferredLibrary}
+        displayNames={displayNames}
+        libraryLabel={libraryLabel}
+      />
+    </form>
+  );
+};

package/src/components/editor/{database → connections/database}/__tests__/as-code.test.ts

@@ -1,8 +1,8 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 import { describe, expect, it } from "vitest";
+import { prefixSecret } from "../../secrets";
 import { type ConnectionLibrary, generateDatabaseCode } from "../as-code";
 import type { DatabaseConnection } from "../schemas";
-import { prefixSecret } from "../secrets";
 
 describe("generateDatabaseCode", () => {
   // Test fixtures

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

@@ -0,0 +1,303 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { useState } from "react";
+import type { FieldValues } from "react-hook-form";
+import type { z } from "zod";
+import { DatabaseLogo, type DBLogoName } from "@/components/databases/icon";
+import { ConnectionForm, SelectorButton, SelectorGrid } from "../components";
+import {
+  ConnectionDisplayNames,
+  type ConnectionLibrary,
+  generateDatabaseCode,
+} from "./as-code";
+import {
+  BigQueryConnectionSchema,
+  ChdbConnectionSchema,
+  ClickhouseConnectionSchema,
+  type DatabaseConnection,
+  DatabricksConnectionSchema,
+  DataFusionConnectionSchema,
+  DuckDBConnectionSchema,
+  IcebergConnectionSchema,
+  MotherDuckConnectionSchema,
+  MySQLConnectionSchema,
+  PostgresConnectionSchema,
+  PySparkConnectionSchema,
+  RedshiftConnectionSchema,
+  SnowflakeConnectionSchema,
+  SQLiteConnectionSchema,
+  SupabaseConnectionSchema,
+  TimeplusConnectionSchema,
+  TrinoConnectionSchema,
+} from "./schemas";
+
+interface ConnectionSchema {
+  name: string;
+  schema: z.ZodType<DatabaseConnection, FieldValues>;
+  color: string;
+  logo: DBLogoName;
+  connectionLibraries: {
+    libraries: ConnectionLibrary[];
+    preferred: ConnectionLibrary;
+  };
+}
+
+const DATABASES = [
+  {
+    name: "PostgreSQL",
+    schema: PostgresConnectionSchema,
+    color: "#336791",
+    logo: "postgres",
+    connectionLibraries: {
+      libraries: ["sqlalchemy", "sqlmodel"],
+      preferred: "sqlalchemy",
+    },
+  },
+  {
+    name: "MySQL",
+    schema: MySQLConnectionSchema,
+    color: "#00758F",
+    logo: "mysql",
+    connectionLibraries: {
+      libraries: ["sqlalchemy", "sqlmodel"],
+      preferred: "sqlalchemy",
+    },
+  },
+  {
+    name: "SQLite",
+    schema: SQLiteConnectionSchema,
+    color: "#003B57",
+    logo: "sqlite",
+    connectionLibraries: {
+      libraries: ["sqlalchemy", "sqlmodel"],
+      preferred: "sqlalchemy",
+    },
+  },
+  {
+    name: "DuckDB",
+    schema: DuckDBConnectionSchema,
+    color: "#FFD700",
+    logo: "duckdb",
+    connectionLibraries: {
+      libraries: ["duckdb"],
+      preferred: "duckdb",
+    },
+  },
+  {
+    name: "MotherDuck",
+    schema: MotherDuckConnectionSchema,
+    color: "#ff9538",
+    logo: "motherduck",
+    connectionLibraries: {
+      libraries: ["duckdb"],
+      preferred: "duckdb",
+    },
+  },
+  {
+    name: "Snowflake",
+    schema: SnowflakeConnectionSchema,
+    color: "#29B5E8",
+    logo: "snowflake",
+    connectionLibraries: {
+      libraries: ["sqlalchemy", "sqlmodel"],
+      preferred: "sqlalchemy",
+    },
+  },
+  {
+    name: "ClickHouse",
+    schema: ClickhouseConnectionSchema,
+    color: "#2C2C1D",
+    logo: "clickhouse",
+    connectionLibraries: {
+      libraries: ["clickhouse_connect"],
+      preferred: "clickhouse_connect",
+    },
+  },
+  {
+    name: "Timeplus",
+    schema: TimeplusConnectionSchema,
+    color: "#B83280",
+    logo: "timeplus",
+    connectionLibraries: {
+      libraries: ["sqlalchemy", "sqlmodel"],
+      preferred: "sqlalchemy",
+    },
+  },
+  {
+    name: "BigQuery",
+    schema: BigQueryConnectionSchema,
+    color: "#4285F4",
+    logo: "bigquery",
+    connectionLibraries: {
+      libraries: ["sqlalchemy", "sqlmodel"],
+      preferred: "sqlalchemy",
+    },
+  },
+  {
+    name: "ClickHouse Embedded",
+    schema: ChdbConnectionSchema,
+    color: "#f2b611",
+    logo: "clickhouse",
+    connectionLibraries: {
+      libraries: ["chdb"],
+      preferred: "chdb",
+    },
+  },
+  {
+    name: "Trino",
+    schema: TrinoConnectionSchema,
+    color: "#d466b6",
+    logo: "trino",
+    connectionLibraries: {
+      libraries: ["sqlalchemy", "sqlmodel"],
+      preferred: "sqlalchemy",
+    },
+  },
+  {
+    name: "DataFusion",
+    schema: DataFusionConnectionSchema,
+    color: "#202A37",
+    logo: "datafusion",
+    connectionLibraries: {
+      libraries: ["ibis"],
+      preferred: "ibis",
+    },
+  },
+  {
+    name: "PySpark",
+    schema: PySparkConnectionSchema,
+    color: "#1C5162",
+    logo: "pyspark",
+    connectionLibraries: {
+      libraries: ["ibis"],
+      preferred: "ibis",
+    },
+  },
+  {
+    name: "Redshift",
+    schema: RedshiftConnectionSchema,
+    color: "#522BAE",
+    logo: "redshift",
+    connectionLibraries: {
+      libraries: ["redshift"],
+      preferred: "redshift",
+    },
+  },
+  {
+    name: "Databricks",
+    schema: DatabricksConnectionSchema,
+    color: "#c41e0c",
+    logo: "databricks",
+    connectionLibraries: {
+      libraries: ["sqlalchemy", "sqlmodel", "ibis"],
+      preferred: "sqlalchemy",
+    },
+  },
+  {
+    name: "Supabase",
+    schema: SupabaseConnectionSchema,
+    color: "#238F5F",
+    logo: "supabase",
+    connectionLibraries: {
+      libraries: ["sqlalchemy", "sqlmodel"],
+      preferred: "sqlalchemy",
+    },
+  },
+] satisfies ConnectionSchema[];
+
+const DATA_CATALOGS = [
+  {
+    name: "Iceberg",
+    schema: IcebergConnectionSchema,
+    color: "#000000",
+    logo: "iceberg",
+    connectionLibraries: {
+      libraries: ["pyiceberg"],
+      preferred: "pyiceberg",
+    },
+  },
+] satisfies ConnectionSchema[];
+
+const ALL_ENTRIES = [...DATABASES, ...DATA_CATALOGS];
+
+const DatabaseSchemaSelector: React.FC<{
+  onSelect: (schema: z.ZodType<DatabaseConnection, FieldValues>) => void;
+}> = ({ onSelect }) => {
+  return (
+    <>
+      <SelectorGrid>
+        {DATABASES.map(({ name, schema, color, logo }) => (
+          <SelectorButton
+            key={name}
+            name={name}
+            color={color}
+            icon={
+              <DatabaseLogo
+                name={logo}
+                className="w-8 h-8 text-white brightness-0 invert dark:invert"
+              />
+            }
+            onSelect={() => onSelect(schema)}
+          />
+        ))}
+      </SelectorGrid>
+      <h4 className="font-semibold text-muted-foreground text-lg flex items-center gap-4 my-2">
+        Data Catalogs
+        <hr className="flex-1" />
+      </h4>
+      <SelectorGrid>
+        {DATA_CATALOGS.map(({ name, schema, color, logo }) => (
+          <SelectorButton
+            key={name}
+            name={name}
+            color={color}
+            icon={
+              <DatabaseLogo
+                name={logo}
+                className="w-8 h-8 text-white brightness-0 invert dark:invert"
+              />
+            }
+            onSelect={() => onSelect(schema)}
+          />
+        ))}
+      </SelectorGrid>
+    </>
+  );
+};
+
+export const AddDatabaseForm: React.FC<{
+  onSubmit: () => void;
+  header?: React.ReactNode;
+}> = ({ onSubmit, header }) => {
+  const [selectedSchema, setSelectedSchema] = useState<z.ZodType<
+    DatabaseConnection,
+    FieldValues
+  > | null>(null);
+
+  if (!selectedSchema) {
+    return (
+      <>
+        {header}
+        <div>
+          <DatabaseSchemaSelector onSelect={setSelectedSchema} />
+        </div>
+      </>
+    );
+  }
+
+  const entry = ALL_ENTRIES.find((e) => e.schema === selectedSchema);
+  const libs = entry?.connectionLibraries;
+
+  return (
+    <ConnectionForm<DatabaseConnection, ConnectionLibrary>
+      schema={selectedSchema}
+      libraries={libs?.libraries ?? []}
+      preferredLibrary={libs?.preferred ?? "sqlalchemy"}
+      displayNames={ConnectionDisplayNames}
+      libraryLabel="Preferred connection library"
+      generateCode={(values, library) => generateDatabaseCode(values, library)}
+      onSubmit={onSubmit}
+      onBack={() => setSelectedSchema(null)}
+    />
+  );
+};

package/src/components/editor/{database → connections/database}/as-code.ts

@@ -2,8 +2,8 @@
 
 import dedent from "string-dedent";
 import { assertNever } from "@/utils/assertNever";
+import { isSecret, unprefixSecret } from "../secrets";
 import { type DatabaseConnection, DatabaseConnectionSchema } from "./schemas";
-import { isSecret, unprefixSecret } from "./secrets";
 
 export type ConnectionLibrary =
   | "sqlmodel"

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

@@ -0,0 +1,100 @@
+// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
+
+exports[`generateStorageCode > Azure > basic connection with account key 1`] = `
+"from obstore.store import AzureStore
+
+store = AzureStore("my-container",
+    account_name="storageaccount",
+    account_key="base64accountkey==",
+)"
+`;
+
+exports[`generateStorageCode > Azure > with secrets 1`] = `
+"from obstore.store import AzureStore
+import os
+
+_account_name = os.environ.get("AZURE_ACCOUNT")
+_account_key = os.environ.get("AZURE_KEY")
+store = AzureStore("my-container",
+    account_name=_account_name,
+    account_key=_account_key,
+)"
+`;
+
+exports[`generateStorageCode > Azure > without account key 1`] = `
+"from obstore.store import AzureStore
+
+store = AzureStore("my-container",
+    account_name="storageaccount",
+)"
+`;
+
+exports[`generateStorageCode > GCS > with service account key 1`] = `
+"from obstore.store import GCSStore
+import json
+
+_credentials = json.loads("""{"type": "service_account", "project_id": "test"}""")
+store = GCSStore("my-bucket",
+    service_account_key=_credentials,
+)"
+`;
+
+exports[`generateStorageCode > GCS > without service account key (default credentials) 1`] = `
+"from obstore.store import GCSStore
+
+store = GCSStore("my-bucket")"
+`;
+
+exports[`generateStorageCode > Google Drive > with browser auth (no credentials) 1`] = `
+"from gdrive_fsspec import GoogleDriveFileSystem
+
+fs = GoogleDriveFileSystem(token="browser")"
+`;
+
+exports[`generateStorageCode > Google Drive > with service account credentials 1`] = `
+"from gdrive_fsspec import GoogleDriveFileSystem
+import json
+
+_creds = json.loads("""{"type": "service_account", "client_email": "test@test.iam.gserviceaccount.com"}""")
+fs = GoogleDriveFileSystem(creds=_creds, token="service_account")"
+`;
+
+exports[`generateStorageCode > S3 > basic connection with all fields 1`] = `
+"from obstore.store import S3Store
+
+store = S3Store("my-bucket",
+    region="us-east-1",
+    access_key_id="AKIAIOSFODNN7EXAMPLE",
+    secret_access_key="wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY",
+)"
+`;
+
+exports[`generateStorageCode > S3 > minimal connection (bucket only) 1`] = `
+"from obstore.store import S3Store
+
+store = S3Store("my-bucket",)"
+`;
+
+exports[`generateStorageCode > S3 > with custom endpoint 1`] = `
+"from obstore.store import S3Store
+
+store = S3Store("my-bucket",
+    region="us-east-1",
+    access_key_id="AKIAIOSFODNN7EXAMPLE",
+    secret_access_key="wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY",
+    endpoint_url="https://minio.example.com:9000",
+)"
+`;
+
+exports[`generateStorageCode > S3 > with secrets 1`] = `
+"from obstore.store import S3Store
+import os
+
+_access_key_id = os.environ.get("AWS_ACCESS_KEY_ID")
+_secret_access_key = os.environ.get("AWS_SECRET_ACCESS_KEY")
+store = S3Store("my-bucket",
+    region="us-east-1",
+    access_key_id=_access_key_id,
+    secret_access_key=_secret_access_key,
+)"
+`;