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

package/src/core/codemirror/code-lens/__tests__/popover.test.ts

@@ -56,6 +56,19 @@ describe("CachePopover getCacheInfo throttling", () => {
     return () => act(() => dispose());
   }
 
+  it("renders content synchronously for CodeMirror's initial measurement", async () => {
+    const dom = document.createElement("div");
+
+    // Keep the mount outside `act()`. It would flush an asynchronous render
+    // and make this assertion pass without `flushSync`.
+    const dispose = mountLensPopover(dom, CACHE_SPEC);
+
+    expect(dom.textContent).toContain("my_cache");
+    await act(async () => {
+      dispose();
+    });
+  });
+
   it("does not refetch on every hover, only after the throttle interval elapses", () => {
     const disposeFirst = mount();
     expect(getCacheInfoMock).toHaveBeenCalledTimes(1);

package/src/core/codemirror/code-lens/popover.tsx

@@ -4,6 +4,7 @@ import { Provider, useAtomValue } from "jotai";
 import { DatabaseZapIcon } from "lucide-react";
 import type React from "react";
 import { useEffect } from "react";
+import { flushSync } from "react-dom";
 import { createRoot } from "react-dom/client";
 import { useLocale } from "react-aria";
 import { ProtocolIcon } from "@/components/storage/components";
@@ -39,11 +40,16 @@ export function mountLensPopover(
   spec: CodeLensSpec,
 ): () => void {
   const root = createRoot(dom);
-  root.render(
-    <Provider store={store}>
-      <LensPopover spec={spec} />
-    </Provider>,
-  );
+  // CodeMirror measures and positions the tooltip as soon as this function
+  // returns. Render now so it measures the popover content, not an empty
+  // container that later grows over the hovered icon.
+  flushSync(() => {
+    root.render(
+      <Provider store={store}>
+        <LensPopover spec={spec} />
+      </Provider>,
+    );
+  });
   // Defer so tearing the tooltip down from inside a React event (e.g. the
   // icon's click handler) doesn't unmount mid-render
   return () => queueMicrotask(() => root.unmount());

package/src/core/codemirror/language/languages/sql/__tests__/renderers.test.tsx

@@ -9,13 +9,18 @@ import type {
 } from "@/core/kernel/messages";
 import { renderColumnInfo, renderTableInfo } from "../renderers";
 
-const geometryColumn = {
+const bogusColumn = {
   name: "geom",
-  type: "geometry" as DataType,
+  type: "bogus_type" as DataType,
   external_type: "geometry",
   sample_values: [],
 } as DataTableColumn;
 
+const geometryColumn = {
+  ...bogusColumn,
+  type: "geometry" as DataType,
+} as DataTableColumn;
+
 describe("SQL renderers", () => {
   it("renders table information with an unrecognized column type", () => {
     const table: DataTable = {
@@ -26,7 +31,7 @@ describe("SQL renderers", () => {
       num_rows: null,
       num_columns: 1,
       variable_name: null,
-      columns: [geometryColumn],
+      columns: [bogusColumn],
       primary_keys: [],
       indexes: [],
     };
@@ -37,8 +42,14 @@ describe("SQL renderers", () => {
   });
 
   it("renders column information with an unrecognized column type", () => {
-    const result = render(renderColumnInfo(geometryColumn));
+    const result = render(renderColumnInfo(bogusColumn));
 
     expect(result.container.querySelector("svg")).not.toBeNull();
   });
+
+  it("renders geometry column information with the map pin icon", () => {
+    const result = render(renderColumnInfo(geometryColumn));
+
+    expect(result.container.querySelector(".lucide-map-pin")).not.toBeNull();
+  });
 });

package/src/core/codemirror/language/languages/sql/renderers.tsx

@@ -41,6 +41,7 @@ const DATA_TYPE_COLORS: Record<DataType, string> = {
   number: "bg-(--purple-4) text-(--purple-11)",
   integer: "bg-(--purple-4) text-(--purple-11)",
   string: "bg-(--blue-4) text-(--blue-11)",
+  geometry: "bg-(--cyan-4) text-(--cyan-11)",
   unknown: "bg-(--slate-4) text-(--slate-11)",
 };
 

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]);
 }