@marimo-team/frontend 0.23.11-dev1 → 0.23.11-dev2

package/src/components/datasources/utils.ts

@@ -4,15 +4,54 @@ import { BigQueryDialect } from "@marimo-team/codemirror-sql/dialects";
 import { isKnownDialect } from "@/core/codemirror/language/languages/sql/utils";
 import type { SQLTableContext } from "@/core/datasets/data-source-connections";
 import { DUCKDB_ENGINE } from "@/core/datasets/engines";
-import type { DataTable, DataType } from "@/core/kernel/messages";
+import type {
+  Database,
+  DatabaseSchema,
+  DataTable,
+  DataType,
+} from "@/core/kernel/messages";
 import { logNever } from "@/utils/assertNever";
 import type { ColumnHeaderStatsKey } from "../data-table/types";
 
+/**
+ * Stable id for a table node in the datasources tree.
+ *
+ * schemaPath already includes the leaf schema for nested namespaces, so use it
+ * when present and fall back to the flat schema name otherwise (avoids
+ * duplicating the leaf, e.g. `top.nested.nested.table`).
+ */
+export function tableUniqueId(
+  sqlTableContext: SQLTableContext | undefined,
+  tableName: string,
+): string {
+  if (!sqlTableContext) {
+    return tableName;
+  }
+  const segments = (
+    sqlTableContext.schemaPath?.length
+      ? sqlTableContext.schemaPath
+      : [sqlTableContext.schema]
+  ).filter(Boolean);
+  return [sqlTableContext.database, ...segments, tableName].join(".");
+}
+
 // Some databases have no schemas, so we don't show it (eg. Clickhouse)
 export function isSchemaless(schemaName: string) {
   return schemaName === "";
 }
 
+// Lazy discovery: the `*_resolved` flags default to `true` and are only `false`
+// when enumeration was deferred. Helper functions to centralize logic
+export function areSchemasResolved(database: Database): boolean {
+  return database.schemas_resolved !== false;
+}
+export function areTablesResolved(schema: DatabaseSchema): boolean {
+  return schema.tables_resolved !== false;
+}
+export function areChildSchemasResolved(schema: DatabaseSchema): boolean {
+  return schema.child_schemas_resolved !== false;
+}
+
 interface SqlCodeFormatter {
   /**
    * Format the table path based on dialect-specific rules

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

@@ -2,9 +2,12 @@
 import { beforeEach, describe, expect, it } from "vitest";
 import { variableName } from "@/__tests__/branded";
 import type { DatabaseSchema, DataTable } from "@/core/kernel/messages";
+import { store } from "@/core/state/jotai";
 import type { VariableName } from "@/core/variables/types";
 import {
+  allTablesAtom,
   type DataSourceConnection,
+  dataSourceConnectionsAtom,
   type DataSourceState,
   exportedForTesting,
   type SQLTableContext,
@@ -575,3 +578,226 @@ describe("add table", () => {
     expect(db1?.schemas.length).toBe(1);
   });
 });
+
+describe("nested namespaces", () => {
+  // Iceberg-style: database "top" with a schemaless schema and a nested
+  // namespace "nested" that has a deferred child namespace "deep".
+  const nestedConnections: DataSourceConnection[] = [
+    {
+      name: "ice" as ConnectionName,
+      source: "iceberg",
+      display_name: "Iceberg",
+      dialect: "iceberg",
+      databases: [
+        {
+          name: "top",
+          dialect: "iceberg",
+          schemas_resolved: true,
+          schemas: [
+            { name: "", tables: [], tables_resolved: true },
+            {
+              name: "nested",
+              tables: [],
+              tables_resolved: false,
+              child_schemas: [],
+              child_schemas_resolved: false,
+            },
+          ],
+        },
+      ],
+    },
+  ];
+
+  let baseState: DataSourceState;
+
+  beforeEach(() => {
+    baseState = addConnection(nestedConnections, initialState());
+  });
+
+  const findSchema = (
+    state: DataSourceState,
+    path: string[],
+  ): DatabaseSchema | undefined => {
+    const conn = state.connectionsMap.get("ice" as ConnectionName);
+    const db = conn?.databases.find((d) => d.name === "top");
+    let schemas = db?.schemas ?? [];
+    let found: DatabaseSchema | undefined;
+    for (const segment of path) {
+      found = schemas.find((s) => s.name === segment);
+      schemas = found?.child_schemas ?? [];
+    }
+    return found;
+  };
+
+  it("sets child namespaces at a nested path", () => {
+    const children: DatabaseSchema[] = [
+      { name: "deep", tables: [], tables_resolved: false },
+    ];
+    const newState = reducer(baseState, {
+      type: "addSchemaList",
+      payload: {
+        schemas: children,
+        sqlSchemaContext: {
+          engine: "ice",
+          database: "top",
+          schemaPath: ["nested"],
+        },
+      },
+    });
+
+    const nested = findSchema(newState, ["nested"]);
+    expect(nested?.child_schemas_resolved).toBe(true);
+    expect(nested?.child_schemas?.map((s) => s.name)).toEqual(["deep"]);
+    // The schemaless sibling is untouched.
+    expect(findSchema(newState, [""])?.name).toBe("");
+  });
+
+  it("sets tables at a nested path", () => {
+    const tables: DataTable[] = [
+      {
+        name: "table4",
+        columns: [],
+        num_columns: 0,
+        num_rows: 0,
+        variable_name: null,
+        source: "iceberg",
+        source_type: "catalog",
+        type: "table",
+      },
+    ];
+    const newState = reducer(baseState, {
+      type: "addTableList",
+      payload: {
+        tables,
+        sqlTableContext: {
+          engine: "ice",
+          database: "top",
+          schema: "nested",
+          dialect: "iceberg",
+          schemaPath: ["nested"],
+        },
+      },
+    });
+
+    const nested = findSchema(newState, ["nested"]);
+    expect(nested?.tables_resolved).toBe(true);
+    expect(nested?.tables.map((t) => t.name)).toEqual(["table4"]);
+  });
+
+  it("replaces a single table at a nested path", () => {
+    const makeTable = (numRows: number): DataTable => ({
+      name: "table4",
+      columns: [],
+      num_columns: 0,
+      num_rows: numRows,
+      variable_name: null,
+      source: "iceberg",
+      source_type: "catalog",
+      type: "table",
+    });
+    const context = {
+      engine: "ice",
+      database: "top",
+      schema: "nested",
+      dialect: "iceberg",
+      schemaPath: ["nested"],
+    };
+    let state = reducer(baseState, {
+      type: "addTableList",
+      payload: { tables: [makeTable(1)], sqlTableContext: context },
+    });
+    state = reducer(state, {
+      type: "addTable",
+      payload: { table: makeTable(42), sqlTableContext: context },
+    });
+
+    const nested = findSchema(state, ["nested"]);
+    expect(nested?.tables).toHaveLength(1);
+    expect(nested?.tables[0].num_rows).toBe(42);
+  });
+
+  it("does not change anything for a missing nested path", () => {
+    const newState = reducer(baseState, {
+      type: "addSchemaList",
+      payload: {
+        schemas: [{ name: "deep", tables: [] }],
+        sqlSchemaContext: {
+          engine: "ice",
+          database: "top",
+          schemaPath: ["does_not_exist"],
+        },
+      },
+    });
+    // The result is unchanged: nested namespace stays unresolved and the
+    // database keeps its two schemas (schemaless + nested).
+    const newDb = newState.connectionsMap
+      .get("ice" as ConnectionName)
+      ?.databases.find((d) => d.name === "top");
+    expect(findSchema(newState, ["nested"])?.child_schemas_resolved).toBe(
+      false,
+    );
+    expect(newDb?.schemas.length).toBe(2);
+  });
+});
+
+describe("allTablesAtom with nested namespaces", () => {
+  it("enumerates tables from nested namespaces", () => {
+    const table = (name: string): DataTable => ({
+      name,
+      columns: [],
+      num_columns: 0,
+      num_rows: 0,
+      variable_name: null,
+      source: "iceberg",
+      source_type: "catalog",
+      type: "table",
+    });
+
+    const state = addConnection(
+      [
+        {
+          name: "ice" as ConnectionName,
+          source: "iceberg",
+          display_name: "Iceberg",
+          dialect: "iceberg",
+          databases: [
+            {
+              name: "top",
+              dialect: "iceberg",
+              schemas_resolved: true,
+              schemas: [
+                {
+                  name: "",
+                  tables: [table("toptable")],
+                  tables_resolved: true,
+                },
+                {
+                  name: "nested",
+                  tables: [table("nestedtable")],
+                  tables_resolved: true,
+                  child_schemas_resolved: true,
+                  child_schemas: [
+                    {
+                      name: "deep",
+                      tables: [table("deeptable")],
+                      tables_resolved: true,
+                      child_schemas: [],
+                      child_schemas_resolved: true,
+                    },
+                  ],
+                },
+              ],
+            },
+          ],
+        },
+      ],
+      initialState(),
+    );
+
+    store.set(dataSourceConnectionsAtom, state);
+    const names = [...store.get(allTablesAtom).values()].map((t) => t.name);
+    expect(names).toContain("toptable");
+    expect(names).toContain("nestedtable");
+    expect(names).toContain("deeptable");
+  });
+});

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

@@ -49,6 +49,9 @@ export interface DataSourceState {
 export interface SQLSchemaContext {
   engine: string;
   database: string;
+  // Parent schema path (relative to `database`) for nested schemas.
+  // Empty/undefined for the database's top level.
+  schemaPath?: string[];
 }
 
 export interface SQLTableContext {
@@ -58,6 +61,49 @@ export interface SQLTableContext {
   dialect: string;
   defaultSchema?: string | null;
   defaultDatabase?: string | null;
+  // Nested schema path (relative to `database`). Empty/undefined at top level.
+  schemaPath?: string[];
+}
+
+/**
+ * Immutably descend `path` (schema segment names) into a nested schema list
+ * and apply `update` to the matching schema. Unmatched branches are unchanged.
+ */
+function updateSchemaAtPath(
+  schemas: DatabaseSchema[],
+  path: string[],
+  update: (schema: DatabaseSchema) => DatabaseSchema,
+): DatabaseSchema[] {
+  if (path.length === 0) {
+    return schemas;
+  }
+  const [head, ...rest] = path;
+  return schemas.map((schema) => {
+    if (schema.name !== head) {
+      return schema;
+    }
+    if (rest.length === 0) {
+      return update(schema);
+    }
+    return {
+      ...schema,
+      child_schemas: updateSchemaAtPath(
+        schema.child_schemas ?? [],
+        rest,
+        update,
+      ),
+    };
+  });
+}
+
+/**
+ * The path (schema/namespace segment names within a database) that locates the
+ * schema holding a set of tables. For nested namespaces this is the
+ * `schemaPath`; otherwise it is the single (possibly schemaless) schema name.
+ */
+function tableSchemaPath(sqlTableContext: SQLTableContext): string[] {
+  const { schemaPath, schema } = sqlTableContext;
+  return schemaPath && schemaPath.length > 0 ? schemaPath : [schema];
 }
 
 function initialState(): DataSourceState {
@@ -159,6 +205,7 @@ const {
       return state;
     }
 
+    const schemaPath = sqlSchemaContext.schemaPath ?? [];
     const newMap = new Map(connectionsMap);
     const newConn: DataSourceConnection = {
       ...conn,
@@ -166,10 +213,22 @@ const {
         if (db.name !== sqlSchemaContext.database) {
           return db;
         }
+        // Top level: replace the database's schemas.
+        if (schemaPath.length === 0) {
+          return {
+            ...db,
+            schemas: schemas,
+            schemas_resolved: true,
+          };
+        }
+        // Nested namespace: replace the child schemas of the namespace at path.
         return {
           ...db,
-          schemas: schemas,
-          schemas_resolved: true,
+          schemas: updateSchemaAtPath(db.schemas, schemaPath, (schema) => ({
+            ...schema,
+            child_schemas: schemas,
+            child_schemas_resolved: true,
+          })),
         };
       }),
     };
@@ -198,6 +257,7 @@ const {
       return state;
     }
 
+    const path = tableSchemaPath(sqlTableContext);
     const newMap = new Map(connectionsMap);
     const newConn: DataSourceConnection = {
       ...conn,
@@ -207,16 +267,11 @@ const {
         }
         return {
           ...db,
-          schemas: db.schemas.map((schema) => {
-            if (schema.name !== sqlTableContext.schema) {
-              return schema;
-            }
-            return {
-              ...schema,
-              tables: tables,
-              tables_resolved: true,
-            };
-          }),
+          schemas: updateSchemaAtPath(db.schemas, path, (schema) => ({
+            ...schema,
+            tables: tables,
+            tables_resolved: true,
+          })),
         };
       }),
     };
@@ -246,6 +301,7 @@ const {
       return state;
     }
 
+    const path = tableSchemaPath(sqlTableContext);
     const newMap = new Map(connectionsMap);
     const newConn: DataSourceConnection = {
       ...conn,
@@ -253,21 +309,15 @@ const {
         if (db.name !== sqlTableContext.database) {
           return db;
         }
-
         return {
           ...db,
-          schemas: db.schemas.map((schema) => {
-            if (schema.name !== sqlTableContext.schema) {
-              return schema;
-            }
-
+          schemas: updateSchemaAtPath(db.schemas, path, (schema) => {
             // If tables array is empty, add the table
             // Otherwise, replace existing table or keep unchanged tables
             const tables =
               schema.tables.length === 0
                 ? [table]
                 : schema.tables.map((t) => (t.name === tableName ? table : t));
-
             return {
               ...schema,
               tables,
@@ -327,9 +377,14 @@ export const allTablesAtom = atom((get) => {
       const isDefaultDb =
         database.name === conn.default_database || conn.databases.length === 1;
 
-      for (const schema of database.schemas) {
-        const isDefaultSchema = schema.name === conn.default_schema;
-        const schemalessDb = isSchemaless(schema.name);
+      // Walk schemas recursively so nested namespaces (e.g. Iceberg
+      // `top.nested`) are enumerated. `segments` is the path of named
+      // (non-schemaless) namespace segments down to this schema.
+      const walkSchema = (schema: DatabaseSchema, segments: string[]): void => {
+        const schemalessDb = segments.length === 0;
+        const isDefaultSchema =
+          segments.length === 1 && segments[0] === conn.default_schema;
+        const schemaQualifier = segments.join(".");
 
         for (const table of schema.tables) {
           let nameToSave: string;
@@ -358,14 +413,14 @@ export const allTablesAtom = atom((get) => {
             continue;
           }
 
-          nameToSave = `${schema.name}.${table.name}`;
+          nameToSave = `${schemaQualifier}.${table.name}`;
 
           if (isDefaultDb && !tableNames.has(nameToSave)) {
             tableNames.set(nameToSave, table);
             continue;
           }
 
-          nameToSave = `${database.name}.${schema.name}.${table.name}`;
+          nameToSave = `${database.name}.${schemaQualifier}.${table.name}`;
 
           if (tableNames.has(nameToSave)) {
             Logger.warn(`Table name collision for ${nameToSave}. Skipping.`);
@@ -373,6 +428,15 @@ export const allTablesAtom = atom((get) => {
             tableNames.set(nameToSave, table);
           }
         }
+
+        // Recurse into nested child namespaces. Children are always named.
+        for (const child of schema.child_schemas ?? []) {
+          walkSchema(child, [...segments, child.name]);
+        }
+      };
+
+      for (const schema of database.schemas) {
+        walkSchema(schema, isSchemaless(schema.name) ? [] : [schema.name]);
       }
     }
   }