@marimo-team/frontend 0.25.3-dev13 → 0.25.3-dev135

package/src/components/datasources/__tests__/column-preview.test.tsx

@@ -1,16 +1,24 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-import { render } from "@testing-library/react";
-import { Provider } from "jotai";
+import { fireEvent, render, screen } from "@testing-library/react";
+import { createStore, Provider } from "jotai";
 import type React from "react";
 import { beforeEach, describe, expect, it, vi } from "vitest";
 import { MockRequestClient } from "@/__mocks__/requests";
+import { TooltipProvider } from "@/components/ui/tooltip";
 import type { SQLTableContext } from "@/core/datasets/data-source-connections";
 import type { DataTable, DataTableColumn } from "@/core/kernel/messages";
 import { requestClientAtom } from "@/core/network/requests";
 import { store } from "@/core/state/jotai";
 import { DatasetColumnPreview } from "../column-preview";
 
+const { createNewCell } = vi.hoisted(() => ({ createNewCell: vi.fn() }));
+
+vi.mock("@/core/cells/cells", async (importOriginal) => ({
+  ...(await importOriginal<typeof import("@/core/cells/cells")>()),
+  useCellActions: () => ({ createNewCell }),
+}));
+
 const wrapper = ({ children }: { children: React.ReactNode }) => (
   <Provider store={store}>{children}</Provider>
 );
@@ -71,7 +79,9 @@ describe("DatasetColumnPreview fullyQualifiedTableName", () => {
   it("qualifies with database and schema for flat engines", () => {
     const client = renderPreview(ctx({ database: "db", schema: "public" }));
     expect(client.previewDatasetColumn).toHaveBeenCalledWith(
-      expect.objectContaining({ fullyQualifiedTableName: "db.public.users" }),
+      expect.objectContaining({
+        fullyQualifiedTableName: '"db"."public"."users"',
+      }),
     );
   });
 
@@ -79,7 +89,7 @@ describe("DatasetColumnPreview fullyQualifiedTableName", () => {
     // Regression: previously produced "db..users".
     const client = renderPreview(ctx({ database: "db", schema: "" }));
     expect(client.previewDatasetColumn).toHaveBeenCalledWith(
-      expect.objectContaining({ fullyQualifiedTableName: "db.users" }),
+      expect.objectContaining({ fullyQualifiedTableName: '"db"."users"' }),
     );
   });
 
@@ -90,8 +100,91 @@ describe("DatasetColumnPreview fullyQualifiedTableName", () => {
     );
     expect(client.previewDatasetColumn).toHaveBeenCalledWith(
       expect.objectContaining({
-        fullyQualifiedTableName: "top.nested.deep.users",
+        fullyQualifiedTableName: '"top"."nested"."deep"."users"',
+      }),
+    );
+  });
+
+  it("quotes a literal dot as one schema segment", () => {
+    const client = renderPreview(
+      ctx({
+        database: "memory",
+        schema: "finance.reports",
+        schemaPath: ["finance.reports"],
+      }),
+    );
+    expect(client.previewDatasetColumn).toHaveBeenCalledWith(
+      expect.objectContaining({
+        fullyQualifiedTableName: '"memory"."finance.reports"."users"',
       }),
     );
   });
+
+  it("escapes quotes inside schema identifiers", () => {
+    const client = renderPreview(
+      ctx({
+        database: "memory",
+        schema: 'fi"nance',
+        schemaPath: ['fi"nance', "reports"],
+      }),
+    );
+    expect(client.previewDatasetColumn).toHaveBeenCalledWith(
+      expect.objectContaining({
+        fullyQualifiedTableName: '"memory"."fi""nance"."reports"."users"',
+      }),
+    );
+  });
+});
+
+describe("DatasetColumnPreview chart actions", () => {
+  it("identifies each column and inserts its corresponding chart", () => {
+    const store = createStore();
+    store.set(requestClientAtom, MockRequestClient.create());
+    const columns: DataTableColumn[] = [
+      {
+        name: "email",
+        type: "string",
+        external_type: "str",
+        sample_values: [],
+      },
+      { name: "age", type: "integer", external_type: "int", sample_values: [] },
+    ];
+    render(
+      <Provider store={store}>
+        <TooltipProvider>
+          {columns.map((column) => (
+            <DatasetColumnPreview
+              key={column.name}
+              table={{ ...table, source_type: "local", engine: null }}
+              column={column}
+              onAddColumnChart={vi.fn()}
+              preview={{
+                table_name: table.name,
+                column_name: column.name,
+                chart_code: `chart_${column.name}`,
+                chart_spec: null,
+                error: null,
+                missing_packages: null,
+                stats: {},
+              }}
+            />
+          ))}
+        </TooltipProvider>
+      </Provider>,
+    );
+
+    for (const column of columns) {
+      fireEvent.click(
+        screen.getByRole("button", {
+          name: `Add chart to notebook for ${column.name}`,
+        }),
+      );
+      expect(createNewCell).toHaveBeenLastCalledWith({
+        code: `chart_${column.name}`,
+        before: false,
+        cellId: "__end__",
+      });
+    }
+    expect(createNewCell).toHaveBeenCalledTimes(2);
+  });
 });

package/src/components/datasources/__tests__/datasources-interactions.test.tsx

@@ -0,0 +1,181 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import {
+  act,
+  fireEvent,
+  render,
+  screen,
+  waitFor,
+} from "@testing-library/react";
+import { createStore, Provider } from "jotai";
+import { beforeAll, describe, expect, it, vi } from "vitest";
+import { MockRequestClient } from "@/__mocks__/requests";
+import { Command, CommandList } from "@/components/ui/command";
+import { TooltipProvider } from "@/components/ui/tooltip";
+import {
+  dataSourceConnectionsAtom,
+  type SQLTableContext,
+} from "@/core/datasets/data-source-connections";
+import { DUCKDB_ENGINE } from "@/core/datasets/engines";
+import { datasetsAtom, exportedForTesting } from "@/core/datasets/state";
+import type { DataTable } from "@/core/kernel/messages";
+import { requestClientAtom } from "@/core/network/requests";
+import {
+  DataSources,
+  hideEmptyDatasourcesAtom,
+  TableList,
+} from "../datasources";
+
+beforeAll(() => {
+  HTMLElement.prototype.scrollIntoView = vi.fn();
+});
+
+const table: DataTable = {
+  name: "reports",
+  columns: [
+    {
+      name: "amount",
+      type: "integer",
+      external_type: "INTEGER",
+      sample_values: [],
+    },
+  ],
+  source: "memory",
+  source_type: "duckdb",
+  type: "table",
+  engine: null,
+  indexes: null,
+  num_columns: 1,
+  num_rows: null,
+  variable_name: null,
+  primary_keys: null,
+};
+
+describe("Datasource column preview responses", () => {
+  it.each([
+    { path: ["main"], qualifiedName: '"memory"."main"."reports"' },
+    {
+      path: ["finance", "reports"],
+      qualifiedName: '"memory"."finance"."reports"."reports"',
+    },
+    {
+      path: ["finance.reports"],
+      qualifiedName: '"memory"."finance.reports"."reports"',
+    },
+    {
+      path: ['fi"nance', "reports"],
+      qualifiedName: '"memory"."fi""nance"."reports"."reports"',
+    },
+  ])(
+    "renders the returned preview for $path",
+    async ({ path, qualifiedName }) => {
+      const store = createStore();
+      const client = MockRequestClient.create();
+      store.set(requestClientAtom, client);
+      const context: SQLTableContext = {
+        engine: DUCKDB_ENGINE,
+        database: "memory",
+        schema: path.at(-1) ?? "",
+        schemaPath: path,
+        dialect: "duckdb",
+      };
+      render(
+        <Provider store={store}>
+          <TooltipProvider>
+            <Command shouldFilter={false}>
+              <CommandList>
+                <TableList tables={[table]} sqlTableContext={context} />
+              </CommandList>
+            </Command>
+          </TooltipProvider>
+        </Provider>,
+      );
+
+      fireEvent.click(await screen.findByText("reports"));
+      fireEvent.click(await screen.findByText("amount"));
+      await waitFor(() => {
+        expect(client.previewDatasetColumn).toHaveBeenCalledWith(
+          expect.objectContaining({ fullyQualifiedTableName: qualifiedName }),
+        );
+      });
+      expect(screen.getByText("Loading...")).toBeInTheDocument();
+
+      const actions = exportedForTesting.createActions((action) => {
+        store.set(datasetsAtom, (state) =>
+          exportedForTesting.reducer(state, action),
+        );
+      });
+      act(() => {
+        actions.addColumnPreview({
+          table_name: qualifiedName,
+          column_name: "amount",
+          stats: { min: 42 },
+          error: null,
+          missing_packages: null,
+          chart_spec: null,
+          chart_code: null,
+        });
+      });
+
+      expect(await screen.findByText("42")).toBeInTheDocument();
+      expect(screen.queryByText("Loading...")).not.toBeInTheDocument();
+    },
+  );
+});
+
+it("activates a schema and a table with the same path independently using the keyboard", async () => {
+  const store = createStore();
+  store.set(requestClientAtom, MockRequestClient.create());
+  store.set(hideEmptyDatasourcesAtom, false);
+  store.set(dataSourceConnectionsAtom, {
+    latestEngineSelected: DUCKDB_ENGINE,
+    connectionsMap: new Map([
+      [
+        DUCKDB_ENGINE,
+        {
+          name: DUCKDB_ENGINE,
+          dialect: "duckdb",
+          source: "duckdb",
+          display_name: "DuckDB (In-Memory)",
+          databases: [
+            {
+              name: "memory",
+              dialect: "duckdb",
+              schemas: [
+                {
+                  name: "finance",
+                  tables: [table],
+                  child_schemas: [{ name: "reports", tables: [] }],
+                },
+              ],
+            },
+          ],
+        },
+      ],
+    ]),
+  });
+  render(
+    <Provider store={store}>
+      <TooltipProvider>
+        <DataSources />
+      </TooltipProvider>
+    </Provider>,
+  );
+  fireEvent.click(await screen.findByText("memory"));
+  fireEvent.click(await screen.findByText("finance"));
+  await waitFor(() => {
+    expect(screen.getAllByText("reports")).toHaveLength(2);
+  });
+
+  const input = screen.getByRole("combobox");
+  input.focus();
+  fireEvent.keyDown(input, { key: "ArrowDown" });
+  fireEvent.keyDown(input, { key: "Enter" });
+  expect(await screen.findByText("No tables found")).toBeInTheDocument();
+  expect(screen.queryByText("amount")).not.toBeInTheDocument();
+
+  fireEvent.keyDown(input, { key: "ArrowDown" });
+  fireEvent.keyDown(input, { key: "Enter" });
+  expect(await screen.findByText("amount")).toBeInTheDocument();
+  expect(screen.getByText("No tables found")).toBeInTheDocument();
+});

package/src/components/datasources/__tests__/table-list.test.tsx

@@ -0,0 +1,118 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { render, screen, waitFor } from "@testing-library/react";
+import { Provider } from "jotai";
+import type React from "react";
+import { describe, expect, it } from "vitest";
+import { MockRequestClient } from "@/__mocks__/requests";
+import type { SQLTableContext } from "@/core/datasets/data-source-connections";
+import { PreviewSQLTableList } from "@/core/datasets/request-registry";
+import type { SQLTableListPreview } from "@/core/kernel/messages";
+import type { RequestId } from "@/core/network/DeferredRequestRegistry";
+import { requestClientAtom } from "@/core/network/requests";
+import { store } from "@/core/state/jotai";
+import { TableList } from "../datasources";
+
+const wrapper = ({ children }: { children: React.ReactNode }) => (
+  <Provider store={store}>{children}</Provider>
+);
+
+const sqlTableContext: SQLTableContext = {
+  engine: "my_engine",
+  database: "my_db",
+  schema: "public",
+  dialect: "postgres",
+  schemaPath: ["public"],
+};
+
+// The kernel answers table-list requests over the websocket, so the mock
+// settles the deferred request instead of returning the response directly.
+function renderDeferredTableList(settle?: (requestId: RequestId) => void) {
+  const client = MockRequestClient.create();
+  client.previewSQLTableList.mockImplementation(async ({ requestId }) => {
+    settle?.(requestId);
+    return null;
+  });
+  store.set(requestClientAtom, client);
+  render(
+    <TableList
+      tables={[]}
+      tablesResolved={false}
+      sqlTableContext={sqlTableContext}
+    />,
+    { wrapper },
+  );
+  return client;
+}
+
+describe("TableList deferred discovery", () => {
+  it("shows the error when the request rejects", async () => {
+    renderDeferredTableList((requestId) =>
+      PreviewSQLTableList.reject(requestId, new Error("engine offline")),
+    );
+
+    expect(screen.getByText("Loading tables...")).toBeInTheDocument();
+
+    await waitFor(() => {
+      expect(screen.getByText("engine offline")).toBeInTheDocument();
+    });
+    expect(screen.queryByText("Loading tables...")).toBeNull();
+  });
+
+  it("shows the error when the response has no tables field", async () => {
+    renderDeferredTableList((requestId) =>
+      PreviewSQLTableList.resolve(requestId, {} as SQLTableListPreview),
+    );
+
+    await waitFor(() => {
+      expect(screen.getByText("No tables available")).toBeInTheDocument();
+    });
+    expect(screen.queryByText("Loading tables...")).toBeNull();
+  });
+
+  it("shows the empty state for an empty table list", async () => {
+    renderDeferredTableList((requestId) =>
+      PreviewSQLTableList.resolve(requestId, {
+        tables: [],
+      } as unknown as SQLTableListPreview),
+    );
+
+    await waitFor(() => {
+      expect(screen.getByText("No tables found")).toBeInTheDocument();
+    });
+    expect(screen.queryByText("Loading tables...")).toBeNull();
+  });
+
+  it("requests the table list with the table context and waits for the kernel", async () => {
+    let pendingRequestId: RequestId | undefined;
+    const client = renderDeferredTableList((requestId) => {
+      pendingRequestId = requestId;
+    });
+
+    await waitFor(() => {
+      expect(client.previewSQLTableList).toHaveBeenCalledTimes(1);
+    });
+    expect(client.previewSQLTableList).toHaveBeenCalledWith(
+      expect.objectContaining({
+        engine: "my_engine",
+        database: "my_db",
+        schema: "public",
+        schemaPath: ["public"],
+      }),
+    );
+    expect(screen.getByText("Loading tables...")).toBeInTheDocument();
+
+    // Settle the request so the process-wide registry does not keep the entry.
+    expect(pendingRequestId).toBeDefined();
+    PreviewSQLTableList.resolve(
+      pendingRequestId as RequestId,
+      {
+        tables: [],
+      } as unknown as SQLTableListPreview,
+    );
+    await waitFor(() => {
+      expect(screen.queryByText("Loading tables...")).toBeNull();
+    });
+    expect(PreviewSQLTableList.requests.size).toBe(0);
+  });
+});

package/src/components/datasources/__tests__/utils.test.ts

@@ -34,6 +34,37 @@ describe("sqlCode", () => {
     name: "email" as const,
   } as DataTableColumn;
 
+  it.each([
+    {
+      path: ["finance", "reports"],
+      expected: '"memory"."finance"."reports"."users"',
+    },
+    { path: ["finance.reports"], expected: '"finance.reports"."users"' },
+    {
+      path: ['fi"nance', "reports"],
+      expected: '"memory"."fi""nance"."reports"."users"',
+    },
+    {
+      path: ["main", "reports"],
+      expected: '"memory"."main"."reports"."users"',
+    },
+  ])("preserves schema path $path in generated SQL", ({ path, expected }) => {
+    const result = sqlCode({
+      table: mockTable,
+      columnName: "*",
+      sqlTableContext: {
+        engine: DUCKDB_ENGINE,
+        dialect: "duckdb",
+        database: "memory",
+        defaultDatabase: "memory",
+        defaultSchema: "main",
+        schema: path.at(-1)!,
+        schemaPath: path,
+      },
+    });
+    expect(result).toContain(`FROM ${expected}`);
+  });
+
   describe("basic SQL generation", () => {
     it("should generate basic SQL without sqlTableContext", () => {
       const result = sqlCode({ table: mockTable, columnName: mockColumn.name });
@@ -484,7 +515,7 @@ describe("tableUniqueId", () => {
 
   it("uses database + schema for flat engines", () => {
     expect(tableUniqueId(ctx({ database: "db", schema: "public" }), "t")).toBe(
-      "db.public.t",
+      '["db","public","t"]',
     );
   });
 
@@ -495,7 +526,7 @@ describe("tableUniqueId", () => {
         ctx({ database: "top", schema: "nested", schemaPath: ["nested"] }),
         "t",
       ),
-    ).toBe("top.nested.t");
+    ).toBe('["top","nested","t"]');
   });
 
   it("includes the full schema path for deeply nested namespaces", () => {
@@ -508,7 +539,7 @@ describe("tableUniqueId", () => {
         }),
         "t",
       ),
-    ).toBe("top.nested.deep.t");
+    ).toBe('["top","nested","deep","t"]');
   });
 
   it("falls back to the flat schema when schemaPath is empty", () => {
@@ -517,16 +548,36 @@ describe("tableUniqueId", () => {
         ctx({ database: "db", schema: "public", schemaPath: [] }),
         "t",
       ),
-    ).toBe("db.public.t");
+    ).toBe('["db","public","t"]');
   });
 
   it("does not emit a double dot for schemaless tables", () => {
     expect(
       tableUniqueId(ctx({ database: "db", schema: "", schemaPath: [] }), "t"),
-    ).toBe("db.t");
+    ).toBe('["db","t"]');
     expect(tableUniqueId(ctx({ database: "db", schema: "" }), "t")).toBe(
-      "db.t",
+      '["db","t"]',
+    );
+  });
+
+  it("keeps literal-dot and nested schemas distinct", () => {
+    const nested = tableUniqueId(
+      ctx({
+        database: "memory",
+        schema: "reports",
+        schemaPath: ["finance", "reports"],
+      }),
+      "sales",
+    );
+    const literal = tableUniqueId(
+      ctx({
+        database: "memory",
+        schema: "finance.reports",
+        schemaPath: ["finance.reports"],
+      }),
+      "sales",
     );
+    expect(nested).not.toBe(literal);
   });
 });
 

package/src/components/datasources/column-preview.tsx

@@ -29,7 +29,7 @@ import { Button } from "../ui/button";
 import { Tooltip } from "../ui/tooltip";
 import { ColumnPreviewContainer } from "./components";
 import { InstallPackageButton } from "./install-package-button";
-import { convertStatsName, sqlCode, tableUniqueId } from "./utils";
+import { convertStatsName, fullyQualifiedTableName, sqlCode } from "./utils";
 
 export const DatasetColumnPreview: React.FC<{
   table: DataTable;
@@ -48,7 +48,10 @@ export const DatasetColumnPreview: React.FC<{
       tableName: table.name,
       columnName: column.name,
       sourceType: table.source_type,
-      fullyQualifiedTableName: tableUniqueId(sqlTableContext, table.name),
+      fullyQualifiedTableName: fullyQualifiedTableName(
+        sqlTableContext,
+        table.name,
+      ),
     });
   };
 
@@ -112,12 +115,16 @@ export const DatasetColumnPreview: React.FC<{
 
   const addDataframeChart = preview.chart_code &&
     table.source_type === "local" && (
-      <AddDataframeChart chartCode={preview.chart_code} />
+      <AddDataframeChart
+        chartCode={preview.chart_code}
+        columnName={column.name}
+      />
     );
 
   const addSQLChart = table.source_type === "duckdb" && (
     <Tooltip content="Add SQL cell" delayDuration={400}>
       <Button
+        aria-label={`Add SQL cell for ${column.name}`}
         variant="outline"
         size="icon"
         className="z-10 bg-background absolute right-1 -top-1"
@@ -238,7 +245,8 @@ export function renderChart(chartSpec: string, theme: Theme) {
 
 export const AddDataframeChart: React.FC<{
   chartCode: string;
-}> = ({ chartCode }) => {
+  columnName: string;
+}> = ({ chartCode, columnName }) => {
   const autoInstantiate = useAtomValue(autoInstantiateAtom);
   const lastFocusedCellId = useLastFocusedCellId();
   const { createNewCell } = useCellActions();
@@ -261,6 +269,7 @@ export const AddDataframeChart: React.FC<{
   return (
     <Tooltip content="Add chart to notebook" delayDuration={400}>
       <Button
+        aria-label={`Add chart to notebook for ${columnName}`}
         variant="outline"
         size="icon"
         className="z-10 bg-background absolute right-1 -top-0.5"