@marimo-team/frontend 0.21.2-dev1 → 0.21.2-dev101

package/src/core/kernel/handlers.ts

@@ -3,7 +3,7 @@ import { deserializeLayout } from "@/components/editor/renderers/plugins";
 import type { LayoutType } from "@/components/editor/renderers/types";
 import { Logger } from "@/utils/Logger";
 import { Objects } from "@/utils/objects";
-import type { CellId, UIElementId } from "../cells/ids";
+import type { UIElementId } from "../cells/ids";
 import { type CellData, createCell } from "../cells/types";
 import { type AppConfig, AppConfigSchema } from "../config/config-schema";
 import { UI_ELEMENT_REGISTRY } from "../dom/uiregistry";
@@ -47,7 +47,7 @@ export function buildCellData(data: KernelReadyData): CellData[] {
     const edited = lastCodeRun ? lastCodeRun !== code : false;
 
     return createCell({
-      id: cellId as CellId,
+      id: cellId,
       code,
       edited,
       name: names[i],
@@ -160,7 +160,10 @@ export function handleKernelReady(
   // If resumed, restore UI element values from kernel and we're done
   if (resumed) {
     for (const [objectId, value] of Objects.entries(ui_values || {})) {
-      UI_ELEMENT_REGISTRY.set(objectId as UIElementId, value);
+      // @ts-expect-error - TODO: We need to fix the types for KernelReadyNotification["ui_values"]
+      // It is currently typed as Record<string, any> and should be Record<UIElementId, unknown>
+      const uiObjectId: UIElementId = objectId;
+      UI_ELEMENT_REGISTRY.set(uiObjectId, value);
     }
     return;
   }
@@ -196,7 +199,7 @@ export function handleRemoveUIElements(
   // This removes the element from the registry to (1) clean-up
   // memory and (2) make sure that the old value doesn't get re-used
   // if the same cell-id is later reused for another element.
-  const cellId = data.cell_id as CellId;
+  const cellId = data.cell_id;
   UI_ELEMENT_REGISTRY.removeElementsByCell(cellId);
   VirtualFileTracker.INSTANCE.removeForCellId(cellId);
 }

package/src/core/kernel/messages.ts

@@ -41,6 +41,8 @@ export type DataColumnPreview = NotificationMessageData<"data-column-preview">;
 export type SQLTablePreview = NotificationMessageData<"sql-table-preview">;
 export type SQLTableListPreview =
   NotificationMessageData<"sql-table-list-preview">;
+export type SQLSchemaListPreview =
+  NotificationMessageData<"sql-schema-list-preview">;
 export type ValidateSQLResult = NotificationMessageData<"validate-sql-result">;
 export type SecretKeysResult = NotificationMessageData<"secret-keys-result">;
 export type StartupLogs = NotificationMessageData<"startup-logs">;

package/src/core/network/DeferredRequestRegistry.ts

@@ -1,10 +1,10 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 /* eslint-disable @typescript-eslint/no-redeclare */
+import type { components } from "@marimo-team/marimo-api";
 import { Deferred } from "@/utils/Deferred";
 import { generateUUID } from "@/utils/uuid";
-import type { TypedString } from "../../utils/typed";
 
-export type RequestId = TypedString<"RequestId">;
+export type RequestId = components["schemas"]["RequestId"];
 export const RequestId = {
   create(): RequestId {
     return generateUUID() as RequestId;

package/src/core/network/__tests__/CachingRequestRegistry.test.ts

@@ -1,17 +1,16 @@
 /* Copyright 2026 Marimo. All rights reserved. */
+
 import { beforeEach, describe, expect, it, vi } from "vitest";
+import { requestId } from "@/__tests__/branded";
 import { CachingRequestRegistry } from "../CachingRequestRegistry";
-import {
-  DeferredRequestRegistry,
-  type RequestId,
-} from "../DeferredRequestRegistry";
+import { DeferredRequestRegistry } from "../DeferredRequestRegistry";
 
 vi.mock("@/utils/uuid", () => ({
   generateUUID: vi.fn().mockReturnValue("uuid"),
 }));
 
 describe("CachingRequestRegistry", () => {
-  const REQUEST_ID = "uuid" as RequestId;
+  const REQUEST_ID = requestId("uuid");
   let makeRequestMock = vi.fn();
   let delegate: DeferredRequestRegistry<unknown, unknown>;
   let caching: CachingRequestRegistry<unknown, unknown>;
@@ -31,7 +30,7 @@ describe("CachingRequestRegistry", () => {
 
     // Resolve first request
     delegate.resolve(REQUEST_ID, "response");
-    await expect(p1).resolves.toBe("response");
+    await expect(p1).resolves.toBe(requestId("response"));
 
     // Second call with equivalent request gets served from cache
     const p2 = caching.request({ a: 1 });
@@ -50,8 +49,8 @@ describe("CachingRequestRegistry", () => {
     expect(p1).toStrictEqual(p2);
 
     // Resolve and ensure both resolve to same result
-    delegate.resolve(REQUEST_ID, "ok");
-    await expect(p1).resolves.toBe("ok");
+    delegate.resolve(REQUEST_ID, requestId("ok"));
+    await expect(p1).resolves.toBe(requestId("ok"));
     await expect(p2).resolves.toBe("ok");
   });
 
@@ -67,7 +66,7 @@ describe("CachingRequestRegistry", () => {
     expect(makeRequestMock).toHaveBeenCalledTimes(2);
 
     // Resolve the second request
-    delegate.resolve(REQUEST_ID, "ok");
-    await expect(p2).resolves.toBe("ok");
+    delegate.resolve(REQUEST_ID, requestId("ok"));
+    await expect(p2).resolves.toBe(requestId("ok"));
   });
 });

package/src/core/network/__tests__/DeferredRequestRegistry.test.ts

@@ -1,16 +1,14 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 import { beforeEach, describe, expect, it, vi } from "vitest";
-import {
-  DeferredRequestRegistry,
-  type RequestId,
-} from "../DeferredRequestRegistry";
+import { requestId } from "@/__tests__/branded";
+import { DeferredRequestRegistry } from "../DeferredRequestRegistry";
 
 vi.mock("@/utils/uuid", () => ({
   generateUUID: vi.fn().mockReturnValue("uuid"),
 }));
 
 describe("DeferredRequestRegistry", () => {
-  const REQUEST_ID = "uuid" as RequestId;
+  const REQUEST_ID = requestId("uuid");
   let makeRequestMock = vi.fn();
   let registry: DeferredRequestRegistry<unknown, unknown>;
 
@@ -26,7 +24,7 @@ describe("DeferredRequestRegistry", () => {
 
     // resolve
     registry.resolve(REQUEST_ID, "response");
-    await expect(promise).resolves.toBe("response");
+    await expect(promise).resolves.toBe(requestId("response"));
   });
 
   it("should handle request failure", async () => {

package/src/core/network/requests-lazy.ts

@@ -112,13 +112,14 @@ const ACTIONS: Record<keyof AllRequests, Action> = {
   previewDatasetColumn: "waitForConnectionOpen",
   previewSQLTable: "waitForConnectionOpen",
   previewSQLTableList: "waitForConnectionOpen",
+  previewSQLSchemaList: "waitForConnectionOpen",
   previewDataSourceConnection: "waitForConnectionOpen",
   validateSQL: "waitForConnectionOpen",
   listStorageEntries: "waitForConnectionOpen",
   downloadStorage: "waitForConnectionOpen",
 
   // Sync operations that wait for connection
-  syncCellIds: "waitForConnectionOpen",
+  sendDocumentTransaction: "waitForConnectionOpen",
   sendCodeCompletionRequest: "waitForConnectionOpen",
 };
 

package/src/core/network/requests-network.ts

@@ -1,6 +1,6 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-import { once } from "lodash-es";
+import { once } from "@/utils/once";
 import { getRuntimeManager } from "../runtime/config";
 import { API, createClientWithRuntimeManager } from "./api";
 import { waitForConnectionOpen } from "./connection";
@@ -49,10 +49,10 @@ export function createNetworkRequests(): EditRequests & RunRequests {
         })
         .then(handleResponseReturnNull);
     },
-    syncCellIds: async (request) => {
+    sendDocumentTransaction: async (request) => {
       await waitForConnectionOpen();
       return getClient()
-        .POST("/api/kernel/sync/cell_ids", {
+        .POST("/api/document/transaction", {
           body: request,
           params: getParams(),
         })
@@ -235,6 +235,14 @@ export function createNetworkRequests(): EditRequests & RunRequests {
         })
         .then(handleResponseReturnNull);
     },
+    previewSQLSchemaList: (request) => {
+      return getClient()
+        .POST("/api/datasources/preview_sql_schema_list", {
+          body: request,
+          params: getParams(),
+        })
+        .then(handleResponseReturnNull);
+    },
     previewDataSourceConnection: (request) => {
       return getClient()
         .POST("/api/datasources/preview_datasource_connection", {

package/src/core/network/requests-static.ts

@@ -36,7 +36,7 @@ export function createStaticRequests(): EditRequests & RunRequests {
       return null;
     },
     sendRestart: throwNotInEditMode,
-    syncCellIds: throwNotInEditMode,
+    sendDocumentTransaction: throwNotInEditMode,
     sendRun: throwNotInEditMode,
     sendRunScratchpad: throwNotInEditMode,
     sendRename: throwNotInEditMode,
@@ -56,6 +56,7 @@ export function createStaticRequests(): EditRequests & RunRequests {
     previewDatasetColumn: throwNotInEditMode,
     previewSQLTable: throwNotInEditMode,
     previewSQLTableList: throwNotInEditMode,
+    previewSQLSchemaList: throwNotInEditMode,
     previewDataSourceConnection: throwNotInEditMode,
     validateSQL: throwNotInEditMode,
     openFile: throwNotInEditMode,

package/src/core/network/requests-toasting.tsx

@@ -21,7 +21,7 @@ export function createErrorToastingRequests(
     sendInstantiate: "Failed to instantiate",
     sendFunctionRequest: "Failed to send function request",
     sendRestart: "Failed to restart",
-    syncCellIds: "Failed to sync cell IDs",
+    sendDocumentTransaction: "Failed to sync document transaction",
     sendRun: "Failed to run",
     sendRunScratchpad: "Failed to run scratchpad",
     sendRename: "Failed to rename",
@@ -41,6 +41,7 @@ export function createErrorToastingRequests(
     previewDatasetColumn: "Failed to fetch data sources",
     previewSQLTable: "Failed to fetch SQL table",
     previewSQLTableList: "Failed to fetch SQL table list",
+    previewSQLSchemaList: "Failed to fetch SQL schema list",
     previewDataSourceConnection: "Failed to preview data source connection",
     validateSQL: "Failed to validate SQL",
     openFile: "Failed to open file",

package/src/core/network/types.ts

@@ -59,6 +59,7 @@ export type PreviewDatasetColumnRequest =
   schemas["PreviewDatasetColumnRequest"];
 export type PreviewSQLTableRequest = schemas["PreviewSQLTableRequest"];
 export type ListSQLTablesRequest = schemas["ListSQLTablesRequest"];
+export type ListSQLSchemasRequest = schemas["ListSQLSchemasRequest"];
 export type ListDataSourceConnectionRequest =
   schemas["ListDataSourceConnectionRequest"];
 export type ValidateSQLRequest = schemas["ValidateSQLRequest"];
@@ -79,7 +80,8 @@ export interface SetCellConfigRequest {
 }
 export type UpdateUIElementRequest = schemas["UpdateUIElementRequest"];
 export type ModelRequest = schemas["ModelRequest"];
-export type UpdateCellIdsRequest = schemas["UpdateCellIdsRequest"];
+export type NotebookDocumentTransactionRequest =
+  schemas["NotebookDocumentTransactionRequest"];
 export type UpdateUserConfigRequest = schemas["UpdateUserConfigRequest"];
 export type ShutdownSessionRequest = schemas["ShutdownSessionRequest"];
 export type Snippet = schemas["Snippet"];
@@ -137,7 +139,9 @@ export interface EditRequests {
   saveAppConfig: (request: SaveAppConfigurationRequest) => Promise<null>;
   saveCellConfig: (request: SetCellConfigRequest) => Promise<null>;
   sendRestart: () => Promise<null>;
-  syncCellIds: (request: UpdateCellIdsRequest) => Promise<null>;
+  sendDocumentTransaction: (
+    request: NotebookDocumentTransactionRequest,
+  ) => Promise<null>;
   sendInstallMissingPackages: (
     request: InstallPackagesRequest,
   ) => Promise<null>;
@@ -146,6 +150,7 @@ export interface EditRequests {
   previewDatasetColumn: (request: PreviewDatasetColumnRequest) => Promise<null>;
   previewSQLTable: (request: PreviewSQLTableRequest) => Promise<null>;
   previewSQLTableList: (request: ListSQLTablesRequest) => Promise<null>;
+  previewSQLSchemaList: (request: ListSQLSchemasRequest) => Promise<null>;
   previewDataSourceConnection: (
     request: ListDataSourceConnectionRequest,
   ) => Promise<null>;

package/src/core/rtc/state.ts

@@ -1,6 +1,6 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 import { atomWithStorage } from "jotai/utils";
-import { once } from "lodash-es";
+import { once } from "@/utils/once";
 import { jotaiJsonStorage } from "@/utils/storage/jotai";
 import { getFeatureFlag } from "../config/feature-flag";
 

package/src/core/saving/state.ts

@@ -1,6 +1,7 @@
 /* Copyright 2026 Marimo. All rights reserved. */
+
+import { dequal as isEqual } from "dequal";
 import { atom } from "jotai";
-import { isEqual } from "lodash-es";
 import { arrayShallowEquals } from "@/utils/arrays";
 import { type NotebookState, notebookAtom } from "../cells/cells";
 import { type LayoutState, layoutStateAtom } from "../layout/layout";

package/src/core/slots/slots.ts

@@ -6,4 +6,5 @@ export const slotsController = new SlotzController();
 export const SlotNames = {
   SIDEBAR: "sidebar",
   CONTEXT_AWARE_PANEL: "context-aware-panel",
+  CONTEXT_AWARE_PANEL_HEADER: "context-aware-panel-header",
 };

package/src/core/state/jotai.ts

@@ -1,6 +1,7 @@
 /* Copyright 2026 Marimo. All rights reserved. */
+
+import { dequal as isEqual } from "dequal";
 import { type Atom, atom, createStore, useStore } from "jotai";
-import { isEqual } from "lodash-es";
 import { useEffect } from "react";
 
 /**

package/src/core/static/__tests__/virtual-file-tracker.test.ts

@@ -1,6 +1,6 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 import { describe, expect, it } from "vitest";
-import type { CellId } from "@/core/cells/ids";
+import { cellId } from "@/__tests__/branded";
 import type { OutputMessage } from "@/core/kernel/messages";
 import { findVirtualFiles, VirtualFileTracker } from "../virtual-file-tracker";
 
@@ -103,32 +103,32 @@ describe("findVirtualFiles", () => {
 describe("VirtualFileTracker", () => {
   it("should track virtual files, append, clear", () => {
     const tracker = VirtualFileTracker.INSTANCE;
-    const cellId = "test-cell-id" as CellId;
+    const cid = cellId("test-cell-id");
     tracker.track({
-      cell_id: cellId,
+      cell_id: cid,
       output: {
         mimetype: "text/html",
         data: "Some text /@file/test-file.js more text",
       } as OutputMessage,
     });
-    expect(tracker.virtualFiles.get(cellId)).toEqual(
+    expect(tracker.virtualFiles.get(cid)).toEqual(
       new Set(["/@file/test-file.js"]),
     );
 
     // can append
     tracker.track({
-      cell_id: cellId,
+      cell_id: cid,
       output: {
         mimetype: "text/html",
         data: "Some text /@file/another-file.txt more text",
       } as OutputMessage,
     });
-    expect(tracker.virtualFiles.get(cellId)).toEqual(
+    expect(tracker.virtualFiles.get(cid)).toEqual(
       new Set(["/@file/test-file.js", "/@file/another-file.txt"]),
     );
 
     // can clear
-    tracker.removeForCellId(cellId);
-    expect(tracker.virtualFiles.get(cellId)).toEqual(undefined);
+    tracker.removeForCellId(cid);
+    expect(tracker.virtualFiles.get(cid)).toEqual(undefined);
   });
 });

package/src/core/static/virtual-file-tracker.ts

@@ -29,7 +29,7 @@ export class VirtualFileTracker {
 
   track(message: Pick<CellMessage, "cell_id" | "output">): void {
     const output = message.output;
-    const cellId = message.cell_id as CellId;
+    const cellId = message.cell_id;
     if (!output) {
       return;
     }

package/src/core/storage/__tests__/state.test.ts

@@ -1,5 +1,6 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 import { beforeEach, describe, expect, it } from "vitest";
+import { variableName } from "@/__tests__/branded";
 import type { VariableName } from "../../variables/types";
 import { exportedForTesting } from "../state";
 import type { StorageEntry, StorageNamespace, StorageState } from "../types";
@@ -7,7 +8,7 @@ import type { StorageEntry, StorageNamespace, StorageState } from "../types";
 const { initialState, reducer, createActions } = exportedForTesting;
 
 function makeNamespace(
-  overrides: Partial<StorageNamespace> & { name: string },
+  overrides: Partial<StorageNamespace> & { name: VariableName },
 ): StorageNamespace {
   return {
     backendType: overrides.backendType ?? "obstore",
@@ -53,8 +54,14 @@ describe("storage state", () => {
 
   describe("setNamespaces", () => {
     it("should add namespaces from an empty state", () => {
-      const ns1 = makeNamespace({ name: "my_s3", protocol: "s3" });
-      const ns2 = makeNamespace({ name: "my_gcs", protocol: "gcs" });
+      const ns1 = makeNamespace({
+        name: variableName("my_s3"),
+        protocol: "s3",
+      });
+      const ns2 = makeNamespace({
+        name: variableName("my_gcs"),
+        protocol: "gcs",
+      });
 
       actions.setNamespaces({ namespaces: [ns1, ns2] });
 
@@ -63,15 +70,18 @@ describe("storage state", () => {
 
     it("should merge namespaces by name, replacing existing ones", () => {
       const nsOld = makeNamespace({
-        name: "my_s3",
+        name: variableName("my_s3"),
         protocol: "s3",
         rootPath: "/old",
       });
-      const nsOther = makeNamespace({ name: "my_gcs", protocol: "gcs" });
+      const nsOther = makeNamespace({
+        name: variableName("my_gcs"),
+        protocol: "gcs",
+      });
       actions.setNamespaces({ namespaces: [nsOld, nsOther] });
 
       const nsUpdated = makeNamespace({
-        name: "my_s3",
+        name: variableName("my_s3"),
         protocol: "s3",
         rootPath: "/new",
       });
@@ -82,10 +92,10 @@ describe("storage state", () => {
     });
 
     it("should add new namespaces alongside existing ones", () => {
-      const ns1 = makeNamespace({ name: "ns1" });
+      const ns1 = makeNamespace({ name: variableName("ns1") });
       actions.setNamespaces({ namespaces: [ns1] });
 
-      const ns2 = makeNamespace({ name: "ns2" });
+      const ns2 = makeNamespace({ name: variableName("ns2") });
       actions.setNamespaces({ namespaces: [ns2] });
 
       expect(state.namespaces).toHaveLength(2);
@@ -101,7 +111,7 @@ describe("storage state", () => {
       });
 
       actions.setNamespaces({
-        namespaces: [makeNamespace({ name: "ns" })],
+        namespaces: [makeNamespace({ name: variableName("ns") })],
       });
 
       expect(state.entriesByPath.get("ns::")).toEqual([entry]);
@@ -205,7 +215,7 @@ describe("storage state", () => {
     });
 
     it("should not affect namespaces", () => {
-      const ns = makeNamespace({ name: "ns" });
+      const ns = makeNamespace({ name: variableName("ns") });
       actions.setNamespaces({ namespaces: [ns] });
 
       actions.setEntries({
@@ -263,7 +273,7 @@ describe("storage state", () => {
     });
 
     it("should not affect namespaces", () => {
-      const ns = makeNamespace({ name: "my_s3" });
+      const ns = makeNamespace({ name: variableName("my_s3") });
       actions.setNamespaces({ namespaces: [ns] });
       actions.setEntries({
         namespace: "my_s3",
@@ -293,31 +303,31 @@ describe("storage state", () => {
 
   describe("filterFromVariables", () => {
     it("should keep namespaces whose variable is still in scope", () => {
-      const ns1 = makeNamespace({ name: "var_a" });
-      const ns2 = makeNamespace({ name: "var_b" });
+      const ns1 = makeNamespace({ name: variableName("var_a") });
+      const ns2 = makeNamespace({ name: variableName("var_b") });
       actions.setNamespaces({ namespaces: [ns1, ns2] });
 
       actions.filterFromVariables([
-        "var_a" as VariableName,
-        "var_b" as VariableName,
+        variableName("var_a"),
+        variableName("var_b"),
       ]);
 
       expect(state.namespaces).toEqual([ns1, ns2]);
     });
 
     it("should remove namespaces whose variable is no longer in scope", () => {
-      const ns1 = makeNamespace({ name: "var_a" });
-      const ns2 = makeNamespace({ name: "var_b" });
+      const ns1 = makeNamespace({ name: variableName("var_a") });
+      const ns2 = makeNamespace({ name: variableName("var_b") });
       actions.setNamespaces({ namespaces: [ns1, ns2] });
 
-      actions.filterFromVariables(["var_a" as VariableName]);
+      actions.filterFromVariables([variableName("var_a")]);
 
       expect(state.namespaces).toEqual([ns1]);
     });
 
     it("should remove all named namespaces when given an empty variable list", () => {
-      const ns1 = makeNamespace({ name: "var_a" });
-      const ns2 = makeNamespace({ name: "var_b" });
+      const ns1 = makeNamespace({ name: variableName("var_a") });
+      const ns2 = makeNamespace({ name: variableName("var_b") });
       actions.setNamespaces({ namespaces: [ns1, ns2] });
 
       actions.filterFromVariables([]);
@@ -326,7 +336,7 @@ describe("storage state", () => {
     });
 
     it("should not affect entriesByPath", () => {
-      const ns = makeNamespace({ name: "ns" });
+      const ns = makeNamespace({ name: variableName("ns") });
       actions.setNamespaces({ namespaces: [ns] });
 
       const entry = makeEntry({ path: "file.txt" });

package/src/core/storage/state.ts

@@ -71,7 +71,7 @@ const {
     const names = new Set(variableNames);
     // Filter out namespaces whose backing variable is no longer in scope
     const namespaces = state.namespaces.filter((ns) => {
-      return names.has(ns.name as VariableName);
+      return names.has(ns.name);
     });
     return { ...state, namespaces };
   },

package/src/core/variables/__tests__/state.test.ts

@@ -1,19 +1,19 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 import { beforeEach, describe, expect, it } from "vitest";
-import type { CellId } from "@/core/cells/ids";
+import { cellId, variableName } from "@/__tests__/branded";
 import { exportedForTesting } from "../state";
-import type { VariableName, Variables } from "../types";
+import type { Variables } from "../types";
 
 const { initialState, reducer, createActions } = exportedForTesting;
 
 const CellIds = {
-  a: "a" as CellId,
-  b: "b" as CellId,
+  a: cellId("a"),
+  b: cellId("b"),
 };
 
 const Names = {
-  x: "x" as VariableName,
-  y: "y" as VariableName,
+  x: variableName("x"),
+  y: variableName("y"),
 };
 
 describe("cell reducer", () => {

package/src/core/variables/types.ts

@@ -1,9 +1,9 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-import type { TypedString } from "../../utils/typed";
+import type { components } from "@marimo-team/marimo-api";
 import type { CellId } from "../cells/ids";
 
-export type VariableName = TypedString<"VariableName">;
+export type VariableName = components["schemas"]["VariableName"];
 
 export interface Variable {
   name: VariableName;

package/src/core/wasm/__tests__/state.test.ts

@@ -2,9 +2,9 @@
 
 import { createRef } from "react";
 import { describe, expect, it } from "vitest";
+import { cellId } from "@/__tests__/branded";
 import type { NotebookState } from "@/core/cells/cells";
 import { initialNotebookState, notebookAtom } from "@/core/cells/cells";
-import type { CellId } from "@/core/cells/ids";
 import { createCell, createCellRuntimeState } from "@/core/cells/types";
 import type { OutputMessage } from "@/core/kernel/messages";
 import { store } from "@/core/state/jotai";
@@ -17,17 +17,17 @@ describe("hasAnyOutputAtom", () => {
   ): NotebookState => ({
     ...initialNotebookState(),
     cellIds: new MultiColumn([
-      CollapsibleTree.from(outputs.map((_, i) => `${i}` as CellId)),
+      CollapsibleTree.from(outputs.map((_, i) => cellId(`${i}`))),
     ]),
     cellData: Object.fromEntries(
       outputs.map((_, i) => [
-        `${i}` as CellId,
-        createCell({ id: `${i}` as CellId }),
+        cellId(`${i}`),
+        createCell({ id: cellId(`${i}`) }),
       ]),
     ),
     cellRuntime: Object.fromEntries(
       outputs.map((output, i) => [
-        `${i}` as CellId,
+        `${i}`,
         createCellRuntimeState({
           output,
           status: "queued",
@@ -36,7 +36,7 @@ describe("hasAnyOutputAtom", () => {
       ]),
     ),
     cellHandles: Object.fromEntries(
-      outputs.map((_, i) => [`${i}` as CellId, createRef()]),
+      outputs.map((_, i) => [cellId(`${i}`), createRef()]),
     ),
   });
 
@@ -90,8 +90,8 @@ describe("hasAnyOutputAtom", () => {
 
   it("should return true when all outputs are idle", () => {
     const notebookState = createNotebookState([null, null]);
-    const cellId0 = "0" as CellId;
-    const cellId1 = "1" as CellId;
+    const cellId0 = cellId("0");
+    const cellId1 = cellId("1");
     // Some idle cell, so returns false
     store.set(notebookAtom, {
       ...notebookState,

package/src/core/wasm/bridge.ts

@@ -523,6 +523,16 @@ export class PyodideBridge implements RunRequests, EditRequests {
     return null;
   };
 
+  previewSQLSchemaList: EditRequests["previewSQLSchemaList"] = async (
+    request,
+  ) => {
+    await this.putControlRequest({
+      type: "list-sql-schemas",
+      ...request,
+    });
+    return null;
+  };
+
   previewDataSourceConnection: EditRequests["previewDataSourceConnection"] =
     async (request) => {
       await this.putControlRequest({
@@ -548,7 +558,7 @@ export class PyodideBridge implements RunRequests, EditRequests {
     return null;
   };
 
-  syncCellIds = () => Promise.resolve(null);
+  sendDocumentTransaction = () => Promise.resolve(null);
 
   addPackage: EditRequests["addPackage"] = async (request) => {
     return this.rpc.proxy.request.addPackage(request);