update datasource connection when manually introspecting

marimo/_messaging/ops.py

@@ -587,11 +587,20 @@ class Datasets(Op, tag="datasets"):
     clear_channel: Optional[DataTableSource] = None
 
 
+class SQLMetadata(msgspec.Struct, tag="sql-metadata"):
+    """Metadata for a SQL database."""
+
+    connection: str
+    database: str
+    schema: str
+
+
 class SQLTablePreview(Op, tag="sql-table-preview"):
     """Preview of a table in a SQL database."""
 
     name: ClassVar[str] = "sql-table-preview"
     request_id: RequestId
+    metadata: SQLMetadata
     table: Optional[DataTable]
     error: Optional[str] = None
 
@@ -601,6 +610,7 @@ class SQLTableListPreview(Op, tag="sql-table-list-preview"):
 
     name: ClassVar[str] = "sql-table-list-preview"
     request_id: RequestId
+    metadata: SQLMetadata
     tables: list[DataTable] = msgspec.field(default_factory=list)
     error: Optional[str] = None
 

marimo/_runtime/runtime.py

@@ -67,6 +67,7 @@
     PackageStatusType,
     RemoveUIElements,
     SecretKeysResult,
+    SQLMetadata,
     SQLTableListPreview,
     SQLTablePreview,
     ValidateSQLResult,
@@ -2454,11 +2455,19 @@ async def preview_sql_table(self, request: PreviewSQLTableRequest) -> None:
         database_name = request.database
         schema_name = request.schema
         table_name = request.table_name
+        sql_metadata = SQLMetadata(
+            connection=variable_name,
+            database=database_name,
+            schema=schema_name,
+        )
 
         engine, error = self.get_engine_catalog(variable_name)
         if error is not None or engine is None:
             SQLTablePreview(
-                request_id=request.request_id, table=None, error=error
+                request_id=request.request_id,
+                table=None,
+                error=error,
+                metadata=sql_metadata,
             ).broadcast()
             return
 
@@ -2470,7 +2479,9 @@ async def preview_sql_table(self, request: PreviewSQLTableRequest) -> None:
             )
 
             SQLTablePreview(
-                request_id=request.request_id, table=table
+                request_id=request.request_id,
+                table=table,
+                metadata=sql_metadata,
             ).broadcast()
         except Exception as e:
             LOGGER.exception(
@@ -2482,6 +2493,7 @@ async def preview_sql_table(self, request: PreviewSQLTableRequest) -> None:
                 request_id=request.request_id,
                 table=None,
                 error="Failed to get table details: " + str(e),
+                metadata=sql_metadata,
             ).broadcast()
 
     @kernel_tracer.start_as_current_span("preview_sql_table_list")
@@ -2499,11 +2511,19 @@ async def preview_sql_table_list(
         variable_name = cast(VariableName, request.engine)
         database_name = request.database
         schema_name = request.schema
+        sql_metadata = SQLMetadata(
+            connection=variable_name,
+            database=database_name,
+            schema=schema_name,
+        )
 
         engine, error = self.get_engine_catalog(variable_name)
         if error is not None or engine is None:
             SQLTableListPreview(
-                request_id=request.request_id, tables=[], error=error
+                request_id=request.request_id,
+                tables=[],
+                error=error,
+                metadata=sql_metadata,
             ).broadcast()
             return
 
@@ -2514,7 +2534,9 @@ async def preview_sql_table_list(
                 include_table_details=False,
             )
             SQLTableListPreview(
-                request_id=request.request_id, tables=table_list
+                request_id=request.request_id,
+                tables=table_list,
+                metadata=sql_metadata,
             ).broadcast()
         except Exception as e:
             LOGGER.exception(
@@ -2524,7 +2546,8 @@ async def preview_sql_table_list(
                 request_id=request.request_id,
                 tables=[],
                 error="Failed to get table list: " + str(e),
-            )
+                metadata=sql_metadata,
+            ).broadcast()
 
     @kernel_tracer.start_as_current_span("preview_datasource_connection")
     async def preview_datasource_connection(

marimo/_server/session/session_view.py

@@ -15,6 +15,8 @@
     Interrupted,
     MessageOperation,
     SendUIElementMessage,
+    SQLTableListPreview,
+    SQLTablePreview,
     StartupLogs,
     UpdateCellCodes,
     UpdateCellIdsRequest,
@@ -32,6 +34,10 @@
     SetUIElementValueRequest,
     SyncGraphRequest,
 )
+from marimo._sql.connection_utils import (
+    update_table_in_connection,
+    update_table_list_in_connection,
+)
 from marimo._sql.engines.duckdb import INTERNAL_DUCKDB_ENGINE
 from marimo._types.ids import CellId_t, WidgetModelId
 from marimo._utils.lists import as_list
@@ -248,6 +254,27 @@ def add_operation(self, operation: MessageOperation) -> None:
                 connections=list(connections.values())
             )
 
+        elif isinstance(operation, SQLTablePreview):
+            sql_table_preview = operation
+            sql_metadata = sql_table_preview.metadata
+            table_preview_connections = self.data_connectors.connections
+            if sql_table_preview.table is not None:
+                update_table_in_connection(
+                    table_preview_connections,
+                    sql_metadata,
+                    sql_table_preview.table,
+                )
+
+        elif isinstance(operation, SQLTableListPreview):
+            sql_table_list_preview = operation
+            sql_metadata = sql_table_list_preview.metadata
+            table_list_connections = self.data_connectors.connections
+            update_table_list_in_connection(
+                table_list_connections,
+                sql_metadata,
+                sql_table_list_preview.tables,
+            )
+
         elif isinstance(operation, UpdateCellIdsRequest):
             self.cell_ids = operation
 

marimo/_sql/connection_utils.py

@@ -0,0 +1,62 @@
+# Copyright 2025 Marimo. All rights reserved.
+
+from marimo._data.models import DataSourceConnection, DataTable
+from marimo._messaging.ops import SQLMetadata
+
+
+def update_table_in_connection(
+    connections: list[DataSourceConnection],
+    sql_metadata: SQLMetadata,
+    updated_table: DataTable,
+) -> None:
+    """Update a table in the connection hierarchy in-place
+
+    Args:
+        connections: List of data source connections
+        sql_metadata: SQL metadata containing connection, database, schema info
+        updated_table: The updated table to replace the existing one
+    """
+    for connection in connections:
+        if connection.name != sql_metadata.connection:
+            continue
+
+        for database in connection.databases:
+            if database.name != sql_metadata.database:
+                continue
+
+            for schema in database.schemas:
+                if schema.name != sql_metadata.schema:
+                    continue
+
+                for i, table in enumerate(schema.tables):
+                    if table.name == updated_table.name:
+                        schema.tables[i] = updated_table
+                        return
+
+
+def update_table_list_in_connection(
+    connections: list[DataSourceConnection],
+    sql_metadata: SQLMetadata,
+    updated_table_list: list[DataTable],
+) -> None:
+    """Update a list of tables in the connection hierarchy, updates in-place.
+
+    Args:
+        connections: List of data source connections
+        sql_metadata: SQL metadata containing connection, database, schema info
+        updated_table_list: The updated list of tables to replace the existing ones
+    """
+    for connection in connections:
+        if connection.name != sql_metadata.connection:
+            continue
+
+        for database in connection.databases:
+            if database.name != sql_metadata.database:
+                continue
+
+            for schema in database.schemas:
+                if schema.name != sql_metadata.schema:
+                    continue
+
+                schema.tables = updated_table_list
+                return

packages/openapi/api.yaml

@@ -2891,6 +2891,25 @@ components:
       - watcher_on_save
       title: RuntimeConfig
       type: object
+    SQLMetadata:
+      description: Metadata for a SQL database.
+      properties:
+        connection:
+          type: string
+        database:
+          type: string
+        schema:
+          type: string
+        type:
+          enum:
+          - sql-metadata
+      required:
+      - type
+      - connection
+      - database
+      - schema
+      title: SQLMetadata
+      type: object
     SQLTableListPreview:
       description: Preview of a list of tables in a schema.
       properties:
@@ -2899,6 +2918,8 @@ components:
           - type: string
           - type: 'null'
           default: null
+        metadata:
+          $ref: '#/components/schemas/SQLMetadata'
         op:
           enum:
           - sql-table-list-preview
@@ -2912,6 +2933,7 @@ components:
       required:
       - op
       - request_id
+      - metadata
       title: SQLTableListPreview
       type: object
     SQLTablePreview:
@@ -2922,6 +2944,8 @@ components:
           - type: string
           - type: 'null'
           default: null
+        metadata:
+          $ref: '#/components/schemas/SQLMetadata'
         op:
           enum:
           - sql-table-preview
@@ -2934,6 +2958,7 @@ components:
       required:
       - op
       - request_id
+      - metadata
       - table
       title: SQLTablePreview
       type: object

packages/openapi/src/api.ts

@@ -4342,13 +4342,25 @@ export interface components {
       /** @enum {unknown} */
       watcher_on_save: "autorun" | "lazy";
     };
+    /**
+     * SQLMetadata
+     * @description Metadata for a SQL database.
+     */
+    SQLMetadata: {
+      connection: string;
+      database: string;
+      schema: string;
+      /** @enum {unknown} */
+      type: "sql-metadata";
+    };
     /**
      * SQLTableListPreview
      * @description Preview of a list of tables in a schema.
      */
     SQLTableListPreview: {
       /** @default null */
       error?: string | null;
+      metadata: components["schemas"]["SQLMetadata"];
       /** @enum {unknown} */
       op: "sql-table-list-preview";
       request_id: string;
@@ -4362,6 +4374,7 @@ export interface components {
     SQLTablePreview: {
       /** @default null */
       error?: string | null;
+      metadata: components["schemas"]["SQLMetadata"];
       /** @enum {unknown} */
       op: "sql-table-preview";
       request_id: string;