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

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

@@ -0,0 +1,166 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+import { describe, expect, it } from "vitest";
+import { prefixSecret } from "../../secrets";
+import { generateStorageCode } from "../as-code";
+import type { StorageConnection } from "../schemas";
+
+describe("generateStorageCode", () => {
+  const baseS3: StorageConnection = {
+    type: "s3",
+    bucket: "my-bucket",
+    region: "us-east-1",
+    access_key_id: "AKIAIOSFODNN7EXAMPLE",
+    secret_access_key: "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY",
+    endpoint_url: undefined,
+  };
+
+  const baseGCS: StorageConnection = {
+    type: "gcs",
+    bucket: "my-bucket",
+    service_account_key: '{"type": "service_account", "project_id": "test"}',
+  };
+
+  const baseAzure: StorageConnection = {
+    type: "azure",
+    container: "my-container",
+    account_name: "storageaccount",
+    account_key: "base64accountkey==",
+  };
+
+  const baseGDrive: StorageConnection = {
+    type: "gdrive",
+    credentials_json:
+      '{"type": "service_account", "client_email": "test@test.iam.gserviceaccount.com"}',
+  };
+
+  describe("S3", () => {
+    it("basic connection with all fields", () => {
+      expect(generateStorageCode(baseS3, "obstore")).toMatchSnapshot();
+    });
+
+    it("minimal connection (bucket only)", () => {
+      const conn: StorageConnection = {
+        type: "s3",
+        bucket: "my-bucket",
+      };
+      expect(generateStorageCode(conn, "obstore")).toMatchSnapshot();
+    });
+
+    it("with custom endpoint", () => {
+      const conn: StorageConnection = {
+        ...baseS3,
+        endpoint_url: "https://minio.example.com:9000",
+      };
+      expect(generateStorageCode(conn, "obstore")).toMatchSnapshot();
+    });
+
+    it("with secrets", () => {
+      const conn: StorageConnection = {
+        type: "s3",
+        bucket: "my-bucket",
+        region: "us-east-1",
+        access_key_id: prefixSecret("AWS_ACCESS_KEY_ID"),
+        secret_access_key: prefixSecret("AWS_SECRET_ACCESS_KEY"),
+      };
+      expect(generateStorageCode(conn, "obstore")).toMatchSnapshot();
+    });
+  });
+
+  describe("GCS", () => {
+    it("with service account key", () => {
+      expect(generateStorageCode(baseGCS, "obstore")).toMatchSnapshot();
+    });
+
+    it("without service account key (default credentials)", () => {
+      const conn: StorageConnection = {
+        type: "gcs",
+        bucket: "my-bucket",
+      };
+      expect(generateStorageCode(conn, "obstore")).toMatchSnapshot();
+    });
+  });
+
+  describe("Azure", () => {
+    it("basic connection with account key", () => {
+      expect(generateStorageCode(baseAzure, "obstore")).toMatchSnapshot();
+    });
+
+    it("without account key", () => {
+      const conn: StorageConnection = {
+        type: "azure",
+        container: "my-container",
+        account_name: "storageaccount",
+      };
+      expect(generateStorageCode(conn, "obstore")).toMatchSnapshot();
+    });
+
+    it("with secrets", () => {
+      const conn: StorageConnection = {
+        type: "azure",
+        container: "my-container",
+        account_name: prefixSecret("AZURE_ACCOUNT"),
+        account_key: prefixSecret("AZURE_KEY"),
+      };
+      expect(generateStorageCode(conn, "obstore")).toMatchSnapshot();
+    });
+  });
+
+  describe("Google Drive", () => {
+    it("with service account credentials", () => {
+      expect(generateStorageCode(baseGDrive, "fsspec")).toMatchSnapshot();
+    });
+
+    it("with browser auth (no credentials)", () => {
+      const conn: StorageConnection = {
+        type: "gdrive",
+      };
+      expect(generateStorageCode(conn, "fsspec")).toMatchSnapshot();
+    });
+  });
+
+  describe("invalid cases", () => {
+    it("throws for empty S3 bucket", () => {
+      expect(() =>
+        generateStorageCode(
+          { type: "s3", bucket: "" } as StorageConnection,
+          "obstore",
+        ),
+      ).toThrow();
+    });
+
+    it("throws for empty GCS bucket", () => {
+      expect(() =>
+        generateStorageCode(
+          { type: "gcs", bucket: "" } as StorageConnection,
+          "obstore",
+        ),
+      ).toThrow();
+    });
+
+    it("throws for empty Azure container", () => {
+      expect(() =>
+        generateStorageCode(
+          {
+            type: "azure",
+            container: "",
+            account_name: "acct",
+          } as StorageConnection,
+          "obstore",
+        ),
+      ).toThrow();
+    });
+
+    it("throws for empty Azure account name", () => {
+      expect(() =>
+        generateStorageCode(
+          {
+            type: "azure",
+            container: "my-container",
+            account_name: "",
+          } as StorageConnection,
+          "obstore",
+        ),
+      ).toThrow();
+    });
+  });
+});

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

@@ -0,0 +1,135 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { useState } from "react";
+import type { FieldValues } from "react-hook-form";
+import type { z } from "zod";
+import { ProtocolIcon } from "@/components/storage/components";
+import { ConnectionForm, SelectorButton, SelectorGrid } from "../components";
+import {
+  generateStorageCode,
+  type StorageLibrary,
+  StorageLibraryDisplayNames,
+} from "./as-code";
+import {
+  AzureStorageSchema,
+  GCSStorageSchema,
+  GoogleDriveStorageSchema,
+  S3StorageSchema,
+  type StorageConnection,
+} from "./schemas";
+
+interface StorageProviderSchema {
+  name: string;
+  schema: z.ZodType<StorageConnection, FieldValues>;
+  color: string;
+  protocol: string;
+  storageLibraries: {
+    libraries: StorageLibrary[];
+    preferred: StorageLibrary;
+  };
+}
+
+const STORAGE_PROVIDERS = [
+  {
+    name: "Amazon S3",
+    schema: S3StorageSchema,
+    color: "#232F3E",
+    protocol: "s3",
+    storageLibraries: {
+      libraries: ["obstore"],
+      preferred: "obstore",
+    },
+  },
+  {
+    name: "Google Cloud Storage",
+    schema: GCSStorageSchema,
+    color: "#4285F4",
+    protocol: "gcs",
+    storageLibraries: {
+      libraries: ["obstore"],
+      preferred: "obstore",
+    },
+  },
+  {
+    name: "Azure Blob Storage",
+    schema: AzureStorageSchema,
+    color: "#0062AD",
+    protocol: "azure",
+    storageLibraries: {
+      libraries: ["obstore"],
+      preferred: "obstore",
+    },
+  },
+  {
+    name: "Google Drive",
+    schema: GoogleDriveStorageSchema,
+    color: "#177834",
+    protocol: "gdrive",
+    storageLibraries: {
+      libraries: ["fsspec"],
+      preferred: "fsspec",
+    },
+  },
+] satisfies StorageProviderSchema[];
+
+const StorageProviderSelector: React.FC<{
+  onSelect: (schema: z.ZodType<StorageConnection, FieldValues>) => void;
+}> = ({ onSelect }) => {
+  return (
+    <SelectorGrid>
+      {STORAGE_PROVIDERS.map(({ name, schema, color, protocol }) => (
+        <SelectorButton
+          key={name}
+          name={name}
+          color={color}
+          icon={
+            <span className="w-8 h-8 flex items-center justify-center">
+              <ProtocolIcon
+                protocol={protocol}
+                className="w-7 h-7 brightness-0 invert"
+              />
+            </span>
+          }
+          onSelect={() => onSelect(schema)}
+        />
+      ))}
+    </SelectorGrid>
+  );
+};
+
+export const AddStorageForm: React.FC<{
+  onSubmit: () => void;
+  header?: React.ReactNode;
+}> = ({ onSubmit, header }) => {
+  const [selectedSchema, setSelectedSchema] = useState<z.ZodType<
+    StorageConnection,
+    FieldValues
+  > | null>(null);
+
+  if (!selectedSchema) {
+    return (
+      <>
+        {header}
+        <div>
+          <StorageProviderSelector onSelect={setSelectedSchema} />
+        </div>
+      </>
+    );
+  }
+
+  const provider = STORAGE_PROVIDERS.find((p) => p.schema === selectedSchema);
+  const libs = provider?.storageLibraries;
+
+  return (
+    <ConnectionForm<StorageConnection, StorageLibrary>
+      schema={selectedSchema}
+      libraries={libs?.libraries ?? []}
+      preferredLibrary={libs?.preferred ?? "obstore"}
+      displayNames={StorageLibraryDisplayNames}
+      libraryLabel="Preferred storage library"
+      generateCode={(values, library) => generateStorageCode(values, library)}
+      onSubmit={onSubmit}
+      onBack={() => setSelectedSchema(null)}
+    />
+  );
+};

package/src/components/editor/connections/storage/as-code.ts

@@ -0,0 +1,188 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import dedent from "string-dedent";
+import { assertNever } from "@/utils/assertNever";
+import { isSecret, unprefixSecret } from "../secrets";
+import { type StorageConnection, StorageConnectionSchema } from "./schemas";
+
+export type StorageLibrary = "obstore" | "fsspec";
+
+export const StorageLibraryDisplayNames: Record<StorageLibrary, string> = {
+  obstore: "Obstore",
+  fsspec: "fsspec",
+};
+
+class SecretContainer {
+  private secrets: Record<string, string> = {};
+
+  get imports(): Set<string> {
+    if (Object.keys(this.secrets).length === 0) {
+      return new Set<string>();
+    }
+    return new Set<string>(["import os"]);
+  }
+
+  print(varName: string, value: string | undefined): string {
+    if (value === undefined || value === "") {
+      return "";
+    }
+    const privateVar = `_${varName}`;
+    if (isSecret(value)) {
+      const envVar = unprefixSecret(value);
+      this.secrets[privateVar] = `os.environ.get("${envVar}")`;
+      return privateVar;
+    }
+    return `"${value}"`;
+  }
+
+  formatSecrets(): string {
+    if (Object.keys(this.secrets).length === 0) {
+      return "";
+    }
+    return Object.entries(this.secrets)
+      .map(([k, v]) => `${k} = ${v}`)
+      .join("\n");
+  }
+}
+
+function generateS3Code(
+  connection: Extract<StorageConnection, { type: "s3" }>,
+  secrets: SecretContainer,
+): { imports: Set<string>; code: string } {
+  const bucket = secrets.print("bucket", connection.bucket);
+  const imports = new Set(["from obstore.store import S3Store"]);
+  const params: string[] = [];
+
+  if (connection.region) {
+    params.push(`    region=${secrets.print("region", connection.region)},`);
+  }
+  if (connection.access_key_id) {
+    params.push(
+      `    access_key_id=${secrets.print("access_key_id", connection.access_key_id)},`,
+    );
+  }
+  if (connection.secret_access_key) {
+    params.push(
+      `    secret_access_key=${secrets.print("secret_access_key", connection.secret_access_key)},`,
+    );
+  }
+  if (connection.endpoint_url) {
+    params.push(
+      `    endpoint_url=${secrets.print("endpoint_url", connection.endpoint_url)},`,
+    );
+  }
+
+  const paramsStr = params.length > 0 ? `\n${params.join("\n")}\n` : "";
+
+  const code = dedent(`
+    store = S3Store(${bucket},${paramsStr})
+  `);
+  return { imports, code };
+}
+
+function generateGCSCode(
+  connection: Extract<StorageConnection, { type: "gcs" }>,
+  secrets: SecretContainer,
+): { imports: Set<string>; code: string } {
+  const bucket = secrets.print("bucket", connection.bucket);
+  const imports = new Set(["from obstore.store import GCSStore"]);
+
+  let code: string;
+  if (connection.service_account_key) {
+    imports.add("import json");
+    code = dedent(`
+      _credentials = json.loads("""${connection.service_account_key}""")
+      store = GCSStore(${bucket},
+          service_account_key=_credentials,
+      )
+    `);
+  } else {
+    code = dedent(`
+      store = GCSStore(${bucket})
+    `);
+  }
+  return { imports, code };
+}
+
+function generateAzureCode(
+  connection: Extract<StorageConnection, { type: "azure" }>,
+  secrets: SecretContainer,
+): { imports: Set<string>; code: string } {
+  const container = secrets.print("container", connection.container);
+  const accountName = secrets.print("account_name", connection.account_name);
+  const imports = new Set(["from obstore.store import AzureStore"]);
+  const params: string[] = [`account_name=${accountName},`];
+
+  if (connection.account_key) {
+    params.push(
+      `account_key=${secrets.print("account_key", connection.account_key)},`,
+    );
+  }
+
+  const paramsStr = params.map((p) => `    ${p}`).join("\n");
+
+  const code = `store = AzureStore(${container},\n${paramsStr}\n)`;
+  return { imports, code };
+}
+
+function generateGDriveCode(
+  connection: Extract<StorageConnection, { type: "gdrive" }>,
+  secrets: SecretContainer,
+): { imports: Set<string>; code: string } {
+  const imports = new Set(["from gdrive_fsspec import GoogleDriveFileSystem"]);
+
+  if (connection.credentials_json) {
+    imports.add("import json");
+    const creds = secrets.print(
+      "credentials_json",
+      connection.credentials_json,
+    );
+    const code = dedent(`
+      _creds = json.loads("""${connection.credentials_json?.startsWith("ENV:") ? `{${creds}}` : connection.credentials_json}""")
+      fs = GoogleDriveFileSystem(creds=_creds, token="service_account")
+    `);
+    return { imports, code };
+  }
+
+  const code = dedent(`
+    fs = GoogleDriveFileSystem(token="browser")
+  `);
+  return { imports, code };
+}
+
+export function generateStorageCode(
+  connection: StorageConnection,
+  _library: StorageLibrary,
+): string {
+  StorageConnectionSchema.parse(connection);
+
+  const secrets = new SecretContainer();
+  let result: { imports: Set<string>; code: string };
+
+  switch (connection.type) {
+    case "s3":
+      result = generateS3Code(connection, secrets);
+      break;
+    case "gcs":
+      result = generateGCSCode(connection, secrets);
+      break;
+    case "azure":
+      result = generateAzureCode(connection, secrets);
+      break;
+    case "gdrive":
+      result = generateGDriveCode(connection, secrets);
+      break;
+    default:
+      assertNever(connection);
+  }
+
+  const allImports = new Set([...secrets.imports, ...result.imports]);
+  const lines = [...allImports].sort();
+  lines.push("");
+  const secretsStr = secrets.formatSecrets();
+  if (secretsStr) {
+    lines.push(secretsStr);
+  }
+  lines.push(result.code.trim());
+  return lines.join("\n");
+}

package/src/components/editor/connections/storage/schemas.ts

@@ -0,0 +1,141 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+import { z } from "zod";
+import { FieldOptions } from "@/components/forms/options";
+
+export const S3StorageSchema = z
+  .object({
+    type: z.literal("s3"),
+    bucket: z
+      .string()
+      .nonempty()
+      .describe(
+        FieldOptions.of({
+          label: "Bucket",
+          placeholder: "my-bucket",
+        }),
+      ),
+    region: z
+      .string()
+      .optional()
+      .describe(
+        FieldOptions.of({
+          label: "Region",
+          placeholder: "us-east-1",
+          optionRegex: ".*region.*",
+        }),
+      ),
+    access_key_id: z
+      .string()
+      .optional()
+      .describe(
+        FieldOptions.of({
+          label: "Access Key ID",
+          inputType: "password",
+          optionRegex: ".*access_key.*",
+        }),
+      ),
+    secret_access_key: z
+      .string()
+      .optional()
+      .describe(
+        FieldOptions.of({
+          label: "Secret Access Key",
+          inputType: "password",
+          optionRegex: ".*secret.*access.*",
+        }),
+      ),
+    endpoint_url: z
+      .string()
+      .optional()
+      .describe(
+        FieldOptions.of({
+          label: "Endpoint URL",
+          placeholder: "https://s3.amazonaws.com",
+        }),
+      ),
+  })
+  .describe(FieldOptions.of({ direction: "two-columns" }));
+
+export const GCSStorageSchema = z
+  .object({
+    type: z.literal("gcs"),
+    bucket: z
+      .string()
+      .nonempty()
+      .describe(
+        FieldOptions.of({
+          label: "Bucket",
+          placeholder: "my-bucket",
+        }),
+      ),
+    service_account_key: z
+      .string()
+      .optional()
+      .describe(
+        FieldOptions.of({
+          label: "Service Account Key (JSON)",
+          inputType: "textarea",
+        }),
+      ),
+  })
+  .describe(FieldOptions.of({ direction: "two-columns" }));
+
+export const AzureStorageSchema = z
+  .object({
+    type: z.literal("azure"),
+    container: z
+      .string()
+      .nonempty()
+      .describe(
+        FieldOptions.of({
+          label: "Container",
+          placeholder: "my-container",
+        }),
+      ),
+    account_name: z
+      .string()
+      .nonempty()
+      .describe(
+        FieldOptions.of({
+          label: "Account Name",
+          placeholder: "storageaccount",
+          optionRegex: ".*account.*",
+        }),
+      ),
+    account_key: z
+      .string()
+      .optional()
+      .describe(
+        FieldOptions.of({
+          label: "Account Key",
+          inputType: "password",
+          optionRegex: ".*azure.*key.*",
+        }),
+      ),
+  })
+  .describe(FieldOptions.of({ direction: "two-columns" }));
+
+export const GoogleDriveStorageSchema = z
+  .object({
+    type: z.literal("gdrive"),
+    credentials_json: z
+      .string()
+      .optional()
+      .describe(
+        FieldOptions.of({
+          label: "Service Account JSON",
+          description: "Leave empty to use browser-based authentication",
+          inputType: "textarea",
+        }),
+      ),
+  })
+  .describe(FieldOptions.of({ direction: "two-columns" }));
+
+export const StorageConnectionSchema = z.discriminatedUnion("type", [
+  S3StorageSchema,
+  GCSStorageSchema,
+  AzureStorageSchema,
+  GoogleDriveStorageSchema,
+]);
+
+export type StorageConnection = z.infer<typeof StorageConnectionSchema>;

package/src/components/storage/components.tsx

@@ -19,8 +19,10 @@ import {
   ImageIcon,
 } from "lucide-react";
 import GoogleCloudIcon from "@/components/databases/icons/google-cloud-storage.svg?inline";
+import GoogleDriveIcon from "@/components/databases/icons/google-drive.svg?inline";
 import type { KnownStorageProtocol } from "@/core/storage/types";
 import { useTheme } from "@/theme/useTheme";
+import { cn } from "@/utils/cn";
 
 export function renderFileIcon(name: string): React.ReactNode {
   const ext = name.split(".").pop()?.toLowerCase();
@@ -67,12 +69,14 @@ const PROTOCOL_ICONS: Record<KnownStorageProtocol, IconEntry> = {
   http: GlobeIcon,
   file: HardDriveIcon,
   "in-memory": DatabaseZapIcon,
+  gdrive: { src: GoogleDriveIcon },
   github: GithubIcon,
 };
 
 export const ProtocolIcon: React.FC<{
   protocol: KnownStorageProtocol | (string & {});
-}> = ({ protocol }) => {
+  className?: string;
+}> = ({ protocol, className }) => {
   const { theme } = useTheme();
   const entry =
     PROTOCOL_ICONS[protocol.toLowerCase() as KnownStorageProtocol] ??
@@ -80,9 +84,11 @@ export const ProtocolIcon: React.FC<{
 
   if ("src" in entry) {
     const src = theme === "dark" && entry.dark ? entry.dark : entry.src;
-    return <img src={src} alt={protocol} className="h-3.5 w-3.5" />;
+    return (
+      <img src={src} alt={protocol} className={cn("h-3.5 w-3.5", className)} />
+    );
   }
 
   const Icon = entry;
-  return <Icon className="h-3.5 w-3.5" />;
+  return <Icon className={cn("h-3.5 w-3.5", className)} />;
 };