fix(py): restore server(data_source=) for Shiny/R parity

pkg-py/CHANGELOG.md

@@ -22,6 +22,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   app_ui = qc.page("Titanic Explorer")  # Core
   ```
 
+* Restored `data_source` (and added `table_name`) parameters on `QueryChat.server()`, matching R's `$server(data_source = )`. This supports the deferred pattern where the data source can only be created inside the Shiny server function (e.g. per-user OAuth-scoped database connections on Posit Connect). (#300)
+
+  ```python
+  qc = QueryChat(None, table_name="my_table")
+
+  def server(input, output, session):
+      qc.server(data_source=conn.table("my_table"), client=chat_client)
+  ```
+
 ### Improvements
 
 * The `"visualize"` tool is now included in the default toolset (`tools=("filter", "query", "visualize")`). If the visualization dependencies are not installed (the `viz` extra), the tool is dropped with a warning instead of raising an `ImportError`.

pkg-py/src/querychat/_querychat_base.py

@@ -103,6 +103,10 @@ def __init__(
         # Track server initialization state for add/remove table validation
         self._server_initialized = False
 
+        # Name to register at .server(data_source=...) time when constructed
+        # with data_source=None (the deferred pattern).
+        self._deferred_table_name: str | None = None
+
         self.tools = normalize_tools(tools, default=DEFAULT_TOOLS)
         self.greeting = greeting.read_text() if isinstance(greeting, Path) else greeting
         self.history = history
@@ -138,6 +142,8 @@ def __init__(
                         "table_name is required when data_source is provided"
                     )
             self.add_table(data_source, table_name, include_in_greeting=True)
+        else:
+            self._deferred_table_name = table_name
 
     def _build_system_prompt(
         self,

pkg-py/src/querychat/_shiny.py

@@ -633,6 +633,8 @@ def page(self, title, *, id: Optional[str] = None, **kwargs):
     def server(
         self,
         *,
+        data_source: IntoFrame | sqlalchemy.Engine | ibis.Table | None = None,
+        table_name: str | None = None,
         client: str | chatlas.Chat | MISSING_TYPE = MISSING,
         history: Optional[bool | HistoryOptions] = None,
         enable_bookmarking: bool | None = None,
@@ -648,6 +650,17 @@ def server(
 
         Parameters
         ----------
+        data_source
+            Optional data source to register for this session, for the
+            deferred pattern where the data source can't be created until the
+            server function runs (e.g., a database connection scoped to
+            per-user OAuth credentials on Posit Connect). Registered under
+            `table_name` if given, otherwise the `table_name` passed to the
+            constructor (when it was created with `data_source=None`), or the
+            first already-registered table.
+        table_name
+            Table name to register `data_source` under. Only used when
+            `data_source` is provided.
         client
             Optional chat client to use for this session. If provided, overrides
             any client set at initialization time for this call only. This is useful
@@ -684,6 +697,27 @@ def server(
                 ".server() must be called within an active Shiny session (i.e., within the server function). "
             )
 
+        if data_source is not None:
+            if table_name is not None:
+                resolved_table_name = table_name
+            elif self._deferred_table_name is not None:
+                resolved_table_name = self._deferred_table_name
+            else:
+                resolved_table_name = next(iter(self._data_sources), None)
+            if resolved_table_name is None:
+                raise ValueError(
+                    "table_name is required when data_source is provided and no "
+                    "table name can be inferred. Pass table_name to .server(), "
+                    "or table_name to the QueryChat constructor, or register a "
+                    "table first with add_table()."
+                )
+            self.add_table(
+                data_source,
+                resolved_table_name,
+                replace=True,
+                include_in_greeting=True,
+            )
+
         self._require_initialized("server")
         resolved_client: chatlas.Chat | None = (
             None

pkg-py/tests/test_server_data_source.py

@@ -0,0 +1,119 @@
+"""Tests for QueryChat.server(data_source=...) parity with R (#300)."""
+
+from unittest.mock import MagicMock
+
+import pandas as pd
+import pytest
+import querychat._shiny as shiny_mod
+
+
+@pytest.fixture(autouse=True)
+def set_dummy_api_key(monkeypatch):
+    monkeypatch.setenv("OPENAI_API_KEY", "sk-dummy-key-for-testing")
+
+
+@pytest.fixture
+def users_df():
+    return pd.DataFrame({"id": [1, 2, 3], "name": ["Alice", "Bob", "Charlie"]})
+
+
+@pytest.fixture
+def other_users_df():
+    return pd.DataFrame({"id": [4, 5], "name": ["Dana", "Eli"]})
+
+
+@pytest.fixture
+def captured_mod_server(monkeypatch):
+    """Patch mod_server and get_current_session; return list of captured kwargs."""
+    calls = []
+
+    def fake_mod_server(*args, **kwargs):
+        calls.append(kwargs)
+        return MagicMock()
+
+    monkeypatch.setattr(shiny_mod, "mod_server", fake_mod_server)
+    monkeypatch.setattr(shiny_mod, "get_current_session", lambda: MagicMock())
+    return calls
+
+
+class TestServerDataSourceRegistersDeferredTable:
+    def test_registers_deferred_table_by_constructor_name(
+        self, users_df, captured_mod_server
+    ):
+        qc = shiny_mod.QueryChat(None, table_name="users")
+        qc.server(data_source=users_df)
+
+        assert qc.table_names() == ["users"]
+        assert list(captured_mod_server[0]["data_sources"].keys()) == ["users"]
+
+    def test_explicit_table_name_overrides_deferred_name(
+        self, users_df, captured_mod_server
+    ):
+        qc = shiny_mod.QueryChat(None, table_name="users")
+        qc.server(data_source=users_df, table_name="people")
+
+        assert qc.table_names() == ["people"]
+
+    def test_falls_back_to_first_existing_table_when_no_name_given(
+        self, users_df, other_users_df, captured_mod_server
+    ):
+        """
+        Mirrors R: server(data_source=) with no deferred/explicit name
+        replaces the first already-registered table.
+        """
+        qc = shiny_mod.QueryChat(users_df, "users")
+        qc.server(data_source=other_users_df)
+
+        assert qc.table_names() == ["users"]
+        registered = captured_mod_server[0]["data_sources"]["users"]
+        assert registered.get_data()["id"].tolist() == [4, 5]
+
+    def test_missing_table_name_raises(self, users_df, captured_mod_server):
+        qc = shiny_mod.QueryChat()
+        with pytest.raises(ValueError, match="table_name"):
+            qc.server(data_source=users_df)
+
+    def test_empty_explicit_table_name_raises_instead_of_falling_back(
+        self, users_df, captured_mod_server
+    ):
+        """
+        An explicit but invalid table_name="" must be validated and rejected,
+        not silently treated as omitted and fall back to the deferred/first
+        table name.
+        """
+        qc = shiny_mod.QueryChat(None, table_name="users")
+        with pytest.raises(ValueError, match="must begin with a letter"):
+            qc.server(data_source=users_df, table_name="")
+
+    def test_data_source_included_in_greeting(self, users_df, captured_mod_server):
+        qc = shiny_mod.QueryChat(None, table_name="users")
+        qc.server(data_source=users_df)
+
+        assert "users" in qc.greeter.tables
+
+    def test_no_data_source_leaves_tables_unchanged(
+        self, users_df, captured_mod_server
+    ):
+        qc = shiny_mod.QueryChat(users_df, "users")
+        qc.server()
+
+        assert qc.table_names() == ["users"]
+
+
+class TestServerDataSourceCurrentSingleSessionLimitation:
+    """
+    Known limitation (tracked by a follow-up PR): server(data_source=...)
+    reuses the public add_table(), so it still hits the "no changes after
+    server initialization" guard on a second session -- the same bug R's
+    existing $server(data_source=) has today. Making this survive a second
+    session is a separate, focused change.
+    """
+
+    def test_second_session_currently_raises(
+        self, users_df, other_users_df, captured_mod_server
+    ):
+        qc = shiny_mod.QueryChat(None, table_name="users")
+        qc.server(data_source=users_df)
+
+        with pytest.raises(RuntimeError, match="Cannot add tables after server"):
+            qc.server(data_source=other_users_df)