@marimo-team/islands 0.24.1-dev3 → 0.24.1-dev30

package/src/core/datasets/__tests__/data-source-discovery.test.ts

@@ -0,0 +1,277 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import { connectionAtom } from "@/core/network/connection";
+import { store } from "@/core/state/jotai";
+import { WebSocketState } from "@/core/websocket/types";
+import {
+  fetchDataSourceDiscovery,
+  invalidateDataSourceDiscovery,
+  isDetectedSourceConnected,
+  matchesDiscoveryGroup,
+  type DetectedDataSource,
+} from "../data-source-discovery";
+import { DiscoverDataSources } from "../request-registry";
+
+function source(
+  overrides: Partial<DetectedDataSource> &
+    Pick<DetectedDataSource, "id" | "integration" | "category" | "displayName">,
+): DetectedDataSource {
+  return {
+    confidence: "high",
+    origins: [],
+    configuration: [],
+    code: "",
+    hidesWhen: { kind: "dialect", substrings: [overrides.integration] },
+    ...overrides,
+  };
+}
+
+describe("fetchDataSourceDiscovery", () => {
+  beforeEach(() => {
+    // Discovery waits for the session's websocket connection to be open
+    // before requesting, so it doesn't race a still-connecting session.
+    store.set(connectionAtom, { state: WebSocketState.OPEN });
+  });
+
+  afterEach(() => {
+    vi.restoreAllMocks();
+    invalidateDataSourceDiscovery();
+  });
+
+  it("shares a single in-flight request across concurrent callers", async () => {
+    const request = vi
+      .spyOn(DiscoverDataSources, "request")
+      .mockResolvedValue({ request_id: "request-id", sources: [] });
+
+    await Promise.all([
+      fetchDataSourceDiscovery(),
+      fetchDataSourceDiscovery(),
+      fetchDataSourceDiscovery(),
+    ]);
+
+    expect(request).toHaveBeenCalledTimes(1);
+  });
+
+  it("re-runs discovery after the cache is invalidated", async () => {
+    const request = vi
+      .spyOn(DiscoverDataSources, "request")
+      .mockResolvedValue({ request_id: "request-id", sources: [] });
+
+    await fetchDataSourceDiscovery();
+    invalidateDataSourceDiscovery();
+    await fetchDataSourceDiscovery();
+
+    expect(request).toHaveBeenCalledTimes(2);
+  });
+
+  it("clears the cache on failure so the next call retries", async () => {
+    const request = vi
+      .spyOn(DiscoverDataSources, "request")
+      .mockRejectedValueOnce(new Error("boom"))
+      .mockResolvedValueOnce({ request_id: "request-id", sources: [] });
+
+    await expect(fetchDataSourceDiscovery()).rejects.toThrow("boom");
+    await expect(fetchDataSourceDiscovery()).resolves.toEqual([]);
+    expect(request).toHaveBeenCalledTimes(2);
+  });
+
+  it("does not clear a newer cache entry when an invalidated request fails", async () => {
+    let rejectFirstRequest: ((error: Error) => void) | undefined;
+    const firstRequest = new Promise<{
+      request_id: string;
+      sources: DetectedDataSource[];
+    }>((_resolve, reject) => {
+      rejectFirstRequest = reject;
+    });
+    const request = vi
+      .spyOn(DiscoverDataSources, "request")
+      .mockReturnValueOnce(firstRequest)
+      .mockResolvedValueOnce({ request_id: "request-id", sources: [] });
+
+    const staleResult = fetchDataSourceDiscovery().catch((error) => error);
+    await vi.waitFor(() => expect(request).toHaveBeenCalledTimes(1));
+    invalidateDataSourceDiscovery();
+    await fetchDataSourceDiscovery();
+
+    rejectFirstRequest?.(new Error("stale request failed"));
+    await staleResult;
+    await fetchDataSourceDiscovery();
+
+    expect(request).toHaveBeenCalledTimes(2);
+  });
+
+  it("returns sources from the kernel", async () => {
+    const sources: DetectedDataSource[] = [
+      source({
+        id: "postgres",
+        integration: "postgres",
+        category: "database",
+        displayName: "PostgreSQL",
+      }),
+      source({
+        id: "mysql",
+        integration: "mysql",
+        category: "database",
+        displayName: "MySQL",
+      }),
+    ];
+    vi.spyOn(DiscoverDataSources, "request").mockResolvedValue({
+      request_id: "request-id",
+      sources,
+    });
+
+    await expect(fetchDataSourceDiscovery()).resolves.toEqual(sources);
+  });
+});
+
+describe("matchesDiscoveryGroup", () => {
+  const postgres = source({
+    id: "postgres",
+    integration: "postgres",
+    category: "database",
+    displayName: "PostgreSQL",
+  });
+  const iceberg = source({
+    id: "iceberg",
+    integration: "pyiceberg",
+    category: "catalog",
+    displayName: "PyIceberg",
+  });
+  const s3 = source({
+    id: "s3",
+    integration: "aws",
+    category: "object-storage",
+    displayName: "S3",
+    hidesWhen: {
+      kind: "storage",
+      protocols: ["s3"],
+      backendTypes: [],
+    },
+  });
+
+  it("keeps every source when no group is given", () => {
+    expect(matchesDiscoveryGroup(postgres)).toBe(true);
+    expect(matchesDiscoveryGroup(iceberg)).toBe(true);
+    expect(matchesDiscoveryGroup(s3)).toBe(true);
+  });
+
+  it("groups catalogs with databases", () => {
+    expect(matchesDiscoveryGroup(postgres, "database")).toBe(true);
+    expect(matchesDiscoveryGroup(iceberg, "database")).toBe(true);
+    expect(matchesDiscoveryGroup(s3, "database")).toBe(false);
+  });
+
+  it("narrows storage to object-storage sources", () => {
+    expect(matchesDiscoveryGroup(postgres, "storage")).toBe(false);
+    expect(matchesDiscoveryGroup(iceberg, "storage")).toBe(false);
+    expect(matchesDiscoveryGroup(s3, "storage")).toBe(true);
+  });
+});
+
+describe("isDetectedSourceConnected", () => {
+  const postgres = source({
+    id: "postgres",
+    integration: "postgres",
+    category: "database",
+    displayName: "PostgreSQL",
+    hidesWhen: { kind: "dialect", substrings: ["postgres"] },
+  });
+  const emptySnapshot = {
+    dialects: [] as string[],
+    storageProtocols: [] as string[],
+    storageBackendTypes: [] as string[],
+  };
+
+  it("does not hide a source when nothing is connected", () => {
+    expect(isDetectedSourceConnected(postgres, emptySnapshot)).toBe(false);
+  });
+
+  it("hides a database suggestion when a matching dialect is live", () => {
+    expect(
+      isDetectedSourceConnected(postgres, {
+        ...emptySnapshot,
+        dialects: ["postgresql"],
+      }),
+    ).toBe(true);
+  });
+
+  it("does not hide a database suggestion for a different dialect", () => {
+    expect(
+      isDetectedSourceConnected(postgres, {
+        ...emptySnapshot,
+        dialects: ["mysql", "duckdb"],
+      }),
+    ).toBe(false);
+  });
+
+  it("hides spark suggestions for either pyspark or spark dialects", () => {
+    const spark = source({
+      id: "spark",
+      integration: "pyspark",
+      category: "catalog",
+      displayName: "Spark",
+      hidesWhen: { kind: "dialect", substrings: ["pyspark", "spark"] },
+    });
+    expect(
+      isDetectedSourceConnected(spark, {
+        ...emptySnapshot,
+        dialects: ["spark"],
+      }),
+    ).toBe(true);
+  });
+
+  it("hides object storage when a matching protocol is live", () => {
+    const s3 = source({
+      id: "s3",
+      integration: "aws",
+      category: "object-storage",
+      displayName: "S3",
+      hidesWhen: {
+        kind: "storage",
+        protocols: ["s3", "cloudflare", "coreweave"],
+        backendTypes: [],
+      },
+    });
+    expect(
+      isDetectedSourceConnected(s3, {
+        ...emptySnapshot,
+        storageProtocols: ["s3"],
+      }),
+    ).toBe(true);
+  });
+
+  it("hides huggingface when the backend type matches", () => {
+    const hf = source({
+      id: "hf",
+      integration: "huggingface",
+      category: "object-storage",
+      displayName: "Hugging Face",
+      hidesWhen: {
+        kind: "storage",
+        protocols: ["hf"],
+        backendTypes: ["huggingface"],
+      },
+    });
+    expect(
+      isDetectedSourceConnected(hf, {
+        ...emptySnapshot,
+        storageBackendTypes: ["huggingface"],
+      }),
+    ).toBe(true);
+  });
+
+  it("does not hide when hidesWhen does not match live connections", () => {
+    expect(
+      isDetectedSourceConnected(
+        source({
+          id: "custom",
+          integration: "custom",
+          category: "database",
+          displayName: "Custom",
+          hidesWhen: { kind: "dialect", substrings: ["custom"] },
+        }),
+        { ...emptySnapshot, dialects: ["postgres"] },
+      ),
+    ).toBe(false);
+  });
+});

package/src/core/datasets/data-source-connections.ts

@@ -4,6 +4,7 @@ import { atom } from "jotai";
 import { isSchemaless } from "@/components/datasources/utils";
 import { createReducerAndAtoms } from "@/utils/createReducer";
 import { Logger } from "@/utils/Logger";
+import { sortBy } from "@/utils/arrays";
 import type {
   DatabaseSchema,
   DataSourceConnection as DataSourceConnectionType,
@@ -13,6 +14,7 @@ import { store } from "../state/jotai";
 import type { VariableName } from "../variables/types";
 import {
   type ConnectionName,
+  DEFAULT_DUCKDB_DATABASE,
   DUCKDB_ENGINE,
   INTERNAL_SQL_ENGINES,
 } from "./engines";
@@ -341,6 +343,40 @@ export const dataConnectionsMapAtom = atom(
   (get) => get(dataSourceConnectionsAtom).connectionsMap,
 );
 
+/**
+ * Get the data connections that are available to the user.
+ * Filters out internal engines if it has no databases or if it has only the in-memory database and no schemas.
+ */
+export const connectionsAtom = atom((get) => {
+  const dataConnections = new Map(get(dataConnectionsMapAtom));
+
+  // Filter out the internal engines if it has no databases
+  // Or if it has only the in-memory database and no schemas
+  for (const engine of INTERNAL_SQL_ENGINES) {
+    const connection = dataConnections.get(engine);
+    if (!connection) {
+      continue;
+    }
+
+    if (connection.databases.length === 0) {
+      dataConnections.delete(engine);
+    }
+
+    if (
+      connection.databases.length === 1 &&
+      connection.databases[0].name === DEFAULT_DUCKDB_DATABASE &&
+      connection.databases[0].schemas.length === 0
+    ) {
+      dataConnections.delete(engine);
+    }
+  }
+
+  // Put internal engines last to prioritize user-defined connections
+  return sortBy([...dataConnections.values()], (connection) =>
+    INTERNAL_SQL_ENGINES.has(connection.name) ? 1 : 0,
+  );
+});
+
 export function setLatestEngineSelected(engine: ConnectionName) {
   const existing = store.get(dataSourceConnectionsAtom);
   // Don't update the map if the engine is not in the map

package/src/core/datasets/data-source-discovery.ts

@@ -1,5 +1,97 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
+import { atom } from "jotai";
 import type { components } from "@marimo-team/marimo-api";
+import { waitForConnectionOpenIfNotebook } from "../network/connection";
+import { store } from "../state/jotai";
+import { DiscoverDataSources } from "./request-registry";
 
 export type DetectedDataSource = components["schemas"]["DetectedDataSource"];
+
+export type DataSourceDiscoveryGroup = "database" | "storage";
+
+/** Category filter for UI surfaces. Catalogs ride with databases. */
+export function matchesDiscoveryGroup(
+  source: DetectedDataSource,
+  group?: DataSourceDiscoveryGroup,
+): boolean {
+  if (!group) {
+    return true;
+  }
+  const isStorage = source.category === "object-storage";
+  return isStorage === (group === "storage");
+}
+
+/**
+ * Live engines and storage used to hide suggestions of the same type.
+ */
+export interface LiveConnectionSnapshot {
+  dialects: readonly string[];
+  storageProtocols: readonly string[];
+  storageBackendTypes: readonly string[];
+}
+
+/**
+ * Whether a detected source already has a live connection of the same type.
+ * Match rules come from the kernel on each suggestion.
+ */
+export function isDetectedSourceConnected(
+  source: DetectedDataSource,
+  snapshot: LiveConnectionSnapshot,
+): boolean {
+  const match = source.hidesWhen;
+  if (match.kind === "dialect") {
+    return match.substrings.some((alias) =>
+      snapshot.dialects.some((dialect) => dialect.includes(alias)),
+    );
+  }
+
+  if (
+    match.protocols.some((protocol) =>
+      snapshot.storageProtocols.includes(protocol),
+    )
+  ) {
+    return true;
+  }
+
+  return match.backendTypes.some((backendType) =>
+    snapshot.storageBackendTypes.includes(backendType),
+  );
+}
+
+/**
+ * Bumped when discovery should be re-fetched (kernel restart).
+ */
+export const dataSourceDiscoveryEpochAtom = atom(0);
+
+/**
+ * Kernel-managed datasource discovery scans the environment once per epoch.
+ * Every consumer shares a single in-flight request.
+ */
+let discoveryPromise: Promise<DetectedDataSource[]> | undefined;
+let discoveryEpoch = 0;
+
+export function fetchDataSourceDiscovery(): Promise<DetectedDataSource[]> {
+  const epoch = store.get(dataSourceDiscoveryEpochAtom);
+  if (!discoveryPromise || discoveryEpoch !== epoch) {
+    discoveryEpoch = epoch;
+    discoveryPromise = waitForConnectionOpenIfNotebook()
+      .then(() => DiscoverDataSources.request({}))
+      .then((result) => result.sources)
+      .catch((error: unknown) => {
+        if (discoveryEpoch === epoch) {
+          discoveryPromise = undefined;
+        }
+        throw error;
+      });
+  }
+  return discoveryPromise;
+}
+
+/**
+ * Clears the shared discovery cache and triggers consumers to refetch.
+ */
+export function invalidateDataSourceDiscovery(): void {
+  discoveryPromise = undefined;
+  store.set(dataSourceDiscoveryEpochAtom, (value) => value + 1);
+}

package/src/core/kernel/messages.ts

@@ -11,6 +11,7 @@ export const DATA_TYPES = [
   "date",
   "datetime",
   "time",
+  "geometry",
   "unknown",
 ] as const;
 export type ModelLifecycle = NotificationMessageData<"model-lifecycle">;

package/src/core/websocket/useMarimoKernelConnection.tsx

@@ -40,6 +40,7 @@ import type { CellData } from "../cells/types";
 import { capabilitiesAtom } from "../config/capabilities";
 import { connectionTransportTypeAtom, useSetAppConfig } from "../config/config";
 import { useDataSourceActions } from "../datasets/data-source-connections";
+import { invalidateDataSourceDiscovery } from "../datasets/data-source-discovery";
 import type { ConnectionName } from "../datasets/engines";
 import {
   DiscoverDataSources,
@@ -251,6 +252,12 @@ export function useMarimoKernelConnection(opts: {
           onError: showBoundary,
           existingCells,
         });
+        if (!msg.data.resumed) {
+          // A freshly started kernel may expose a different environment
+          // (e.g. new env vars); re-run discovery instead of reusing stale
+          // suggestions from the previous kernel instance.
+          invalidateDataSourceDiscovery();
+        }
         setKioskMode(msg.data.kiosk);
         // A freshly started kernel has no breakpoints of its own; re-push
         // the client's set so they still apply, and clear the stale
@@ -311,6 +318,7 @@ export function useMarimoKernelConnection(opts: {
       }
 
       case "variables":
+        const variableNames = msg.data.variables.map((v) => v.name);
         setVariables(
           msg.data.variables.map((v) => ({
             name: v.name,
@@ -318,9 +326,9 @@ export function useMarimoKernelConnection(opts: {
             usedBy: v.used_by,
           })),
         );
-        filterDatasetsFromVariables(msg.data.variables.map((v) => v.name));
-        filterDataSourcesFromVariables(msg.data.variables.map((v) => v.name));
-        filterStorageFromVariables(msg.data.variables.map((v) => v.name));
+        filterDatasetsFromVariables(variableNames);
+        filterDataSourcesFromVariables(variableNames);
+        filterStorageFromVariables(variableNames);
         return;
       case "variable-values":
         setMetadata(

package/src/hooks/useDataSourceDiscovery.ts

@@ -1,18 +1,59 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-import type { DetectedDataSource } from "@/core/datasets/data-source-discovery";
-import { DiscoverDataSources } from "@/core/datasets/request-registry";
+import { useAtomValue } from "jotai";
+import { useMemo } from "react";
+import { dataConnectionsMapAtom } from "@/core/datasets/data-source-connections";
+import {
+  dataSourceDiscoveryEpochAtom,
+  type DataSourceDiscoveryGroup,
+  type DetectedDataSource,
+  fetchDataSourceDiscovery,
+  isDetectedSourceConnected,
+  matchesDiscoveryGroup,
+} from "@/core/datasets/data-source-discovery";
+import { storageNamespacesAtom } from "@/core/storage/state";
 import { useAsyncData } from "./useAsyncData";
 
-export async function loadDataSourceDiscovery(): Promise<DetectedDataSource[]> {
-  const result = await DiscoverDataSources.request({});
-  return result.sources;
-}
-
 /**
  * Reusable UI-facing hook for kernel-managed datasource discovery.
- * Consumers decide how to render, filter, or act on the detected model.
  */
 export function useDataSourceDiscovery() {
-  return useAsyncData(loadDataSourceDiscovery, []);
+  const epoch = useAtomValue(dataSourceDiscoveryEpochAtom);
+  return useAsyncData(fetchDataSourceDiscovery, [epoch]);
+}
+
+/**
+ * Detected sources from the kernel environment that are not already backed
+ * by a live connection, optionally narrowed to "database" or "storage".
+ */
+export function useDetectedDataSources(
+  group?: DataSourceDiscoveryGroup,
+): DetectedDataSource[] {
+  const { data } = useDataSourceDiscovery();
+  const connections = useAtomValue(dataConnectionsMapAtom);
+  const storageNamespaces = useAtomValue(storageNamespacesAtom);
+
+  return useMemo(() => {
+    if (!data) {
+      return [];
+    }
+
+    const snapshot = {
+      dialects: [...connections.values()].map((connection) =>
+        connection.dialect.toLowerCase(),
+      ),
+      storageProtocols: storageNamespaces.map((namespace) =>
+        namespace.protocol.toLowerCase(),
+      ),
+      storageBackendTypes: storageNamespaces.map((namespace) =>
+        namespace.backendType.toLowerCase(),
+      ),
+    };
+
+    return data.filter(
+      (source) =>
+        matchesDiscoveryGroup(source, group) &&
+        !isDetectedSourceConnected(source, snapshot),
+    );
+  }, [connections, data, group, storageNamespaces]);
 }