update datasource connection when manually introspecting

Author: Light2Dark

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

pixi.lock

@@ -5,7 +5,7 @@
 from unittest.mock import patch
 
 from marimo._ast.cell import RuntimeStateType
-from marimo._data.models import DataTable, DataTableColumn
+from marimo._data.models import Database, DataTable, DataTableColumn, Schema
 from marimo._messaging.cell_output import CellChannel, CellOutput
 from marimo._messaging.msgspec_encoder import asdict as serialize
 from marimo._messaging.ops import (
@@ -15,6 +15,9 @@
     DataSourceConnections,
     InstallingPackageAlert,
     SendUIElementMessage,
+    SQLMetadata,
+    SQLTableListPreview,
+    SQLTablePreview,
     StartupLogs,
     UpdateCellCodes,
     UpdateCellIdsRequest,
@@ -32,7 +35,7 @@
 )
 from marimo._server.session.session_view import SessionView
 from marimo._sql.engines.duckdb import INTERNAL_DUCKDB_ENGINE
-from marimo._types.ids import CellId_t, WidgetModelId
+from marimo._types.ids import CellId_t, RequestId, VariableName, WidgetModelId
 from marimo._utils.parse_dataclass import parse_raw
 
 cell_id = CellId_t("cell_1")
@@ -558,6 +561,125 @@ def test_add_data_source_connections(session_view: SessionView) -> None:
     assert INTERNAL_DUCKDB_ENGINE in session_view_names
 
 
+def test_add_sql_table_previews() -> None:
+    session_view = SessionView()
+
+    # Add initial connections
+    session_view.add_raw_operation(
+        serialize_kernel_message(
+            DataSourceConnections(
+                connections=[
+                    DataSourceConnection(
+                        source="duckdb",
+                        name="connection1",
+                        dialect="duckdb",
+                        display_name="duckdb (connection1)",
+                        databases=[
+                            Database(
+                                name="db1",
+                                dialect="duckdb",
+                                schemas=[
+                                    Schema(
+                                        name="db1",
+                                        tables=[
+                                            DataTable(
+                                                name="table1",
+                                                source_type="connection",
+                                                source="db1",
+                                                columns=[],
+                                                num_rows=0,
+                                                num_columns=0,
+                                                variable_name=None,
+                                            )
+                                        ],
+                                    )
+                                ],
+                            )
+                        ],
+                    )
+                ],
+            )
+        )
+    )
+
+    session_view_connections = session_view.data_connectors.connections
+    assert session_view_connections[0].databases[0].schemas[0].tables == [
+        DataTable(
+            source_type="connection",
+            source="db1",
+            name="table1",
+            num_rows=0,
+            num_columns=0,
+            variable_name=None,
+            columns=[],
+        )
+    ]
+
+    session_view.add_raw_operation(
+        serialize_kernel_message(
+            SQLTablePreview(
+                metadata=SQLMetadata(
+                    connection="connection1", database="db1", schema="db1"
+                ),
+                request_id=RequestId("request_id"),
+                table=DataTable(
+                    name="table1",
+                    source_type="connection",
+                    source="db1",
+                    num_rows=10,  # Updated
+                    num_columns=0,
+                    variable_name=None,
+                    columns=[],
+                ),
+            )
+        )
+    )
+    session_view_connections = session_view.data_connectors.connections
+    assert (
+        session_view_connections[0].databases[0].schemas[0].tables[0].num_rows
+        == 10
+    )
+
+    # Add sql table preview list
+    session_view.add_raw_operation(
+        serialize_kernel_message(
+            SQLTableListPreview(
+                metadata=SQLMetadata(
+                    connection="connection1", database="db1", schema="db1"
+                ),
+                request_id=RequestId("request_id"),
+                tables=[
+                    DataTable(
+                        name="table2",
+                        source_type="connection",
+                        source="db1",
+                        num_rows=20,
+                        num_columns=10,
+                        variable_name=VariableName("var"),
+                        columns=[],
+                    )
+                ],
+            )
+        )
+    )
+
+    assert session_view_connections[0].databases[0].schemas[0].tables == [
+        DataTable(
+            source_type="connection",
+            source="db1",
+            name="table2",
+            num_rows=20,
+            num_columns=10,
+            variable_name=VariableName("var"),
+            columns=[],
+            engine=None,
+            type="table",
+            primary_keys=None,
+            indexes=None,
+        )
+    ]
+
+
 def test_add_cell_op(session_view: SessionView) -> None:
     session_view.add_raw_operation(
         serialize_kernel_message(