diff --git a/pkg-py/CHANGELOG.md b/pkg-py/CHANGELOG.md index 0d994f22..e8be2fdc 100644 --- a/pkg-py/CHANGELOG.md +++ b/pkg-py/CHANGELOG.md @@ -31,7 +31,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 qc.server(data_source=conn.table("my_table"), client=chat_client) ``` - Registering a table this way is no longer blocked by an earlier session having already registered one. + Registering a table this way is no longer blocked by an earlier session having already registered one, and no longer tears down a still-in-use data source from an earlier session (replacing it via `server(data_source=)` leaves the replaced source's cleanup to whoever created it). ### Improvements diff --git a/pkg-py/src/querychat/_querychat_base.py b/pkg-py/src/querychat/_querychat_base.py index 851fab66..b0a7fb8a 100644 --- a/pkg-py/src/querychat/_querychat_base.py +++ b/pkg-py/src/querychat/_querychat_base.py @@ -505,6 +505,7 @@ def _add_or_replace_table( *, replace: bool, include_in_greeting: bool, + cleanup_replaced: bool = True, ) -> None: """ Stage a table and rebuild the system prompt/executor cache. @@ -513,6 +514,16 @@ def _add_or_replace_table( directly by ``.server(data_source=...)`` so that each session can register (or replace) its own table even after an earlier session's ``.server()`` call has already set ``_server_initialized``. + + ``cleanup_replaced=False`` must be used for that per-session + replacement: a table replaced here may still be in active use by an + earlier, already-running session (e.g. its own + ``DataSourceExecutor`` holds a live reference to it), so closing/ + disposing it here would pull the resource out from under that + session. Cleaning it up is then the caller's own responsibility + (e.g. via ``session.on_ended()`` in the code that created it). The + default (``True``) preserves :meth:`add_table`'s existing behavior, + where a config-time replacement has exactly one owner. """ if not isinstance(include_in_greeting, bool): raise TypeError( @@ -549,11 +560,12 @@ def _add_or_replace_table( old_source = self._data_sources.get(table_name) self._data_sources = next_data_sources - if old_source is not None and old_source is not normalized: + if cleanup_replaced and old_source is not None and old_source is not normalized: old_source.cleanup() if self._query_executor is not None: - with contextlib.suppress(Exception): - self._query_executor.cleanup() + if cleanup_replaced: + with contextlib.suppress(Exception): + self._query_executor.cleanup() self._query_executor = None if include_in_greeting and table_name not in self.greeter.tables: diff --git a/pkg-py/src/querychat/_shiny.py b/pkg-py/src/querychat/_shiny.py index 3871aa70..0e2d1b6e 100644 --- a/pkg-py/src/querychat/_shiny.py +++ b/pkg-py/src/querychat/_shiny.py @@ -716,6 +716,7 @@ def server( resolved_table_name, replace=True, include_in_greeting=True, + cleanup_replaced=False, ) self._require_initialized("server") diff --git a/pkg-py/tests/test_server_data_source.py b/pkg-py/tests/test_server_data_source.py index 936355cd..125eb96a 100644 --- a/pkg-py/tests/test_server_data_source.py +++ b/pkg-py/tests/test_server_data_source.py @@ -1,6 +1,6 @@ """Tests for QueryChat.server(data_source=...) parity with R (#300).""" -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import pandas as pd import pytest @@ -129,3 +129,71 @@ def test_add_table_still_blocked_after_server_init( with pytest.raises(RuntimeError, match="Cannot add tables after server"): qc.add_table(other_users_df, "other") + + +class TestServerDataSourceCleanupSafety: + def test_second_session_does_not_clean_up_first_sessions_source( + self, users_df, other_users_df, captured_mod_server + ): + """ + A second session's server(data_source=...) call must not tear down + the DataSource object an earlier, still-running session's own + DataSourceExecutor holds a live reference to (e.g. closing a DuckDB + connection or disposing a SQLAlchemy engine out from under it). + """ + qc = shiny_mod.QueryChat(None, table_name="users") + + qc.server(data_source=users_df) + first_source = qc._data_sources["users"] + + with patch.object(first_source, "cleanup") as mock_cleanup: + qc.server(data_source=other_users_df) + mock_cleanup.assert_not_called() + + def test_public_add_table_replace_still_cleans_up_old_source( + self, users_df, other_users_df + ): + """ + Config-time add_table(replace=True) (before any session starts) has + exactly one owner for the replaced table, so its existing + cleanup-on-replace behavior must be unchanged. + """ + qc = shiny_mod.QueryChat(users_df, "users") + first_source = qc._data_sources["users"] + + with patch.object(first_source, "cleanup") as mock_cleanup: + qc.add_table(other_users_df, "users", replace=True) + mock_cleanup.assert_called_once() + + def test_second_session_does_not_clean_up_first_sessions_query_executor( + self, users_df, other_users_df, captured_mod_server + ): + """ + A second session's server(data_source=...) call must not close the + cached QueryExecutor an earlier, still-running session's chat has + already captured (e.g. via _create_session_client) and is actively + querying through. + """ + qc = shiny_mod.QueryChat(None, table_name="users") + + qc.server(data_source=users_df) + first_executor = qc._require_query_executor("test") + + with patch.object(first_executor, "cleanup") as mock_cleanup: + qc.server(data_source=other_users_df) + mock_cleanup.assert_not_called() + + def test_public_add_table_replace_still_cleans_up_old_query_executor( + self, users_df, other_users_df + ): + """ + Config-time add_table(replace=True) has exactly one owner, so its + existing cleanup-on-replace behavior for the cached executor must be + unchanged. + """ + qc = shiny_mod.QueryChat(users_df, "users") + first_executor = qc._require_query_executor("test") + + with patch.object(first_executor, "cleanup") as mock_cleanup: + qc.add_table(other_users_df, "users", replace=True) + mock_cleanup.assert_called_once()