diff --git a/pkg-py/CHANGELOG.md b/pkg-py/CHANGELOG.md index 635fb3d4..f76d334f 100644 --- a/pkg-py/CHANGELOG.md +++ b/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`. diff --git a/pkg-py/src/querychat/_querychat_base.py b/pkg-py/src/querychat/_querychat_base.py index 3de20789..8c8cb599 100644 --- a/pkg-py/src/querychat/_querychat_base.py +++ b/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, diff --git a/pkg-py/src/querychat/_shiny.py b/pkg-py/src/querychat/_shiny.py index 661b98f1..7bb5fd6e 100644 --- a/pkg-py/src/querychat/_shiny.py +++ b/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 diff --git a/pkg-py/tests/test_server_data_source.py b/pkg-py/tests/test_server_data_source.py new file mode 100644 index 00000000..19efbb75 --- /dev/null +++ b/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)