From 7eccd843f92e409f0e5491fe0bd1c30f2dc98229 Mon Sep 17 00:00:00 2001 From: Carson Date: Sat, 12 Sep 2026 18:09:48 -0500 Subject: [PATCH 01/35] feat(py): add TableSet value object and build_query_executor() --- pkg-py/src/querychat/_query_executor.py | 25 ++++++++ pkg-py/src/querychat/_table_set.py | 55 ++++++++++++++++++ pkg-py/tests/test_table_set.py | 76 +++++++++++++++++++++++++ 3 files changed, 156 insertions(+) create mode 100644 pkg-py/src/querychat/_table_set.py create mode 100644 pkg-py/tests/test_table_set.py diff --git a/pkg-py/src/querychat/_query_executor.py b/pkg-py/src/querychat/_query_executor.py index f1c4e2287..a3f302bdd 100644 --- a/pkg-py/src/querychat/_query_executor.py +++ b/pkg-py/src/querychat/_query_executor.py @@ -19,6 +19,8 @@ from ._utils import check_query if TYPE_CHECKING: + from collections.abc import Mapping + from ._datasource import DataFrameSource, DataSource, PolarsLazySource @@ -314,3 +316,26 @@ def get_dataframe_backend_name(source: DataFrameSource) -> str: return nw.get_native_namespace( nw.from_native(source.get_data(), eager_only=True) ).__name__ + + +def build_query_executor(sources: Mapping[str, DataSource]) -> QueryExecutor: + """Pick the executor for a compatible group of sources.""" + from ._datasource import DataFrameSource, PolarsLazySource + + validate_source_group_compatibility(dict(sources)) + + if len(sources) == 1: + return DataSourceExecutor(dict(sources)) + + first_source = next(iter(sources.values())) + + if isinstance(first_source, DataFrameSource): + return DuckDBExecutor( + {n: s for n, s in sources.items() if isinstance(s, DataFrameSource)} + ) + if isinstance(first_source, PolarsLazySource): + return PolarsSQLExecutor( + {n: s for n, s in sources.items() if isinstance(s, PolarsLazySource)} + ) + + return DataSourceExecutor(dict(sources)) diff --git a/pkg-py/src/querychat/_table_set.py b/pkg-py/src/querychat/_table_set.py new file mode 100644 index 000000000..b33930c81 --- /dev/null +++ b/pkg-py/src/querychat/_table_set.py @@ -0,0 +1,55 @@ +"""Immutable bundle of the tables a chat can query.""" + +from __future__ import annotations + +from functools import cached_property +from types import MappingProxyType +from typing import TYPE_CHECKING + +from ._query_executor import QueryExecutor, build_query_executor + +if TYPE_CHECKING: + from collections.abc import Mapping + + from ._datasource import DataSource + from ._system_prompt import QueryChatSystemPrompt + + +class TableSet: + """ + The tables a chat can query, plus the prompt and executor built from them. + + A ``TableSet`` is never mutated after construction. ``QueryChatBase`` + holds one built from ``add_table()`` calls; ``QueryChat.server()`` derives + a second one when a session registers its own table, so a running session + never observes changes made after it started. + """ + + def __init__( + self, + data_sources: Mapping[str, DataSource], + system_prompt: QueryChatSystemPrompt, + ) -> None: + if not data_sources: + raise ValueError("TableSet requires at least one data source") + self.data_sources: Mapping[str, DataSource] = MappingProxyType( + dict(data_sources) + ) + self.system_prompt = system_prompt + + @cached_property + def executor(self) -> QueryExecutor: + return build_query_executor(self.data_sources) + + @property + def executor_built(self) -> bool: + return "executor" in self.__dict__ + + @property + def table_names(self) -> list[str]: + return list(self.data_sources) + + def cleanup_executor(self) -> None: + """Close the executor if it was ever built. Never touches data sources.""" + if self.executor_built: + self.executor.cleanup() diff --git a/pkg-py/tests/test_table_set.py b/pkg-py/tests/test_table_set.py new file mode 100644 index 000000000..118ef7e9a --- /dev/null +++ b/pkg-py/tests/test_table_set.py @@ -0,0 +1,76 @@ +"""Tests for the TableSet value object.""" + +import duckdb +import pandas as pd +import pytest +from querychat._query_executor import DataSourceExecutor, DuckDBExecutor +from querychat._querychat_base import normalize_data_source +from querychat._system_prompt import QueryChatSystemPrompt +from querychat._table_set import TableSet + + +def make_table_set(**frames: pd.DataFrame) -> TableSet: + sources = {name: normalize_data_source(df, name) for name, df in frames.items()} + prompt = QueryChatSystemPrompt(prompt_template=None, data_sources=sources) + return TableSet(sources, prompt) + + +@pytest.fixture +def users(): + return pd.DataFrame({"id": [1, 2], "name": ["a", "b"]}) + + +@pytest.fixture +def orders(): + return pd.DataFrame({"id": [1], "user_id": [2]}) + + +def test_requires_at_least_one_source(): + prompt = QueryChatSystemPrompt(prompt_template=None, data_sources={}) + with pytest.raises(ValueError, match="at least one"): + TableSet({}, prompt) + + +def test_data_sources_is_read_only(users): + ts = make_table_set(users=users) + with pytest.raises(TypeError): + ts.data_sources["other"] = ts.data_sources["users"] # type: ignore[index] + + +def test_table_names_preserve_insertion_order(users, orders): + with pytest.warns(UserWarning, match="without a data_dict"): + ts = make_table_set(users=users, orders=orders) + assert ts.table_names == ["users", "orders"] + + +def test_executor_is_lazy_and_cached(users): + ts = make_table_set(users=users) + assert ts.executor_built is False + first = ts.executor + assert ts.executor_built is True + assert ts.executor is first + + +def test_single_table_uses_data_source_executor(users): + assert isinstance(make_table_set(users=users).executor, DataSourceExecutor) + + +def test_multi_dataframe_uses_duckdb_executor(users, orders): + with pytest.warns(UserWarning, match="without a data_dict"): + ts = make_table_set(users=users, orders=orders) + assert isinstance(ts.executor, DuckDBExecutor) + + +def test_cleanup_executor_is_noop_when_never_built(users): + ts = make_table_set(users=users) + ts.cleanup_executor() + assert ts.executor_built is False + + +def test_cleanup_executor_closes_built_duckdb_executor(users, orders): + with pytest.warns(UserWarning, match="without a data_dict"): + ts = make_table_set(users=users, orders=orders) + executor = ts.executor + ts.cleanup_executor() + with pytest.raises(duckdb.ConnectionException): + executor.execute_query("SELECT 1") From d9ccdfb828a231a2de4937226518eb2b1f7701c6 Mon Sep 17 00:00:00 2001 From: Carson Date: Sat, 12 Sep 2026 18:13:51 -0500 Subject: [PATCH 02/35] fix(py): SQLAlchemySource.cleanup() no longer disposes the caller's engine --- pkg-py/src/querychat/_datasource.py | 23 ++++++++--------------- pkg-py/tests/test_datasource.py | 16 ++++++++++++++++ 2 files changed, 24 insertions(+), 15 deletions(-) diff --git a/pkg-py/src/querychat/_datasource.py b/pkg-py/src/querychat/_datasource.py index cffcf1545..546bd425a 100644 --- a/pkg-py/src/querychat/_datasource.py +++ b/pkg-py/src/querychat/_datasource.py @@ -298,15 +298,12 @@ def get_data(self) -> IntoFrameT: @abstractmethod def cleanup(self) -> None: """ - Clean up resources associated with the data source. - - This method should clean up any connections or resources used by the - data source. - - Returns - ------- - None + Release resources this data source created. + Only resources querychat created are closed here (for example the + in-memory DuckDB connection a ``DataFrameSource`` opens). Connections, + engines, and backends passed in by the caller are never closed; their + lifecycle stays with the caller. """ def get_data_description(self) -> str: @@ -819,15 +816,11 @@ def _get_connection(self) -> Connection: def cleanup(self) -> None: """ - Dispose of the SQLAlchemy engine. - - Returns - ------- - None + No-op: the SQLAlchemy engine is owned by the caller. + Dispose it yourself with ``engine.dispose()`` when your application + shuts down. """ - if self._engine: - self._engine.dispose() class PolarsLazySource(DataSource["pl.LazyFrame"]): diff --git a/pkg-py/tests/test_datasource.py b/pkg-py/tests/test_datasource.py index 41256f99a..1f17389fb 100644 --- a/pkg-py/tests/test_datasource.py +++ b/pkg-py/tests/test_datasource.py @@ -568,3 +568,19 @@ def test_dataframe_source_get_schema_unchanged(sample_df) -> None: schema = source.get_schema(categorical_threshold=10) assert "Table: test" in schema assert "Columns:" in schema + + +def test_sqlalchemy_cleanup_does_not_dispose_caller_owned_engine( + test_db_engine, monkeypatch +): + """The engine belongs to the caller; querychat must not dispose it.""" + from unittest.mock import MagicMock + + dispose = MagicMock() + monkeypatch.setattr(test_db_engine, "dispose", dispose) + source = SQLAlchemySource(test_db_engine, "test_table") + + source.cleanup() + + dispose.assert_not_called() + assert len(source.get_data()) > 0 From 2240b6fd4e1b73bf5b2bff7677a8ded374ceff6f Mon Sep 17 00:00:00 2001 From: Carson Date: Sat, 12 Sep 2026 18:14:18 -0500 Subject: [PATCH 03/35] feat(r): add internal TableSet class and validate_source_group_compatibility() --- pkg-r/R/QueryExecutor.R | 10 +++ pkg-r/R/TableSet.R | 76 +++++++++++++++++++++ pkg-r/tests/testthat/helper-fixtures.R | 18 +++++ pkg-r/tests/testthat/test-TableSet.R | 93 ++++++++++++++++++++++++++ 4 files changed, 197 insertions(+) create mode 100644 pkg-r/R/TableSet.R create mode 100644 pkg-r/tests/testthat/test-TableSet.R diff --git a/pkg-r/R/QueryExecutor.R b/pkg-r/R/QueryExecutor.R index 1347e3d88..440e81a88 100644 --- a/pkg-r/R/QueryExecutor.R +++ b/pkg-r/R/QueryExecutor.R @@ -257,3 +257,13 @@ check_source_compatibility <- function(existing_sources, new_source, new_name) { invisible(NULL) } + +# Validates that every source in a group is compatible with the others. +validate_source_group_compatibility <- function(data_sources) { + existing <- list() + for (name in names(data_sources)) { + check_source_compatibility(existing, data_sources[[name]], name) + existing[[name]] <- data_sources[[name]] + } + invisible(NULL) +} diff --git a/pkg-r/R/TableSet.R b/pkg-r/R/TableSet.R new file mode 100644 index 000000000..50334ce49 --- /dev/null +++ b/pkg-r/R/TableSet.R @@ -0,0 +1,76 @@ +# The tables a chat can query, plus the system prompt and query executor +# built from them. Never mutated after construction: QueryChat holds one +# built from $add_table() calls, and $server(data_source = ) derives a second +# one per session so a running session never observes later changes. +# +# Not exported. +TableSet <- R6::R6Class( + "TableSet", + private = list( + .data_sources = NULL, + .system_prompt = NULL, + .data_description = NULL, + .executor = NULL + ), + public = list( + initialize = function( + data_sources, + system_prompt, + data_description = NULL + ) { + if (length(data_sources) == 0) { + cli::cli_abort("{.cls TableSet} requires at least one data source.") + } + nms <- names(data_sources) + if (is.null(nms) || any(!nzchar(nms))) { + cli::cli_abort("{.arg data_sources} must be a named list.") + } + private$.data_sources <- data_sources + private$.system_prompt <- system_prompt + private$.data_description <- data_description + }, + + executor = function() { + if (is.null(private$.executor)) { + private$.executor <- build_query_executor(private$.data_sources) + } + private$.executor + }, + + executor_built = function() { + !is.null(private$.executor) + }, + + table_names = function() { + names(private$.data_sources) + }, + + # Closes the executor if it was ever built. Never touches data sources. + cleanup_executor = function() { + if (!is.null(private$.executor)) { + private$.executor$cleanup() + } + invisible(NULL) + } + ), + active = list( + data_sources = function(value) { + if (!missing(value)) { + cli::cli_abort("{.field data_sources} is read-only.") + } + private$.data_sources + }, + system_prompt = function(value) { + if (!missing(value)) { + cli::cli_abort("{.field system_prompt} is read-only.") + } + private$.system_prompt + }, + data_description = function(value) { + if (!missing(value)) { + cli::cli_abort("{.field data_description} is read-only.") + } + private$.data_description + } + ) +) diff --git a/pkg-r/tests/testthat/helper-fixtures.R b/pkg-r/tests/testthat/helper-fixtures.R index 6ee167ee9..3fab7caa8 100644 --- a/pkg-r/tests/testthat/helper-fixtures.R +++ b/pkg-r/tests/testthat/helper-fixtures.R @@ -95,6 +95,24 @@ local_data_frame_source <- function( df_source } +# Build a TableSet for tests, deferring executor cleanup to `env`. +local_table_set <- function(data_sources, env = parent.frame()) { + sp <- suppressWarnings( + # The multi-table "consider a data_dict" advice is noise for these tests. + QueryChatSystemPrompt$new( + prompt_template = system.file( + "prompts", + "prompt.md", + package = "querychat" + ), + data_sources = data_sources + ) + ) + ts <- TableSet$new(data_sources, sp) + withr::defer(ts$cleanup_executor(), envir = env) + ts +} + local_recording_data_frame_source <- function( data = new_test_df(), table_name = "test_table", diff --git a/pkg-r/tests/testthat/test-TableSet.R b/pkg-r/tests/testthat/test-TableSet.R new file mode 100644 index 000000000..cbe25549a --- /dev/null +++ b/pkg-r/tests/testthat/test-TableSet.R @@ -0,0 +1,93 @@ +test_that("TableSet requires at least one data source", { + sp <- QueryChatSystemPrompt$new( + prompt_template = system.file( + "prompts", + "prompt.md", + package = "querychat" + ), + data_sources = list() + ) + expect_error(TableSet$new(list(), sp), "at least one") +}) + +test_that("TableSet requires a named list of data sources", { + skip_if_no_dataframe_engine() + ds <- local_data_frame_source(new_test_df()) + sp <- QueryChatSystemPrompt$new( + prompt_template = system.file( + "prompts", + "prompt.md", + package = "querychat" + ), + data_sources = list(test_table = ds) + ) + expect_error(TableSet$new(list(ds), sp), "named") +}) + +test_that("TableSet$table_names() preserves registration order", { + skip_if_not_installed("duckdb") + a <- local_data_frame_source(new_test_df(), "a") + b <- local_data_frame_source(new_users_df(), "b") + ts <- local_table_set(list(a = a, b = b)) + expect_equal(ts$table_names(), c("a", "b")) +}) + +test_that("TableSet builds its executor lazily and caches it", { + skip_if_no_dataframe_engine() + ds <- local_data_frame_source(new_test_df()) + ts <- local_table_set(list(test_table = ds)) + + expect_false(ts$executor_built()) + first <- ts$executor() + expect_true(ts$executor_built()) + expect_identical(ts$executor(), first) + expect_s3_class(first, "DataSourceExecutor") +}) + +test_that("TableSet uses DuckDBExecutor for several data frames", { + skip_if_not_installed("duckdb") + a <- local_data_frame_source(new_test_df(), "a") + b <- local_data_frame_source(new_users_df(), "b") + ts <- local_table_set(list(a = a, b = b)) + expect_s3_class(ts$executor(), "DuckDBExecutor") +}) + +test_that("TableSet$cleanup_executor() is a no-op before the executor is built", { + skip_if_no_dataframe_engine() + ds <- local_data_frame_source(new_test_df()) + ts <- local_table_set(list(test_table = ds)) + expect_no_error(ts$cleanup_executor()) + expect_false(ts$executor_built()) +}) + +test_that("TableSet$cleanup_executor() closes a built DuckDB executor", { + skip_if_not_installed("duckdb") + a <- local_data_frame_source(new_test_df(), "a") + b <- local_data_frame_source(new_users_df(), "b") + ts <- local_table_set(list(a = a, b = b)) + ex <- ts$executor() + ts$cleanup_executor() + expect_error(ex$execute_query("SELECT 1")) +}) + +test_that("TableSet fields are read-only", { + skip_if_no_dataframe_engine() + ds <- local_data_frame_source(new_test_df()) + ts <- local_table_set(list(test_table = ds)) + expect_error(ts$data_sources <- list(), "read-only") + expect_error(ts$system_prompt <- NULL, "read-only") + expect_error(ts$data_description <- "x", "read-only") +}) + +test_that("validate_source_group_compatibility() rejects mixed source types", { + skip_if_no_dataframe_engine() + skip_if_not_installed("RSQLite") + df_source <- local_data_frame_source(new_test_df(), "a") + db <- local_sqlite_connection(table_name = "b") + db_source <- DBISource$new(db$conn, "b") + expect_error( + validate_source_group_compatibility(list(a = df_source, b = db_source)), + "same type" + ) + expect_no_error(validate_source_group_compatibility(list(a = df_source))) +}) From 9f712a78864f184ca296d8cb27f0b95da4293e4d Mon Sep 17 00:00:00 2001 From: Carson Date: Sat, 12 Sep 2026 18:19:24 -0500 Subject: [PATCH 04/35] fix(r): DBISource/TblSqlSource cleanup() no longer disconnect the caller's connection --- pkg-r/R/DBISource.R | 8 +++----- pkg-r/R/DataSource.R | 5 ++++- pkg-r/R/TblSqlSource.R | 4 ++-- pkg-r/man/DBISource.Rd | 5 +++-- pkg-r/man/DataSource.Rd | 5 ++++- pkg-r/man/TblSqlSource.Rd | 4 ++-- pkg-r/tests/testthat/test-DBISource.R | 10 ++++++++++ pkg-r/tests/testthat/test-TblSqlSource.R | 9 +++++++++ 8 files changed, 37 insertions(+), 13 deletions(-) diff --git a/pkg-r/R/DBISource.R b/pkg-r/R/DBISource.R index 725c02d8c..fedbf3a15 100644 --- a/pkg-r/R/DBISource.R +++ b/pkg-r/R/DBISource.R @@ -213,13 +213,11 @@ DBISource <- R6::R6Class( }, #' @description - #' Disconnect from the database + #' No-op: the DBI connection is owned by the caller. Disconnect it + #' yourself with `DBI::dbDisconnect()` when your application shuts down. #' - #' @return NULL (invisibly) + #' @return `NULL` (invisibly) cleanup = function() { - if (!is.null(private$conn) && DBI::dbIsValid(private$conn)) { - DBI::dbDisconnect(private$conn) - } invisible(NULL) } ) diff --git a/pkg-r/R/DataSource.R b/pkg-r/R/DataSource.R index fe4c1539d..dc0594abc 100644 --- a/pkg-r/R/DataSource.R +++ b/pkg-r/R/DataSource.R @@ -112,7 +112,10 @@ DataSource <- R6::R6Class( }, #' @description - #' Clean up resources (close connections, etc.) + #' Release resources this data source created. Only resources querychat + #' opened itself are closed (for example the in-memory DuckDB connection a + #' [DataFrameSource] creates). Connections passed in by the caller are + #' never closed; their lifecycle stays with the caller. #' #' @return NULL (invisibly) cleanup = function() { diff --git a/pkg-r/R/TblSqlSource.R b/pkg-r/R/TblSqlSource.R index 7ab55e7ac..5c9af2124 100644 --- a/pkg-r/R/TblSqlSource.R +++ b/pkg-r/R/TblSqlSource.R @@ -188,9 +188,9 @@ TblSqlSource <- R6::R6Class( }, #' @description - #' Clean up resources (close connections, etc.) + #' No-op: the connection behind the `tbl_sql` is owned by the caller. #' - #' @return NULL (invisibly) + #' @return `NULL` (invisibly) cleanup = function() { super$cleanup() } diff --git a/pkg-r/man/DBISource.Rd b/pkg-r/man/DBISource.Rd index 539af25a8..9c1e622be 100644 --- a/pkg-r/man/DBISource.Rd +++ b/pkg-r/man/DBISource.Rd @@ -204,14 +204,15 @@ all original table columns (default: \code{FALSE})} \if{html}{\out{}} \if{latex}{\out{\hypertarget{method-DBISource-cleanup}{}}} \subsection{\code{DBISource$cleanup()}}{ - Disconnect from the database + No-op: the DBI connection is owned by the caller. Disconnect it +yourself with \code{DBI::dbDisconnect()} when your application shuts down. \subsection{Usage}{ \if{html}{\out{
}} \preformatted{DBISource$cleanup()} \if{html}{\out{
}} } \subsection{Returns}{ - NULL (invisibly) + \code{NULL} (invisibly) } } diff --git a/pkg-r/man/DataSource.Rd b/pkg-r/man/DataSource.Rd index 5ecd8ceee..dcabbfa4e 100644 --- a/pkg-r/man/DataSource.Rd +++ b/pkg-r/man/DataSource.Rd @@ -176,7 +176,10 @@ Subclasses may override this to provide metadata-derived descriptions \if{html}{\out{}} \if{latex}{\out{\hypertarget{method-DataSource-cleanup}{}}} \subsection{\code{DataSource$cleanup()}}{ - Clean up resources (close connections, etc.) + Release resources this data source created. Only resources querychat +opened itself are closed (for example the in-memory DuckDB connection a +\link{DataFrameSource} creates). Connections passed in by the caller are +never closed; their lifecycle stays with the caller. \subsection{Usage}{ \if{html}{\out{
}} \preformatted{DataSource$cleanup()} diff --git a/pkg-r/man/TblSqlSource.Rd b/pkg-r/man/TblSqlSource.Rd index f15a1cccc..fad878df3 100644 --- a/pkg-r/man/TblSqlSource.Rd +++ b/pkg-r/man/TblSqlSource.Rd @@ -221,14 +221,14 @@ all original table columns (default: \code{FALSE})} \if{html}{\out{}} \if{latex}{\out{\hypertarget{method-TblSqlSource-cleanup}{}}} \subsection{\code{TblSqlSource$cleanup()}}{ - Clean up resources (close connections, etc.) + No-op: the connection behind the \code{tbl_sql} is owned by the caller. \subsection{Usage}{ \if{html}{\out{
}} \preformatted{TblSqlSource$cleanup()} \if{html}{\out{
}} } \subsection{Returns}{ - NULL (invisibly) + \code{NULL} (invisibly) } } diff --git a/pkg-r/tests/testthat/test-DBISource.R b/pkg-r/tests/testthat/test-DBISource.R index 487052ff4..351a28ed0 100644 --- a/pkg-r/tests/testthat/test-DBISource.R +++ b/pkg-r/tests/testthat/test-DBISource.R @@ -102,3 +102,13 @@ describe("DBISource$test_query()", { expect_type(result$bool_col, "integer") }) }) + +test_that("DBISource$cleanup() leaves the caller's connection open", { + db <- local_sqlite_connection() + source <- DBISource$new(db$conn, "test_table") + + source$cleanup() + + expect_true(DBI::dbIsValid(db$conn)) + expect_equal(nrow(source$get_data()), 5L) +}) diff --git a/pkg-r/tests/testthat/test-TblSqlSource.R b/pkg-r/tests/testthat/test-TblSqlSource.R index d1abcedf2..e5666e33a 100644 --- a/pkg-r/tests/testthat/test-TblSqlSource.R +++ b/pkg-r/tests/testthat/test-TblSqlSource.R @@ -300,6 +300,15 @@ describe("TblSqlSource edge cases - Category B: Column Naming Issues", { }) }) +test_that("TblSqlSource$cleanup() leaves the caller's connection open", { + source <- local_tbl_sql_source(new_test_df()) + conn <- dbplyr::remote_con(source$get_data()) + + source$cleanup() + + expect_true(DBI::dbIsValid(conn)) +}) + describe("TblSqlSource edge cases - Category C: ORDER BY behavior", { it("handles ORDER BY without LIMIT", { source <- local_tbl_sql_source() From d9fa92096de268237cf5164275c5892ae91562dd Mon Sep 17 00:00:00 2001 From: Carson Date: Sat, 12 Sep 2026 18:22:27 -0500 Subject: [PATCH 05/35] fix: move TblSqlSource cleanup test to end of file per brief requirement --- pkg-r/tests/testthat/test-TblSqlSource.R | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/pkg-r/tests/testthat/test-TblSqlSource.R b/pkg-r/tests/testthat/test-TblSqlSource.R index e5666e33a..02b6e76b3 100644 --- a/pkg-r/tests/testthat/test-TblSqlSource.R +++ b/pkg-r/tests/testthat/test-TblSqlSource.R @@ -300,15 +300,6 @@ describe("TblSqlSource edge cases - Category B: Column Naming Issues", { }) }) -test_that("TblSqlSource$cleanup() leaves the caller's connection open", { - source <- local_tbl_sql_source(new_test_df()) - conn <- dbplyr::remote_con(source$get_data()) - - source$cleanup() - - expect_true(DBI::dbIsValid(conn)) -}) - describe("TblSqlSource edge cases - Category C: ORDER BY behavior", { it("handles ORDER BY without LIMIT", { source <- local_tbl_sql_source() @@ -395,3 +386,12 @@ describe("TblSqlSource edge cases - Category C: ORDER BY behavior", { expect_equal(collected$value[5], 10) }) }) + +test_that("TblSqlSource$cleanup() leaves the caller's connection open", { + source <- local_tbl_sql_source(new_test_df()) + conn <- dbplyr::remote_con(source$get_data()) + + source$cleanup() + + expect_true(DBI::dbIsValid(conn)) +}) From 5775b4b65517b778ae80637340252e344209714f Mon Sep 17 00:00:00 2001 From: Carson Date: Sat, 12 Sep 2026 18:35:29 -0500 Subject: [PATCH 06/35] refactor(py): give each Shiny session a read-only TableSet instead of mutating QueryChat server(data_source=) now derives a session-local TableSet and closes what it created on session end. Removes live-session counting, retired-resource deferral, and the greeting snapshot path. Late add_table() warns and parks the superseded set for cleanup(); replace/remove after sessions start raises. --- pkg-py/src/querychat/_querychat_base.py | 472 +++++++------------- pkg-py/src/querychat/_querychat_greeter.py | 57 +-- pkg-py/src/querychat/_shiny.py | 102 +++-- pkg-py/src/querychat/_shiny_module.py | 25 +- pkg-py/tests/test_base.py | 82 +++- pkg-py/tests/test_cleanup.py | 70 +-- pkg-py/tests/test_deferred_shiny.py | 11 +- pkg-py/tests/test_multi_table.py | 99 ++-- pkg-py/tests/test_multi_table_frameworks.py | 2 +- pkg-py/tests/test_pin_source.py | 10 +- pkg-py/tests/test_querychat.py | 54 +-- pkg-py/tests/test_server_data_source.py | 333 ++++++-------- pkg-py/tests/test_shiny_module.py | 29 +- pkg-py/tests/test_state.py | 2 +- 14 files changed, 573 insertions(+), 775 deletions(-) diff --git a/pkg-py/src/querychat/_querychat_base.py b/pkg-py/src/querychat/_querychat_base.py index cda17bc18..532ea6bf9 100644 --- a/pkg-py/src/querychat/_querychat_base.py +++ b/pkg-py/src/querychat/_querychat_base.py @@ -23,14 +23,7 @@ SQLAlchemySource, ) from ._pin_source import PinSource, is_pins_board -from ._query_executor import ( - DataSourceExecutor, - DuckDBExecutor, - PolarsSQLExecutor, - QueryExecutor, - check_source_compatibility, - validate_source_group_compatibility, -) +from ._query_executor import QueryExecutor, validate_source_group_compatibility from ._querychat_core import ( AppState, AppStateDict, @@ -39,6 +32,7 @@ ) from ._querychat_greeter import QueryChatGreeter from ._system_prompt import QueryChatSystemPrompt +from ._table_set import TableSet from ._utils import MISSING, MISSING_TYPE, is_ibis_backend, is_ibis_table from ._viz_utils import has_viz_deps, has_viz_tool from .tools import ( @@ -52,7 +46,7 @@ ) if TYPE_CHECKING: - from collections.abc import Callable + from collections.abc import Callable, Mapping from ibis.backends.sql import SQLBackend from narwhals.stable.v1.typing import IntoFrame @@ -64,6 +58,7 @@ TOOL_GROUPS = Literal["filter", "update", "query", "visualize"] DEFAULT_TOOLS: tuple[TOOL_GROUPS, ...] = ("filter", "query", "visualize") +TABLE_NAME_PATTERN = r"^[a-zA-Z][a-zA-Z0-9_]*$" class QueryChatBase(Generic[IntoFrameT]): @@ -95,22 +90,11 @@ def __init__( history: Optional[bool | HistoryOptions] = None, ): self._data_dicts: list[DataDict] = _normalize_data_dicts(data_dict) - - # Multi-table storage: dict of data sources keyed by table name - self._data_sources: dict[str, DataSource] = {} - self._query_executor: QueryExecutor | None = None - - # Live Shiny session count. Shared-state guards key off this: an - # ended session can no longer be using a resource it registered. - self._active_sessions = 0 - - # Sources/executors replaced while sessions were still live. Their - # cleanup is deferred until the last live session ends (or cleanup() - # runs) so a still-running session never loses a resource it uses. - self._retired_resources: list[DataSource | QueryExecutor] = [] - - # Name to register at .server(data_source=...) time when constructed - # with data_source=None (the deferred pattern). + self._table_set: TableSet | None = None + # Instance sets swapped out by add_table()/add_tables() after a session + # started. A running session may still hold one, so cleanup() closes them. + self._superseded_table_sets: list[TableSet] = [] + self._sessions_started = False self._deferred_table_name: str | None = None self.tools = normalize_tools(tools, default=DEFAULT_TOOLS) @@ -136,7 +120,6 @@ def __init__( self._base_client = client self._client_console = None - self._system_prompt: QueryChatSystemPrompt | None = None self._greeter: QueryChatGreeter | None = None if data_source is not None: @@ -151,32 +134,50 @@ def __init__( else: # Validate now: a bad deferred name would otherwise surface at # .server() registration time, after the module id is built. - if table_name is not None and not re.match( - r"^[a-zA-Z][a-zA-Z0-9_]*$", table_name - ): - raise ValueError( - "Table name must begin with a letter and contain only " - "letters, numbers, and underscores" - ) + if table_name is not None: + check_table_name(table_name) self._deferred_table_name = table_name - def _build_system_prompt( - self, - *, - data_sources: dict[str, DataSource] | None = None, - ) -> None: - """Build/rebuild the system prompt from current or staged data sources.""" - next_data_sources = self._data_sources if data_sources is None else data_sources + @property + def _data_sources(self) -> Mapping[str, DataSource]: + """Read-only view of the instance's registered tables.""" + if self._table_set is None: + return {} + return self._table_set.data_sources + + def _require_table_set(self, method_name: str) -> TableSet: + if self._table_set is None: + raise RuntimeError( + f"At least one data source must be set before calling {method_name}(). " + "Either pass data_source to __init__() or call add_table()." + ) + return self._table_set - if not next_data_sources: - raise RuntimeError("Cannot build system prompt without data_source") + def _require_initialized(self, method_name: str) -> None: + self._require_table_set(method_name) + + def _require_query_executor(self, method_name: str) -> QueryExecutor: + return self._require_table_set(method_name).executor - client_has_history = ( + def _build_table_set(self, sources: dict[str, DataSource]) -> TableSet: + validate_source_group_compatibility(sources) + prompt = QueryChatSystemPrompt( + prompt_template=self._prompt_template, + data_sources=sources, + data_description=self._data_description, + extra_instructions=self._extra_instructions, + categorical_threshold=self._categorical_threshold, + data_dicts=self._data_dicts, + ) + return TableSet(sources, prompt) + + def _warn_if_prompt_rebuilt_with_history(self) -> None: + has_history = ( self._base_client is not None and bool(self._base_client.get_turns()) ) or ( self._client_console is not None and bool(self._client_console.get_turns()) ) - if client_has_history: + if has_history: warnings.warn( "System prompt rebuilt after chat history exists. " "This invalidates any prompt caching from prior turns. " @@ -185,46 +186,37 @@ def _build_system_prompt( stacklevel=3, ) - self._system_prompt = QueryChatSystemPrompt( - prompt_template=self._prompt_template, - data_sources=next_data_sources, - data_description=self._data_description, - extra_instructions=self._extra_instructions, - categorical_threshold=self._categorical_threshold, - data_dicts=self._data_dicts, - ) - - def _build_query_executor( - self, *, data_sources: dict[str, DataSource] | None = None - ) -> QueryExecutor: - """Build a query executor from current or staged data sources.""" - sources = self._data_sources if data_sources is None else data_sources - - validate_source_group_compatibility(sources) - - if len(sources) == 1: - return DataSourceExecutor(dict(sources)) - - first_source = next(iter(sources.values())) - - if isinstance(first_source, DataFrameSource): - return DuckDBExecutor( - {n: s for n, s in sources.items() if isinstance(s, DataFrameSource)} - ) - if isinstance(first_source, PolarsLazySource): - return PolarsSQLExecutor( - {n: s for n, s in sources.items() if isinstance(s, PolarsLazySource)} - ) - - return DataSourceExecutor(dict(sources)) - - def _require_initialized(self, method_name: str) -> None: - """Raise if no data sources have been registered.""" - if not self._data_sources: + def _check_late_change(self, method_name: str, *, destructive: bool) -> None: + if not self._sessions_started: + return + if destructive: raise RuntimeError( - f"At least one data source must be set before calling {method_name}(). " - "Either pass data_source to __init__() or call add_table()." + f"Cannot call {method_name}() to replace or remove a table while " + "sessions may be using it. Configure all tables before calling " + ".server() or .app()." ) + warnings.warn( + f"{method_name}() called after a session has started. Sessions that " + "are already running keep the tables they started with; only new " + "sessions will see this change.", + UserWarning, + stacklevel=3, + ) + + def _swap_table_set(self, new_set: TableSet, *, replaced: list[DataSource]) -> None: + old_set, self._table_set = self._table_set, new_set + if old_set is None: + return + if self._sessions_started: + # A running session may still query through old_set's executor; + # cleanup() closes it. `replaced` is always empty here because + # _check_late_change() rejects replacement once sessions started. + self._superseded_table_sets.append(old_set) + return + with contextlib.suppress(Exception): + old_set.cleanup_executor() + for source in replaced: + source.cleanup() def _require_single_table(self, method_name: str) -> None: """Raise if multiple tables are registered, directing to per-table API.""" @@ -235,17 +227,6 @@ def _require_single_table(self, method_name: str) -> None: f"Use .table('name').{method_name}() for per-table access." ) - def _require_query_executor(self, method_name: str) -> QueryExecutor: - """Return the cached executor, building it lazily on first use.""" - if self._query_executor is None: - if not self._data_sources: - raise RuntimeError( - f"query executor must be set before calling {method_name}(). " - "Set the data_source first so querychat can build an executor." - ) - self._query_executor = self._build_query_executor() - return self._query_executor - def _create_client(self, base: chatlas.Chat | None = None) -> chatlas.Chat: """Clone a Chat from ``base`` or the resolved ``_base_client``.""" if base is None: @@ -279,6 +260,7 @@ def _close_owned_client(self, client: chatlas.Chat) -> None: def _create_session_client( self, + table_set: TableSet, *, base: chatlas.Chat | None = None, tools: TOOL_GROUPS | tuple[TOOL_GROUPS, ...] | MISSING_TYPE | None = MISSING, @@ -287,28 +269,27 @@ def _create_session_client( visualize: Callable[[VisualizeData], None] | None = None, handoff_available: bool = False, ) -> chatlas.Chat: - """Create a fresh, fully-configured Chat.""" + """Create a fresh Chat configured for ``table_set``.""" chat = self._create_client(base) resolved_tools = normalize_tools(tools, default=self.tools) - - if self._system_prompt is not None: - chat.system_prompt = self._system_prompt.render( - resolved_tools, - handoff_available=handoff_available, - ) + chat.system_prompt = table_set.system_prompt.render( + resolved_tools, + handoff_available=handoff_available, + ) if resolved_tools is None: return chat - executor = self._require_query_executor("_create_session_client") + executor = table_set.executor + table_names = table_set.table_names + multi_table = len(table_names) > 1 - # Always register the schema tool (for all non-None tool sets) chat.register_tool( tool_get_schema( self._data_dicts, executor, - list(self._data_sources.keys()), + table_names, self._categorical_threshold, ) ) @@ -316,30 +297,20 @@ def _create_session_client( if "update" in resolved_tools: update_fn = update_dashboard or (lambda _: None) user_reset = reset_dashboard or (lambda _table: None) - chat.register_tool( tool_update_dashboard( - executor, - list(self._data_sources.keys()), - update_fn, - multi_table=len(self._data_sources) > 1, + executor, table_names, update_fn, multi_table=multi_table ) ) - chat.register_tool( - tool_reset_dashboard(user_reset, list(self._data_sources.keys())) - ) + chat.register_tool(tool_reset_dashboard(user_reset, table_names)) if "query" in resolved_tools: - chat.register_tool( - tool_query(executor, multi_table=len(self._data_sources) > 1) - ) + chat.register_tool(tool_query(executor, multi_table=multi_table)) if "visualize" in resolved_tools: viz_fn = visualize or (lambda _: None) chat.register_tool( - tool_visualize( - executor, viz_fn, multi_table=len(self._data_sources) > 1 - ) + tool_visualize(executor, viz_fn, multi_table=multi_table) ) return chat @@ -374,8 +345,8 @@ def client( A configured chat client. """ - self._require_initialized("client") return self._create_session_client( + self._require_table_set("client"), tools=tools, update_dashboard=update_dashboard, reset_dashboard=reset_dashboard, @@ -399,13 +370,12 @@ def client_factory( prompt: str | Path, base: chatlas.Chat | None = None, *, - data_sources: dict[str, DataSource] | None = None, + table_set: TableSet | None = None, ) -> chatlas.Chat: + resolved = table_set if table_set is not None else self._table_set sp = QueryChatSystemPrompt( prompt_template=prompt, - data_sources=( - self._data_sources if data_sources is None else data_sources - ), + data_sources=dict(resolved.data_sources) if resolved else {}, data_description=self._data_description, extra_instructions=None, categorical_threshold=self._categorical_threshold, @@ -438,10 +408,7 @@ def console( @property def system_prompt(self) -> str: """Get the system prompt.""" - self._require_initialized("system_prompt") - if self._system_prompt is None: - raise RuntimeError("System prompt not initialized") - return self._system_prompt.render(self.tools) + return self._require_table_set("system_prompt").system_prompt.render(self.tools) @property def data_source(self) -> DataSource: @@ -502,91 +469,37 @@ def add_table( ValueError If table_name already exists (and replace=False) or is invalid. RuntimeError - If called while a server session is active. + If called to replace or remove an existing table after a session has started. - """ - if self._active_sessions > 0: - raise RuntimeError( - "Cannot add tables while a server session is active. " - "Add all tables before calling .server() or .app()." - ) - self._add_or_replace_table( - data_source, - table_name, - replace=replace, - include_in_greeting=include_in_greeting, - ) - - def _add_or_replace_table( - self, - data_source: IntoFrame | sqlalchemy.Engine | BaseBoard, - table_name: str, - *, - replace: bool, - include_in_greeting: bool, - cleanup_replaced: bool = True, - ) -> None: - """ - Stage a table and rebuild the system prompt/executor cache. - - Guard-free core of :meth:`add_table`, also called directly by - ``.server(data_source=...)`` so each session can register its own - table even while earlier sessions are still running. - - ``cleanup_replaced=False`` is for that per-session path: the - replaced table may still be in active use by an earlier, - still-running session, so cleaning it up here would pull the - resource out from under it. The retired source and cached executor - are retained and cleaned up once the last live session ends - (see ``_mark_server_initialized``). """ if not isinstance(include_in_greeting, bool): raise TypeError( "include_in_greeting must be True or False, got " f"{type(include_in_greeting).__name__}." ) + check_table_name(table_name, data_source=data_source) - if not is_pins_board(data_source) and not re.match( - r"^[a-zA-Z][a-zA-Z0-9_]*$", table_name - ): - raise ValueError( - "Table name must begin with a letter and contain only " - "letters, numbers, and underscores" - ) - - if table_name in self._data_sources and not replace: + exists = table_name in self._data_sources + if exists and not replace: raise ValueError(f"Table '{table_name}' already exists") + self._check_late_change("add_table", destructive=exists) normalized = normalize_data_source(data_source, table_name) try: - other_sources = { - name: source - for name, source in self._data_sources.items() - if name != table_name - } - check_source_compatibility(other_sources, normalized, table_name) - next_data_sources = dict(self._data_sources) - next_data_sources[table_name] = normalized - - self._build_system_prompt(data_sources=next_data_sources) + self._warn_if_prompt_rebuilt_with_history() + new_set = self._build_table_set({**self._data_sources, table_name: normalized}) except Exception: - cleanup_failed_staged_source(data_source, normalized) + if normalized is not data_source: + normalized.cleanup() raise 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: - old_source.cleanup() - else: - self._retired_resources.append(old_source) - if self._query_executor is not None: - if cleanup_replaced: - with contextlib.suppress(Exception): - self._query_executor.cleanup() - else: - self._retired_resources.append(self._query_executor) - self._query_executor = None + replaced = ( + [old_source] + if old_source is not None and old_source is not normalized + else [] + ) + self._swap_table_set(new_set, replaced=replaced) if include_in_greeting and table_name not in self.greeter.tables: self.greeter.tables = [*self.greeter.tables, table_name] @@ -631,7 +544,7 @@ def add_tables( # noqa: PLR0912 If the resolved table list is empty, any name is invalid, or any name already exists (and ``replace=False``). RuntimeError - If called while a server session is active. + If called to replace or remove an existing table after a session has started. Examples -------- @@ -651,12 +564,6 @@ def add_tables( # noqa: PLR0912 >>> qc.add_tables(backend) """ - if self._active_sessions > 0: - raise RuntimeError( - "Cannot add tables while a server session is active. " - "Add all tables before calling .server() or .app()." - ) - if isinstance(data_source, sqlalchemy.Engine): if tables is None: tables = sqlalchemy.inspect(data_source).get_table_names() @@ -680,13 +587,11 @@ def normalized_builder(name: str) -> DataSource: raise ValueError("No tables found in database") for table_name in tables: - if not re.match(r"^[a-zA-Z][a-zA-Z0-9_]*$", table_name): - raise ValueError( - "Table name must begin with a letter and contain only " - "letters, numbers, and underscores" - ) - if table_name in self._data_sources and not replace: - raise ValueError(f"Table '{table_name}' already exists") + check_table_name(table_name) + existing = [name for name in tables if name in self._data_sources] + if existing and not replace: + raise ValueError(f"Table '{existing[0]}' already exists") + self._check_late_change("add_tables", destructive=bool(existing)) if isinstance(include_in_greeting, bool): greeting_names = list(tables) if include_in_greeting else [] @@ -701,26 +606,16 @@ def normalized_builder(name: str) -> DataSource: ) normalized = {name: normalized_builder(name) for name in tables} - - staged: dict[str, DataSource] = {} - for name, source in normalized.items(): - other_sources = {n: s for n, s in self._data_sources.items() if n != name} - check_source_compatibility({**other_sources, **staged}, source, name) - staged[name] = source - - next_data_sources = {**self._data_sources, **normalized} - self._build_system_prompt(data_sources=next_data_sources) - - for name, normalized_source in normalized.items(): - old_source = self._data_sources.get(name) - if old_source is not None and old_source is not normalized_source: - old_source.cleanup() - - self._data_sources = next_data_sources - if self._query_executor is not None: - with contextlib.suppress(Exception): - self._query_executor.cleanup() - self._query_executor = None + self._warn_if_prompt_rebuilt_with_history() + new_set = self._build_table_set({**self._data_sources, **normalized}) + + replaced = [ + old + for name in tables + if (old := self._data_sources.get(name)) is not None + and old is not normalized[name] + ] + self._swap_table_set(new_set, replaced=replaced) new_greeting = list(self.greeter.tables) for name in greeting_names: @@ -742,14 +637,10 @@ def remove_table(self, table_name: str) -> None: ValueError If table doesn't exist or is the last remaining table. RuntimeError - If called while a server session is active. + If called to replace or remove an existing table after a session has started. """ - if self._active_sessions > 0: - raise RuntimeError( - "Cannot remove tables while a server session is active. " - "Configure all tables before calling .server() or .app()." - ) + self._check_late_change("remove_table", destructive=True) if table_name not in self._data_sources: available = ", ".join(self._data_sources.keys()) @@ -760,78 +651,38 @@ def remove_table(self, table_name: str) -> None: "Cannot remove last table. At least one table is required." ) - removed_source = self._data_sources[table_name] - next_data_sources = dict(self._data_sources) - del next_data_sources[table_name] + removed = self._data_sources[table_name] + remaining = {n: s for n, s in self._data_sources.items() if n != table_name} + self._warn_if_prompt_rebuilt_with_history() + new_set = self._build_table_set(remaining) + self._swap_table_set(new_set, replaced=[removed]) - self._build_system_prompt(data_sources=next_data_sources) - self._data_sources = next_data_sources if self._greeter is not None: self._greeter.tables = [n for n in self._greeter.tables if n != table_name] - if self._query_executor is not None: - with contextlib.suppress(Exception): - self._query_executor.cleanup() - self._query_executor = None - removed_source.cleanup() - - def _mark_server_initialized(self, session) -> None: - """ - Track a newly started session until it ends. - - The add/remove_table guards and cleanup-on-replace in - ``server(data_source=...)`` key off the number of *live* sessions: - a session that has ended can no longer be using a replaced resource. - """ - self._active_sessions += 1 - - def untrack_session() -> None: - self._active_sessions -= 1 - if self._active_sessions == 0: - self._flush_retired_resources() - - session.on_ended(untrack_session) - - def _flush_retired_resources(self) -> None: - """ - Clean up resources retired while sessions were still live. - - Retired sources/executors may still be in use by a live session, so - this only runs once no sessions remain (or from ``cleanup()``). - """ - retired = self._retired_resources - self._retired_resources = [] - for resource in retired: - # Best-effort: one failing cleanup must not leave the rest open. - with contextlib.suppress(Exception): - resource.cleanup() def cleanup(self) -> None: """ - Clean up resources held by this object. + Clean up resources this object created. - This closes the query executor and all data sources (e.g., DuckDB - connections). It also closes the chatlas client, but only if - querychat created it (i.e., `client` was `None` or a string spec like - `"openai/gpt-4o"`, including per-call overrides such as - `.server(client="openai")`). A user-supplied `chatlas.Chat` instance - is never closed here -- its lifecycle remains the caller's - responsibility. + Closes the query executors and data-source connections querychat + opened (in-memory DuckDB), including those of table sets superseded + by a late ``add_table()``. Connections, engines, and backends you + passed in are never closed. Also closes the chatlas client, but only + if querychat created it from a spec (``client=None`` or a string such + as ``"openai/gpt-4o"``); a ``chatlas.Chat`` you supplied is left open. Safe to call multiple times. In long-lived applications, call this - when the app shuts down (e.g., via `atexit`). + when the app shuts down (e.g., via ``atexit``). """ - if self._query_executor is not None: - self._query_executor.cleanup() - for source in self._data_sources.values(): - source.cleanup() - self._flush_retired_resources() + for superseded in self._superseded_table_sets: + warn_on_failure(superseded.cleanup_executor, "query executor") + self._superseded_table_sets.clear() + if self._table_set is not None: + warn_on_failure(self._table_set.cleanup_executor, "query executor") + for source in self._table_set.data_sources.values(): + warn_on_failure(source.cleanup, "data source") for client in self._owned_clients: - # Best-effort: one provider's close() failing must not leave the - # remaining owned clients open. - try: - client.close() - except Exception as e: # noqa: PERF203 (teardown of a few clients, not a hot loop) - warnings.warn(f"Failed to close chatlas client: {e}", stacklevel=2) + warn_on_failure(client.close, "chatlas client") self._owned_clients.clear() @@ -879,22 +730,23 @@ def normalize_data_source( ) -def cleanup_failed_staged_source( - original_source: IntoFrame | sqlalchemy.Engine | BaseBoard | DataSource, - normalized_source: DataSource, -) -> None: - """ - Clean up transient resources created during a failed staged rebuild. - - DataFrameSource and PinSource both allocate disposable connections during - normalization. SQLAlchemySource wraps a caller-owned engine, while - PolarsLazySource and IbisSource do not allocate disposable resources here. - """ - if isinstance(original_source, (DataSource, sqlalchemy.Engine)): +def check_table_name(table_name: str, *, data_source: object = None) -> None: + """Reject SQL table names querychat can't safely interpolate. Pins are exempt.""" + if data_source is not None and is_pins_board(data_source): return + if not re.match(TABLE_NAME_PATTERN, table_name): + raise ValueError( + "Table name must begin with a letter and contain only " + "letters, numbers, and underscores" + ) + - if isinstance(normalized_source, (DataFrameSource, PinSource)): - normalized_source.cleanup() +def warn_on_failure(fn: Callable[[], None], what: str) -> None: + """Run a teardown step, warning instead of raising so the rest still runs.""" + try: + fn() + except Exception as e: + warnings.warn(f"Failed to clean up {what}: {e}", stacklevel=3) def resolve_client(client: str | chatlas.Chat | None) -> chatlas.Chat: diff --git a/pkg-py/src/querychat/_querychat_greeter.py b/pkg-py/src/querychat/_querychat_greeter.py index f5d4422db..bd93c347d 100644 --- a/pkg-py/src/querychat/_querychat_greeter.py +++ b/pkg-py/src/querychat/_querychat_greeter.py @@ -12,13 +12,14 @@ import chatlas - from ._datasource import DataSource + from ._table_set import TableSet class QueryChatGreeter: """Controls greeting generation for a QueryChat instance. Access via ``qc.greeter``.""" def __init__(self, client_factory: Callable[..., chatlas.Chat]) -> None: + """``client_factory`` signature: ``(tables, prompt, base, *, table_set)``.""" self._client_factory = client_factory self._tables: list[str] = [] self._prompt: str | Path = Path(__file__).parent / "prompts" / "greeting.md" @@ -53,9 +54,27 @@ def prompt(self) -> str | Path: def prompt(self, value: str | Path) -> None: self._prompt = value - def build_client(self, base: chatlas.Chat | None = None) -> chatlas.Chat: - """Build a greeting chat client using the injected factory.""" - return self._client_factory(self._tables, self._prompt, base) + def build_client( + self, + base: chatlas.Chat | None = None, + *, + tables: list[str] | None = None, + table_set: TableSet | None = None, + ) -> chatlas.Chat: + """ + Build a greeting chat client using the injected factory. + + ``tables`` and ``table_set`` default to this greeter's configured + tables and the owning QueryChat's instance table set. ``server()`` + passes explicit values so a lazily generated greeting describes the + session that requested it, not whatever the instance holds later. + """ + return self._client_factory( + self._tables if tables is None else tables, + self._prompt, + base, + table_set=table_set, + ) def generate( self, @@ -66,31 +85,13 @@ def generate( """Generate a greeting using the greeting system prompt.""" return str(self.build_client(base).chat(GREETING_PROMPT, echo=echo)) - async def generate_async(self, *, base: chatlas.Chat | None = None): - """Stream a greeting response from the greeting client.""" - client = self.build_client(base) - return await client.stream_async(GREETING_PROMPT, echo="none") - - async def _generate_async_snapshot( + async def generate_async( self, *, - base: chatlas.Chat | None, - tables: list[str] | None, - data_sources: dict[str, DataSource], + base: chatlas.Chat | None = None, + tables: list[str] | None = None, + table_set: TableSet | None = None, ): - """ - Stream a greeting response from an explicit session snapshot. - - Internal counterpart to :meth:`generate_async`, used by - ``mod_server()``. The snapshot matters because greeting generation - is scheduled lazily: by the time it runs, a later session's - ``.server(data_source=...)`` call may have already mutated the - shared live state. - """ - client = self._client_factory( - self._tables if tables is None else tables, - self._prompt, - base, - data_sources=data_sources, - ) + """Stream a greeting response from the greeting client.""" + client = self.build_client(base, tables=tables, table_set=table_set) return await client.stream_async(GREETING_PROMPT, echo="none") diff --git a/pkg-py/src/querychat/_shiny.py b/pkg-py/src/querychat/_shiny.py index ebdf2c2c8..89f6e0d02 100644 --- a/pkg-py/src/querychat/_shiny.py +++ b/pkg-py/src/querychat/_shiny.py @@ -14,7 +14,14 @@ from shiny import App, Inputs, Outputs, Session, reactive, render, req, ui from ._icons import bs_icon -from ._querychat_base import DEFAULT_TOOLS, TOOL_GROUPS, QueryChatBase +from ._querychat_base import ( + DEFAULT_TOOLS, + TOOL_GROUPS, + QueryChatBase, + check_table_name, + normalize_data_source, + warn_on_failure, +) from ._shiny_module import ( CHAT_ID, ServerValues, @@ -38,7 +45,9 @@ from narwhals.stable.v1.typing import IntoFrame from ._data_dict import DataDict + from ._datasource import DataSource from ._table_accessor import TableAccessor + from ._table_set import TableSet class QueryChat(QueryChatBase[IntoFrameT]): @@ -312,7 +321,7 @@ def app( A Shiny App object that can be run with `app.run()` or served with `shiny run`. """ - self._require_initialized("app") + self._require_table_set("app") resolved_history: bool | HistoryOptions = ( history if history is not None @@ -407,15 +416,15 @@ def app_ui(request): ) def app_server(input: Inputs, output: Outputs, session: Session): - self._mark_server_initialized(session) if enable_bookmarking: session.bookmark.exclude.extend(["reset_query", "sql_editor"]) + self._sessions_started = True + table_set = self._require_table_set("app") vals = mod_server( self.id, - data_sources=dict(self._data_sources), - executor=self._require_query_executor("server"), + table_set=table_set, greeting=self.greeting, - client=self._create_session_client, + client=lambda **kw: self._create_session_client(table_set, **kw), history=resolved_history, tools=self.tools, greeter=self.greeter, @@ -631,7 +640,7 @@ def page(self, title, *, id: Optional[str] = None, **kwargs): **kwargs, ) - def server( + def server( # noqa: PLR0912 self, *, data_source: IntoFrame | sqlalchemy.Engine | ibis.Table | None = None, @@ -658,7 +667,10 @@ def server( 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. + first already-registered table. The table is registered for this + session only: the instance's own tables are not modified, a + same-named instance table is shadowed for this session, and the + session's data source is cleaned up when the session ends. table_name Table name to register `data_source` under. Only used when `data_source` is provided. @@ -698,6 +710,11 @@ def server( ".server() must be called within an active Shiny session (i.e., within the server function). " ) + self._sessions_started = True + table_set: TableSet | None = self._table_set + greeting_tables = list(self.greeter.tables) + session_source: DataSource | None = None + if data_source is not None: if table_name is not None: resolved_table_name = table_name @@ -712,29 +729,41 @@ def server( "or table_name to the QueryChat constructor, or register a " "table first with add_table()." ) - self._add_or_replace_table( - data_source, - resolved_table_name, - replace=True, - include_in_greeting=True, - # A live session may still be using the replaced source, - # so defer its cleanup until no sessions are active. - cleanup_replaced=self._active_sessions == 0, - ) + check_table_name(resolved_table_name, data_source=data_source) + session_source = normalize_data_source(data_source, resolved_table_name) + try: + table_set = self._build_table_set( + {**self._data_sources, resolved_table_name: session_source} + ) + except Exception: + if session_source is not data_source: + session_source.cleanup() + raise + if resolved_table_name not in greeting_tables: + greeting_tables.append(resolved_table_name) + + if table_set is None: + table_set = self._require_table_set("server") - self._require_initialized("server") resolved_client: chatlas.Chat | None = ( None if isinstance(client, MISSING_TYPE) else self._resolve_override_client(client) ) if resolved_client is not None and not isinstance(client, chatlas.Chat): - # Owned overrides are session-scoped: close and untrack them when - # this session ends rather than holding them open until cleanup(). session.on_ended(lambda: self._close_owned_client(resolved_client)) + if session_source is not None: + session_set, owned_source = table_set, session_source + + def cleanup_session() -> None: + warn_on_failure(session_set.cleanup_executor, "session query executor") + warn_on_failure(owned_source.cleanup, "session data source") + + session.on_ended(cleanup_session) + def create_session_client(**kwargs) -> chatlas.Chat: - return self._create_session_client(base=resolved_client, **kwargs) + return self._create_session_client(table_set, base=resolved_client, **kwargs) if enable_bookmarking is not None: warnings.warn( @@ -760,18 +789,16 @@ def create_session_client(**kwargs) -> chatlas.Chat: ) ) - self._mark_server_initialized(session) return mod_server( id or self.id, - data_sources=dict(self._data_sources), - executor=self._require_query_executor("server"), + table_set=table_set, greeting=self.greeting, client=create_session_client, history=resolved_history, tools=self.tools, greeter=self.greeter, greeting_base=resolved_client, - greeting_tables=list(self.greeter.tables), + greeting_tables=greeting_tables, ) @@ -1020,26 +1047,30 @@ def __init__( ) self._enable_bookmarking = enable_bookmarking + self._server_attempted = False self._vals: ServerValues[IntoFrameT] | None = None def _ensure_server_started(self) -> None: """ Start the Shiny module server if not already started. - Called lazily from ui()/sidebar() and the reactive accessors so that - module-level add_table() calls (which happen after __init__ but before - sidebar()/ui()) can complete before server initialization locks the - table set. + Called lazily from ui()/sidebar()/page() and the reactive accessors so + module-level add_table() calls, which run after __init__ but before + the UI is built, are included. Express re-executes the app file per + session, so this instance only ever sees one real session; a single + flag guarantees mod_server() runs at most once even if the first + attempt raised. """ - if self._active_sessions > 0: + if self._server_attempted: return session = get_current_session() if session is None or isinstance(session, ExpressStubSession): return - if not self._data_sources: + if self._table_set is None: return - self._require_initialized("_ensure_server_started") - self._mark_server_initialized(session) + self._server_attempted = True + self._sessions_started = True + table_set = self._table_set resolved_history: bool | HistoryOptions = ( self.history if self.history is not None @@ -1051,10 +1082,9 @@ def _ensure_server_started(self) -> None: ) self._vals = mod_server( self.id, - data_sources=dict(self._data_sources), - executor=self._require_query_executor("_ensure_server_started"), + table_set=table_set, greeting=self.greeting, - client=self._create_session_client, + client=lambda **kw: self._create_session_client(table_set, **kw), history=resolved_history, tools=self.tools, greeter=self.greeter, diff --git a/pkg-py/src/querychat/_shiny_module.py b/pkg-py/src/querychat/_shiny_module.py index b238b4191..74db9d6be 100644 --- a/pkg-py/src/querychat/_shiny_module.py +++ b/pkg-py/src/querychat/_shiny_module.py @@ -30,6 +30,7 @@ from ._datasource import DataSource from ._query_executor import QueryExecutor from ._querychat_greeter import QueryChatGreeter + from ._table_set import TableSet from ._viz_tools import VisualizeData from .types import UpdateDashboardData @@ -240,8 +241,7 @@ def mod_server( output: Outputs, session: Session, *, - data_sources: dict[str, DataSource[IntoFrameT]] | None, - executor: QueryExecutor | None, + table_set: TableSet | None, greeting: str | None, client: Callable[..., chatlas.Chat], history: bool | HistoryOptions, @@ -249,7 +249,10 @@ def mod_server( greeter: QueryChatGreeter, greeting_base: chatlas.Chat | None = None, greeting_tables: list[str] | None = None, -) -> ServerValues[IntoFrameT]: + # TableSet erases each DataSource's frame type for uniform executor handling, + # so IntoFrameT can't be bound from a parameter here -- only from callers' + # annotated return-type usage (e.g. QueryChat.server() -> ServerValues[IntoFrameT]). +) -> ServerValues[IntoFrameT]: # pyright: ignore[reportInvalidTypeVarUse] if not callable(client): raise TypeError("mod_server() requires a callable client factory.") @@ -299,14 +302,14 @@ def build_chat_client() -> chatlas.Chat: ) # Short-circuit for stub sessions (e.g. 1st run of an Express app) - # data_sources may be None during stub session for deferred pattern + # table_set may be None during stub session for deferred pattern if session.is_stub_session(): # Mock the error that would otherwise occur in a real session def _stub_df(): raise RuntimeError("RuntimeError: No current reactive context") stub_client = ( - _DeferredStubChatClient() if data_sources is None else build_chat_client() + _DeferredStubChatClient() if table_set is None else build_chat_client() ) return ServerValues( @@ -315,17 +318,19 @@ def _stub_df(): title=ReactiveStringOrNone(None), tables={}, client=stub_client, - data_sources=data_sources or {}, + data_sources=dict(table_set.data_sources) if table_set else {}, current_table=ReactiveStringOrNone(None), ) - # Real session requires data_sources and executor - if data_sources is None or executor is None: + if table_set is None: raise RuntimeError( "At least one table must be registered before the session starts. " "Call add_table() before server(), or pass the data to the QueryChat constructor." ) + data_sources = dict(table_set.data_sources) + executor = table_set.executor + for name, source in data_sources.items(): table_states[name] = _make_table_state(source, executor) @@ -345,10 +350,10 @@ async def _make_greeting(): GreetWarning, stacklevel=1, ) - stream = await greeter._generate_async_snapshot( + stream = await greeter.generate_async( base=greeting_base, tables=greeting_tables, - data_sources=data_sources, + table_set=table_set, ) return shinychat.chat_greeting(stream, persistent=True) diff --git a/pkg-py/tests/test_base.py b/pkg-py/tests/test_base.py index 2e2d425e7..59b0f6372 100644 --- a/pkg-py/tests/test_base.py +++ b/pkg-py/tests/test_base.py @@ -4,6 +4,7 @@ import warnings from pathlib import Path from typing import Any +from unittest.mock import patch import chatlas import narwhals.stable.v1 as nw @@ -265,6 +266,7 @@ def test_session_client_advertises_handoff_without_registering_tool( qc = QueryChatBase(sample_df, "test_table") client = qc._create_session_client( + qc._table_set, tools=None, handoff_available=True, ) @@ -370,12 +372,6 @@ def test_empty_list_raises(self, multi_table_engine): with pytest.raises(ValueError, match="No tables found"): qc.add_tables(multi_table_engine, []) - def test_after_server_raises(self, multi_table_engine): - qc = QueryChatBase() - qc._active_sessions = 1 - with pytest.raises(RuntimeError, match="Cannot add tables while a server session"): - qc.add_tables(multi_table_engine) - def test_system_prompt_built_exactly_once(self, multi_table_engine): qc = QueryChatBase() with warnings.catch_warnings(record=True) as w: @@ -475,3 +471,77 @@ def test_history_stored_verbatim_no_default_substitution(): qc_explicit_true = QueryChatBase(df, "a_table3", history=True) assert qc_explicit_true.history is True + + +class TestLateConfigurationChanges: + def test_add_new_table_after_sessions_started_warns(self, sample_df): + qc = QueryChatBase(sample_df, "users") + qc._sessions_started = True + + with pytest.warns(UserWarning, match="after a session has started"): + qc.add_table(sample_df, "other") + + assert qc.table_names() == ["users", "other"] + + def test_add_new_table_after_sessions_started_parks_old_set(self, sample_df): + qc = QueryChatBase(sample_df, "users") + old_set = qc._table_set + old_set.executor # noqa: B018 (forces the cached executor to build) + qc._sessions_started = True + + with ( # noqa: PT031 (asserts before/after cleanup() inside the same warns block) + pytest.warns(UserWarning, match="after a session has started"), + patch.object(old_set, "cleanup_executor") as cleanup_executor, + ): + qc.add_table(sample_df, "other") + cleanup_executor.assert_not_called() + assert qc._superseded_table_sets == [old_set] + + qc.cleanup() + cleanup_executor.assert_called_once() + assert qc._superseded_table_sets == [] + + def test_replace_after_sessions_started_raises(self, sample_df): + qc = QueryChatBase(sample_df, "users") + qc._sessions_started = True + + with pytest.raises(RuntimeError, match="replace or remove"): + qc.add_table(sample_df, "users", replace=True) + assert qc._superseded_table_sets == [] + + def test_remove_after_sessions_started_raises(self, sample_df): + qc = QueryChatBase(sample_df, "users") + qc.add_table(sample_df, "other") + qc._sessions_started = True + + with pytest.raises(RuntimeError, match="replace or remove"): + qc.remove_table("other") + assert qc.table_names() == ["users", "other"] + + def test_add_tables_replace_after_sessions_started_raises( + self, multi_table_engine + ): + qc = QueryChatBase() + qc.add_tables(multi_table_engine) + qc._sessions_started = True + + with pytest.raises(RuntimeError, match="replace or remove"): + qc.add_tables(multi_table_engine, replace=True) + + def test_before_sessions_start_replace_closes_old_source_immediately( + self, sample_df + ): + qc = QueryChatBase(sample_df, "users") + old_source = qc._data_sources["users"] + + with patch.object(old_source, "cleanup") as cleanup: + qc.add_table(sample_df, "users", replace=True) + cleanup.assert_called_once() + + def test_data_sources_accessor_is_read_only(self, sample_df): + qc = QueryChatBase(sample_df, "users") + with pytest.raises(TypeError): + qc._data_sources["x"] = object() # type: ignore[index] + + def test_data_sources_accessor_is_empty_when_deferred(self): + assert dict(QueryChatBase(None, "users")._data_sources) == {} diff --git a/pkg-py/tests/test_cleanup.py b/pkg-py/tests/test_cleanup.py index cbb2fb894..19c3fe10e 100644 --- a/pkg-py/tests/test_cleanup.py +++ b/pkg-py/tests/test_cleanup.py @@ -185,74 +185,6 @@ def test_owned_override_tracked_and_closed_by_cleanup( assert override.provider._client.is_closed() -class TestRetiredResourceCleanup: - """Resources replaced while sessions are live are cleaned once they end.""" - - @pytest.fixture - def ended_callbacks(self, monkeypatch): - callbacks = [] - fake_session = MagicMock() - fake_session.on_ended = callbacks.append - monkeypatch.setattr(shiny_mod, "get_current_session", lambda: fake_session) - monkeypatch.setattr(shiny_mod, "mod_server", lambda *args, **kwargs: None) - return callbacks - - def test_replacement_with_live_session_defers_cleanup( - self, sample_df, ended_callbacks - ): - qc = shiny_mod.QueryChat(sample_df, "users") - qc.server() - old_source = qc._data_sources["users"] - old_executor = qc._query_executor - - # Second session replaces the table while the first is still live - replacement = sample_df.copy() - qc.server(data_source=replacement) - - assert old_source in qc._retired_resources - assert old_executor in qc._retired_resources - with ( - patch.object(old_source, "cleanup") as source_cleanup, - patch.object(old_executor, "cleanup") as executor_cleanup, - ): - for cb in ended_callbacks: - cb() - source_cleanup.assert_called_once() - executor_cleanup.assert_called_once() - assert qc._retired_resources == [] - - def test_replacement_without_live_session_cleans_immediately( - self, sample_df, ended_callbacks - ): - qc = shiny_mod.QueryChat(sample_df, "users") - qc.server() - for cb in ended_callbacks: - cb() - old_source = qc._data_sources["users"] - - with patch.object(old_source, "cleanup") as source_cleanup: - qc.server(data_source=sample_df.copy()) - source_cleanup.assert_called_once() - assert qc._retired_resources == [] - - def test_cleanup_cleans_retired_resources(self, sample_df, ended_callbacks): - qc = shiny_mod.QueryChat(sample_df, "users") - qc.server() - old_source = qc._data_sources["users"] - old_executor = qc._query_executor - qc.server(data_source=sample_df.copy()) - - # Sessions never end: cleanup() still releases retired resources - with ( - patch.object(old_source, "cleanup") as source_cleanup, - patch.object(old_executor, "cleanup") as executor_cleanup, - ): - qc.cleanup() - source_cleanup.assert_called_once() - executor_cleanup.assert_called_once() - assert qc._retired_resources == [] - - class TestCleanupDataSources: """Existing executor/source cleanup behavior is preserved.""" @@ -278,7 +210,7 @@ def test_cleanup_closes_remaining_clients_after_close_failure( override = qc._resolve_override_client("openai") with ( patch.object(qc._base_client, "close", side_effect=RuntimeError("boom")), - pytest.warns(UserWarning, match="Failed to close chatlas client"), + pytest.warns(UserWarning, match="Failed to clean up chatlas client"), ): qc.cleanup() assert override.provider._client.is_closed() diff --git a/pkg-py/tests/test_deferred_shiny.py b/pkg-py/tests/test_deferred_shiny.py index 28228861f..7e2f81c57 100644 --- a/pkg-py/tests/test_deferred_shiny.py +++ b/pkg-py/tests/test_deferred_shiny.py @@ -125,9 +125,6 @@ def test_multiple_server_overrides_do_not_leak_into_shared_state(self, sample_df with session_context(ExpressStubSession()): qc.server(client=first_override) - # Reset live-session count for sequential test - qc._active_sessions = 0 - with session_context(ExpressStubSession()): qc.server(client=second_override) @@ -161,7 +158,7 @@ def test_server_not_initialized_after_init_stub_session(self, orders_df): """No session may be tracked after __init__ in a stub session.""" with session_context(ExpressStubSession()): qc = ExpressQueryChat(orders_df, "orders") - assert qc._active_sessions == 0 + assert qc._server_attempted is False def test_ensure_server_started_noop_during_stub_session( self, orders_df, monkeypatch @@ -196,11 +193,11 @@ def fake_mod_server(*args, **kwargs): mock_session.ns = Root with session_context(mock_session): qc = ExpressQueryChat(orders_df, "orders") - assert qc._active_sessions == 0 + assert qc._server_attempted is False qc._ensure_server_started() assert len(called) == 1 - assert qc._active_sessions == 1 + assert qc._server_attempted is True def test_ensure_server_started_idempotent(self, orders_df, monkeypatch): """_ensure_server_started() called twice starts server only once.""" @@ -228,7 +225,7 @@ def test_add_table_after_init_then_server_started_in_real_session( started_with_sources: list[list[str]] = [] def fake_mod_server(*args, **kwargs): - started_with_sources.append(list(kwargs["data_sources"].keys())) + started_with_sources.append(list(kwargs["table_set"].data_sources.keys())) return MagicMock() monkeypatch.setattr("querychat._shiny.mod_server", fake_mod_server) diff --git a/pkg-py/tests/test_multi_table.py b/pkg-py/tests/test_multi_table.py index a15a068e3..96f24bbf2 100644 --- a/pkg-py/tests/test_multi_table.py +++ b/pkg-py/tests/test_multi_table.py @@ -2,6 +2,7 @@ import os import tempfile +from collections.abc import Mapping from pathlib import Path from unittest.mock import patch @@ -13,6 +14,7 @@ from querychat._query_executor import ( DataSourceExecutor, DuckDBExecutor, + build_query_executor, check_source_compatibility, ) from querychat._querychat_base import QueryChatBase, normalize_data_source @@ -169,12 +171,12 @@ class TestMultiSourceStorage: """Tests for multi-source storage infrastructure.""" def test_single_table_stored_in_data_sources(self, orders_df): - """Test that single table is stored in _data_sources dict.""" + """Test that single table is stored in the _data_sources mapping.""" qc = QueryChat(orders_df, "orders", greeting="Hello!") - # Should have _data_sources dict with one entry + # Should have a read-only _data_sources mapping with one entry assert hasattr(qc, "_data_sources") - assert isinstance(qc._data_sources, dict) + assert isinstance(qc._data_sources, Mapping) assert "orders" in qc._data_sources assert len(qc._data_sources) == 1 @@ -215,10 +217,10 @@ def test_add_table_invalid_name_raises(self, orders_df, customers_df): def test_add_table_after_server_raises(self, orders_df, customers_df): """Test that adding table after server init raises error.""" qc = QueryChat(orders_df, "orders", greeting="Hello!") - qc._active_sessions = 1 # Simulate a live session + qc._sessions_started = True # Simulate a live session - with pytest.raises(RuntimeError, match="Cannot add tables while a server session"): - qc.add_table(customers_df, "customers") + with pytest.raises(RuntimeError, match="replace or remove"): + qc.add_table(orders_df, "orders", replace=True) class TestRemoveTable: @@ -251,9 +253,9 @@ def test_remove_table_after_server_raises(self, orders_df, customers_df): """Test that removing table after server init raises error.""" qc = QueryChat(orders_df, "orders", greeting="Hello!") qc.add_table(customers_df, "customers") - qc._active_sessions = 1 + qc._sessions_started = True - with pytest.raises(RuntimeError, match="Cannot remove tables while a server session"): + with pytest.raises(RuntimeError, match="replace or remove"): qc.remove_table("customers") @@ -408,7 +410,7 @@ def test_add_table_replace_validates_compatibility(self, orders_qc, customers_df class TestBuildQueryExecutor: def test_executor_none_before_first_use(self, orders_qc): - assert orders_qc._query_executor is None + assert orders_qc._table_set.executor_built is False def test_single_table_uses_data_source_executor(self, orders_qc): assert isinstance(orders_qc._require_query_executor("test"), DataSourceExecutor) @@ -418,15 +420,19 @@ def test_multi_dataframe_uses_duckdb_executor(self, orders_qc, customers_df): assert isinstance(orders_qc._require_query_executor("test"), DuckDBExecutor) def test_executor_invalidated_on_add_table(self, orders_qc, customers_df): - orders_qc._require_query_executor("test") # build it + before = orders_qc._table_set + orders_qc._require_query_executor("test") orders_qc.add_table(customers_df, "customers") - assert orders_qc._query_executor is None + assert orders_qc._table_set is not before + assert orders_qc._table_set.executor_built is False def test_executor_invalidated_on_remove_table(self, orders_qc, customers_df): orders_qc.add_table(customers_df, "customers") - orders_qc._require_query_executor("test") # build it + before = orders_qc._table_set + orders_qc._require_query_executor("test") orders_qc.remove_table("customers") - assert orders_qc._query_executor is None + assert orders_qc._table_set is not before + assert orders_qc._table_set.executor_built is False def test_executor_cached_after_build(self, orders_qc): first = orders_qc._require_query_executor("test") @@ -445,46 +451,18 @@ def test_cleanup_includes_executor(self, orders_qc, customers_df): def test_deferred_executor_is_none(self): qc = QueryChatBase(None, "test") - assert qc._query_executor is None + assert qc._table_set is None def test_rejects_inconsistent_internal_source_group(self, orders_qc): - orders_qc._data_sources["customers"] = normalize_data_source( - pl.DataFrame( - { - "id": [101, 102], - "name": ["Alice", "Bob"], - } + sources = { + "orders": orders_qc._data_sources["orders"], + "customers": normalize_data_source( + pl.DataFrame({"id": [101, 102], "name": ["Alice", "Bob"]}), + "customers", ), - "customers", - ) - + } with pytest.raises(ValueError, match="same DataFrame backend"): - orders_qc._build_query_executor() - - def test_cached_executor_survives_direct_source_mutation( - self, orders_qc, customers_df - ): - """Executor built lazily is not invalidated by direct _data_sources mutation.""" - orders_qc.add_table(customers_df, "customers") - built = orders_qc._require_query_executor("test") - - # Directly corrupt _data_sources (bypassing add_table) — executor should - # not be affected since invalidation only happens through add/remove_table. - orders_qc._data_sources["customers"] = normalize_data_source( - pl.DataFrame({"id": [101, 102], "name": ["Alice", "Bob"]}), - "customers", - ) - - assert orders_qc._query_executor is built - result = built.execute_query( - """ - SELECT customers.name, orders.amount - FROM orders - JOIN customers ON orders.customer_id = customers.id - WHERE orders.id = 1 - """ - ) - assert result.to_dict("records") == [{"name": "Alice", "amount": 100.0}] + build_query_executor(sources) def test_add_table_failure_cleans_staged_source_and_preserves_state( self, orders_qc, customers_df, monkeypatch @@ -508,7 +486,7 @@ def fail_compat(*a, **kw): raise ValueError("compat check failed") monkeypatch.setattr( - "querychat._querychat_base.check_source_compatibility", + "querychat._query_executor.check_source_compatibility", fail_compat, ) @@ -546,7 +524,7 @@ def fail_compat(*a, **kw): raise ValueError("compat check failed") monkeypatch.setattr( - "querychat._querychat_base.check_source_compatibility", + "querychat._query_executor.check_source_compatibility", fail_compat, ) @@ -925,8 +903,8 @@ def test_state_dict_mixin_df_warns_multi_table(self, orders_df, customers_df): class DummyAccessor(StateDictQueryChat): def __init__(self): - self._data_sources = dict(qc._data_sources) - self._query_executor = qc._query_executor + self._table_set = qc._table_set + self._query_executor = qc._require_query_executor("test") self.greeting = None def _require_initialized(self, _m): @@ -954,8 +932,8 @@ def test_state_dict_mixin_sql_warns_multi_table(self, orders_df, customers_df): class DummyAccessor(StateDictQueryChat): def __init__(self): - self._data_sources = dict(qc._data_sources) - self._query_executor = qc._query_executor + self._table_set = qc._table_set + self._query_executor = qc._require_query_executor("test") self.greeting = None def _require_initialized(self, _m): @@ -993,8 +971,8 @@ def test_state_dict_mixin_title_warns_multi_table(self, orders_df, customers_df) class DummyAccessor(StateDictQueryChat): def __init__(self): - self._data_sources = dict(qc._data_sources) - self._query_executor = qc._query_executor + self._table_set = qc._table_set + self._query_executor = qc._require_query_executor("test") self.greeting = None def _require_initialized(self, _m): @@ -1034,8 +1012,8 @@ def test_state_dict_mixin_with_table_kwarg_still_works( class DummyAccessor(StateDictQueryChat): def __init__(self): - self._data_sources = dict(qc._data_sources) - self._query_executor = qc._query_executor + self._table_set = qc._table_set + self._query_executor = qc._require_query_executor("test") self.greeting = None def _require_initialized(self, _m): @@ -1135,7 +1113,6 @@ def test_list_dicts_appear_in_system_prompt(self, orders_df, customers_df) -> No ) qc = QueryChat(orders_df, "orders", data_dict=[dd1, dd2]) qc.add_table(customers_df, "customers") - qc._build_system_prompt() - rendered = qc._system_prompt.render(qc.tools) + rendered = qc._table_set.system_prompt.render(qc.tools) assert 'name="sales"' in rendered assert 'name="people"' in rendered diff --git a/pkg-py/tests/test_multi_table_frameworks.py b/pkg-py/tests/test_multi_table_frameworks.py index 03c5b8be0..daf5f72a1 100644 --- a/pkg-py/tests/test_multi_table_frameworks.py +++ b/pkg-py/tests/test_multi_table_frameworks.py @@ -105,7 +105,7 @@ def _make_accessor(self, orders_df, customers_df): class DummyAccessor(StateDictQueryChat): def __init__(self): - self._data_sources = dict(qc._data_sources) + self._table_set = qc._table_set self._query_executor = qc._require_query_executor("test") self.greeting = None diff --git a/pkg-py/tests/test_pin_source.py b/pkg-py/tests/test_pin_source.py index ebb715a0f..704cef77c 100644 --- a/pkg-py/tests/test_pin_source.py +++ b/pkg-py/tests/test_pin_source.py @@ -226,7 +226,7 @@ def test_auto_fills_data_description(self, board, sample_df): ps = PinSource(board, "cars") qc = QueryChat(data_source=ps, table_name="cars", greeting="Hi") try: - prompt = qc._system_prompt.render(qc.tools) + prompt = qc._table_set.system_prompt.render(qc.tools) assert "Motor Trend Cars" in prompt assert "Road test data" in prompt finally: @@ -249,7 +249,7 @@ def test_explicit_description_overrides_pin_metadata(self, board, sample_df): data_description="Custom description", ) try: - prompt = qc._system_prompt.render(qc.tools) + prompt = qc._table_set.system_prompt.render(qc.tools) assert "Custom description" in prompt assert "Motor Trend Cars" not in prompt finally: @@ -273,7 +273,7 @@ def test_explicit_description_survives_source_change(self, board, sample_df): ) try: qc.add_table(sample_df, "cars", replace=True) - prompt = qc._system_prompt.render(qc.tools) + prompt = qc._table_set.system_prompt.render(qc.tools) assert "Custom description" in prompt assert "Motor Trend Cars" not in prompt finally: @@ -291,11 +291,11 @@ def test_clears_auto_description_on_source_change(self, board, sample_df): ps = PinSource(board, "cars") qc = QueryChat(data_source=ps, table_name="cars", greeting="Hi") try: - prompt_before = qc._system_prompt.render(qc.tools) + prompt_before = qc._table_set.system_prompt.render(qc.tools) assert "Motor Trend Cars" in prompt_before qc.add_table(sample_df, "cars", replace=True) - prompt_after = qc._system_prompt.render(qc.tools) + prompt_after = qc._table_set.system_prompt.render(qc.tools) assert "Motor Trend Cars" not in prompt_after finally: qc.cleanup() diff --git a/pkg-py/tests/test_querychat.py b/pkg-py/tests/test_querychat.py index 6f3e356c5..03b26bd42 100644 --- a/pkg-py/tests/test_querychat.py +++ b/pkg-py/tests/test_querychat.py @@ -2,20 +2,15 @@ import os import tempfile from pathlib import Path -from unittest.mock import patch +from unittest.mock import AsyncMock, MagicMock, patch import ibis -import narwhals.stable.v1 as nw import pandas as pd import polars as pl import pytest from querychat import QueryChat -from querychat._datasource import ( - DataFrameSource, - DataSource, - IbisSource, - PolarsLazySource, -) +from querychat._datasource import IbisSource, PolarsLazySource +from querychat._querychat_base import normalize_data_source from sqlalchemy import create_engine, text @@ -440,10 +435,10 @@ def test_remove_table_prunes_greeter_tables(sqlite_engine): class TestGreeterSnapshotOverrides: """ - _generate_async_snapshot() renders from an explicit tables/data_sources - snapshot instead of live shared state, which a later Shiny session may - have mutated before an earlier session's async greeting runs. The - public build_client()/generate()/generate_async() API is unaffected. + generate_async() renders from an explicit tables/table_set snapshot + instead of live shared state, which a later Shiny session may have + mutated before an earlier session's async greeting runs. The public + build_client()/generate()/generate_async() API is unaffected. """ def test_build_client_uses_live_state(self, sample_df): @@ -463,8 +458,8 @@ async def fake_stream_async(self, *args, **kwargs): with patch("chatlas.Chat.stream_async", fake_stream_async): asyncio.run( - qc.greeter._generate_async_snapshot( - base=None, tables=["test_table"], data_sources=qc._data_sources + qc.greeter.generate_async( + base=None, tables=["test_table"], table_set=qc._table_set ) ) @@ -473,25 +468,16 @@ async def fake_stream_async(self, *args, **kwargs): def test_snapshot_data_sources_override_ignores_live_data_sources(self, sample_df): qc = QueryChat(sample_df, "test_table") - other_df = pd.DataFrame({"z": [1, 2, 3]}) - snapshot: dict[str, DataSource] = { - "other_table": DataFrameSource( - nw.from_native(other_df, eager_only=True), "other_table" - ) - } - seen: dict[str, str | None] = {} - - async def fake_stream_async(self, *args, **kwargs): - seen["system_prompt"] = self.system_prompt - return "stream" + other_set = qc._build_table_set( + {"other": normalize_data_source(pd.DataFrame({"x": [1]}), "other")} + ) + seen = {} - with patch("chatlas.Chat.stream_async", fake_stream_async): - asyncio.run( - qc.greeter._generate_async_snapshot( - base=None, tables=["other_table"], data_sources=snapshot - ) - ) + def factory(tables, prompt, base=None, *, table_set=None): + seen["tables"] = tables + seen["table_set"] = table_set + return MagicMock(stream_async=AsyncMock(return_value="stream")) - assert seen["system_prompt"] is not None - assert "other_table" in seen["system_prompt"] - assert "test_table" not in seen["system_prompt"] + qc.greeter._client_factory = factory + asyncio.run(qc.greeter.generate_async(tables=["other"], table_set=other_set)) + assert seen == {"tables": ["other"], "table_set": other_set} diff --git a/pkg-py/tests/test_server_data_source.py b/pkg-py/tests/test_server_data_source.py index f71a635e8..3a197f459 100644 --- a/pkg-py/tests/test_server_data_source.py +++ b/pkg-py/tests/test_server_data_source.py @@ -68,6 +68,26 @@ def next_session(): return sessions +@pytest.fixture +def session_runs(monkeypatch): + """Each server() call gets a fresh FakeSession; mod_server kwargs are captured.""" + sessions: list[FakeSession] = [] + calls: list[dict] = [] + + def fake_mod_server(*args, **kwargs): + calls.append(kwargs) + return MagicMock() + + def next_session(): + session = FakeSession() + sessions.append(session) + return session + + monkeypatch.setattr(shiny_mod, "mod_server", fake_mod_server) + monkeypatch.setattr(shiny_mod, "get_current_session", next_session) + return sessions, calls + + class TestServerDataSourceRegistersDeferredTable: def test_registers_deferred_table_by_constructor_name( self, users_df, captured_mod_server @@ -75,8 +95,8 @@ def test_registers_deferred_table_by_constructor_name( 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"] + assert qc.table_names() == [] + assert captured_mod_server[0]["table_set"].table_names == ["users"] def test_explicit_table_name_overrides_deferred_name( self, users_df, captured_mod_server @@ -84,20 +104,17 @@ def test_explicit_table_name_overrides_deferred_name( qc = shiny_mod.QueryChat(None, table_name="users") qc.server(data_source=users_df, table_name="people") - assert qc.table_names() == ["people"] + assert captured_mod_server[0]["table_set"].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. - """ + """server(data_source=) with no deferred/explicit name shadows the first registered table for this session.""" 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 qc._data_sources["users"].get_data()["id"].tolist() == [1, 2, 3] + registered = captured_mod_server[0]["table_set"].data_sources["users"] assert registered.get_data()["id"].tolist() == [4, 5] def test_missing_table_name_raises(self, users_df, captured_mod_server): @@ -122,7 +139,7 @@ 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 + assert captured_mod_server[0]["greeting_tables"] == ["users"] def test_no_data_source_leaves_tables_unchanged( self, users_df, captured_mod_server @@ -133,235 +150,163 @@ def test_no_data_source_leaves_tables_unchanged( assert qc.table_names() == ["users"] -class TestServerDataSourceSurvivesSecondSession: - def test_second_session_does_not_raise( - self, users_df, other_users_df, captured_mod_server - ): - qc = shiny_mod.QueryChat(None, table_name="users") - qc.server(data_source=users_df) - - qc.server(data_source=other_users_df) # must not raise - - assert list(captured_mod_server[1]["data_sources"].keys()) == ["users"] - - def test_add_table_still_blocked_after_server_init( - self, users_df, other_users_df, captured_mod_server +class TestServerDataSourceSessionIsolation: + def test_instance_tables_unchanged_by_session_registration( + self, users_df, captured_mod_server ): - """ - The public add_table() guard must remain intact; only the - server(data_source=...) path bypasses it. - """ qc = shiny_mod.QueryChat(None, table_name="users") qc.server(data_source=users_df) - with pytest.raises(RuntimeError, match="Cannot add tables while a server session"): - qc.add_table(other_users_df, "other") - + assert qc.table_names() == [] + assert captured_mod_server[0]["table_set"].table_names == ["users"] -class TestServerDataSourceCleanupSafety: - def test_second_session_does_not_clean_up_first_sessions_source( + def test_each_session_gets_its_own_source( self, users_df, other_users_df, captured_mod_server ): - """ - An earlier, still-running session's executor holds a live - reference to the source a later session's registration replaces. - """ 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 replacement has a single owner, so cleanup-on-replace - is unchanged on the public path. - """ - qc = shiny_mod.QueryChat(users_df, "users") - first_source = qc._data_sources["users"] + qc.server(data_source=other_users_df) - with patch.object(first_source, "cleanup") as mock_cleanup: - qc.add_table(other_users_df, "users", replace=True) - mock_cleanup.assert_called_once() + first = captured_mod_server[0]["table_set"].data_sources["users"] + second = captured_mod_server[1]["table_set"].data_sources["users"] + assert first is not second + assert first.get_data()["id"].tolist() == [1, 2, 3] + assert second.get_data()["id"].tolist() == [4, 5] - def test_second_session_does_not_clean_up_first_sessions_query_executor( + def test_session_table_shadows_config_time_table( self, users_df, other_users_df, captured_mod_server ): - """ - An earlier, still-running session's chat has already captured the - cached executor and may be querying through it. - """ - 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 replacement has a single owner, so executor cleanup - is unchanged on the public path. - """ 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() + config_source = qc._data_sources["users"] - def test_first_server_call_cleans_up_constructor_registered_source( - self, users_df, other_users_df, captured_mod_server - ): - """No session can still be using it, so cleanup-on-replace holds.""" - qc = shiny_mod.QueryChat(users_df, "users") - constructor_source = qc._data_sources["users"] + qc.server(data_source=other_users_df) - with patch.object(constructor_source, "cleanup") as mock_cleanup: - qc.server(data_source=other_users_df) - mock_cleanup.assert_called_once() + assert qc._data_sources["users"] is config_source + session_source = captured_mod_server[0]["table_set"].data_sources["users"] + assert session_source is not config_source + assert session_source.get_data()["id"].tolist() == [4, 5] - def test_second_session_skips_cleanup_even_when_first_cleaned_up( + def test_sessions_do_not_see_each_others_tables( self, users_df, other_users_df, captured_mod_server ): - qc = shiny_mod.QueryChat(users_df, "users") - - qc.server(data_source=other_users_df) - session1_source = qc._data_sources["users"] - - third_df = pd.DataFrame({"id": [7, 8, 9], "name": ["F", "G", "H"]}) - with patch.object(session1_source, "cleanup") as mock_cleanup: - qc.server(data_source=third_df) - mock_cleanup.assert_not_called() + qc = shiny_mod.QueryChat() + qc.add_table(users_df, "orders") + qc.server(data_source=other_users_df, table_name="returns") + qc.server(data_source=pd.DataFrame({"id": [7]}), table_name="orders") -class TestServerDataSourceSessionLifecycle: - """ - The hazard behind cleanup-on-replace and the add/remove_table guards is - *live* sessions, not past ones: a session that has ended can no longer - be using a resource it registered. - """ + assert captured_mod_server[0]["table_set"].table_names == ["orders", "returns"] + assert captured_mod_server[1]["table_set"].table_names == ["orders"] - def test_ended_sessions_source_is_cleaned_up_on_replace( - self, users_df, other_users_df, fake_sessions + def test_greeting_tables_snapshot_is_per_session( + self, users_df, captured_mod_server ): qc = shiny_mod.QueryChat(None, table_name="users") qc.server(data_source=users_df) - first_source = qc._data_sources["users"] - - fake_sessions[0].end() - - with patch.object(first_source, "cleanup") as mock_cleanup: - qc.server(data_source=other_users_df) - mock_cleanup.assert_called_once() - - def test_replaced_source_survives_while_any_session_is_live( - self, users_df, other_users_df, fake_sessions - ): - qc = shiny_mod.QueryChat(None, table_name="users") - qc.server(data_source=users_df) # session 1 - qc.server(data_source=other_users_df) # session 2 replaces s1's source - second_source = qc._data_sources["users"] - fake_sessions[0].end() # s1 ends; s2 still live - - third_df = pd.DataFrame({"id": [7], "name": ["F"]}) - with patch.object(second_source, "cleanup") as mock_cleanup: - qc.server(data_source=third_df) # session 3 replaces s2's source - mock_cleanup.assert_not_called() + assert captured_mod_server[0]["greeting_tables"] == ["users"] + assert qc.greeter.tables == [] - def test_add_table_allowed_once_all_sessions_have_ended( - self, users_df, other_users_df, fake_sessions + def test_greeter_build_client_forwards_table_set( + self, users_df, other_users_df, captured_mod_server ): qc = shiny_mod.QueryChat(None, table_name="users") qc.server(data_source=users_df) + qc.server(data_source=other_users_df) - fake_sessions[0].end() + seen = [] - qc.add_table(other_users_df, "other") # must not raise - assert qc.table_names() == ["users", "other"] + def factory(tables, prompt, base=None, *, table_set=None): + seen.append(table_set) + return MagicMock() + qc.greeter._client_factory = factory + first_set = captured_mod_server[0]["table_set"] + qc.greeter.build_client(tables=["users"], table_set=first_set) + assert seen == [first_set] -class TestServerDataSourceGreetingSnapshot: - def test_server_passes_greeting_tables_snapshot_to_mod_server( - self, users_df, captured_mod_server + +class TestServerDataSourceSessionCleanup: + def test_ending_a_session_closes_only_its_own_source_and_executor( + self, users_df, other_users_df, session_runs ): - """ - Greeting generation runs lazily, after a later session may have - mutated the live greeter.tables -- hence the call-time snapshot. - """ + sessions, calls = session_runs qc = shiny_mod.QueryChat(None, table_name="users") qc.server(data_source=users_df) - - assert captured_mod_server[0]["greeting_tables"] == ["users"] - - -class TestServerDataSourceMixedWithConfigTimeAddTable: - def test_unnamed_registration_replaces_config_time_table( - self, users_df, other_users_df, captured_mod_server + qc.server(data_source=other_users_df) + set_a, set_b = calls[0]["table_set"], calls[1]["table_set"] + src_a, src_b = set_a.data_sources["users"], set_b.data_sources["users"] + + with ( + patch.object(src_a, "cleanup") as cleanup_a, + patch.object(src_b, "cleanup") as cleanup_b, + patch.object(set_a, "cleanup_executor") as exec_a, + patch.object(set_b, "cleanup_executor") as exec_b, + ): + sessions[1].end() + cleanup_b.assert_called_once() + exec_b.assert_called_once() + cleanup_a.assert_not_called() + exec_a.assert_not_called() + + sessions[0].end() + cleanup_a.assert_called_once() + exec_a.assert_called_once() + + def test_session_without_data_source_owns_nothing(self, users_df, session_runs): + sessions, _calls = session_runs + qc = shiny_mod.QueryChat(users_df, "users") + qc.server() + config_source = qc._data_sources["users"] + + with ( + patch.object(config_source, "cleanup") as cleanup, + patch.object(qc._table_set, "cleanup_executor") as cleanup_executor, + ): + sessions[0].end() + cleanup.assert_not_called() + cleanup_executor.assert_not_called() + + def test_config_time_source_survives_until_cleanup( + self, users_df, other_users_df, session_runs ): - qc = shiny_mod.QueryChat() - qc.add_table(users_df, "orders") - + sessions, _calls = session_runs + qc = shiny_mod.QueryChat(users_df, "users") + config_source = qc._data_sources["users"] qc.server(data_source=other_users_df) - # Same table name, but the session's data replaces the config-time data - sources = captured_mod_server[0]["data_sources"] - assert list(sources.keys()) == ["orders"] - assert sources["orders"].get_data()["id"].tolist() == [4, 5] + with patch.object(config_source, "cleanup") as cleanup: + sessions[0].end() + cleanup.assert_not_called() + qc.cleanup() + cleanup.assert_called_once() - def test_replacing_config_time_table_on_first_server_call_cleans_it_up( - self, users_df, other_users_df, captured_mod_server + def test_failed_registration_closes_its_source_and_leaves_instance_untouched( + self, users_df, session_runs, monkeypatch ): - """ - No session is running yet, so the replaced source has a single owner - and cleanup-on-replace still holds (only later sessions skip it). - """ - qc = shiny_mod.QueryChat() - qc.add_table(users_df, "orders") - config_source = qc._data_sources["orders"] + import duckdb + import polars as pl - with patch.object(config_source, "cleanup") as mock_cleanup: - qc.server(data_source=other_users_df) - mock_cleanup.assert_called_once() + sessions, calls = session_runs + qc = shiny_mod.QueryChat(users_df, "users") + created = [] + real_normalize = shiny_mod.normalize_data_source - def test_explicit_table_name_adds_alongside_config_time_table( - self, users_df, other_users_df, captured_mod_server - ): - qc = shiny_mod.QueryChat() - qc.add_table(users_df, "orders") + def spy(data_source, table_name): + source = real_normalize(data_source, table_name) + created.append(source) + return source - qc.server(data_source=other_users_df, table_name="returns") + monkeypatch.setattr(shiny_mod, "normalize_data_source", spy) - sources = captured_mod_server[0]["data_sources"] - assert list(sources.keys()) == ["orders", "returns"] - # The config-time table's own data is untouched - assert sources["orders"].get_data()["id"].tolist() == [1, 2, 3] - assert sources["returns"].get_data()["id"].tolist() == [4, 5] + with pytest.raises(ValueError, match="same DataFrame backend"): + qc.server(data_source=pl.DataFrame({"id": [1]}), table_name="other") - def test_later_session_snapshot_includes_earlier_sessions_table( - self, users_df, other_users_df, captured_mod_server - ): - """The registry is shared and cumulative across sessions.""" - qc = shiny_mod.QueryChat() - qc.add_table(users_df, "orders") + (session_source,) = created + with pytest.raises(duckdb.ConnectionException): + session_source.execute_query("SELECT 1") + assert qc.table_names() == ["users"] + assert sessions[0]._ended_callbacks == [] - # Session 1 adds its own table alongside the config-time one - qc.server(data_source=other_users_df, table_name="returns") - # Session 2 replaces "orders" only -- but still sees session 1's table - third_df = pd.DataFrame({"id": [7, 8, 9]}) - qc.server(data_source=third_df, table_name="orders") - - sources = captured_mod_server[1]["data_sources"] - assert list(sources.keys()) == ["orders", "returns"] - assert sources["orders"].get_data()["id"].tolist() == [7, 8, 9] - assert sources["returns"].get_data()["id"].tolist() == [4, 5] + qc.server() + assert calls[-1]["table_set"].table_names == ["users"] diff --git a/pkg-py/tests/test_shiny_module.py b/pkg-py/tests/test_shiny_module.py index 8bc884f3f..e03764744 100644 --- a/pkg-py/tests/test_shiny_module.py +++ b/pkg-py/tests/test_shiny_module.py @@ -3,6 +3,7 @@ from __future__ import annotations import os +from types import SimpleNamespace from unittest.mock import patch import pytest @@ -11,6 +12,12 @@ from shiny import ui +def fake_table_set(source, executor): + return SimpleNamespace( + data_sources={"t": source}, executor=executor, table_names=["t"] + ) + + @pytest.fixture(autouse=True) def set_dummy_api_key(): old = os.environ.get("OPENAI_API_KEY") @@ -131,8 +138,7 @@ def fake_chat_constructor( fake_input, MagicMock(), fake_session, - data_sources={"t": fake_source}, - executor=fake_executor, + table_set=fake_table_set(fake_source, fake_executor), greeting=None, client=client_factory, history=True, @@ -184,8 +190,9 @@ def client_factory(**kwargs): return MagicMock(spec=["stream_async"]) fake_greeter = MagicMock() - fake_greeter._generate_async_snapshot = AsyncMock(return_value=MagicMock()) + fake_greeter.generate_async = AsyncMock(return_value=MagicMock()) fake_greeting_base = MagicMock() + table_set = fake_table_set(fake_source, fake_executor) inner_fn = _unwrap_module_server(mod_server) @@ -207,8 +214,7 @@ def client_factory(**kwargs): fake_input, MagicMock(), fake_session, - data_sources={"t": fake_source}, - executor=fake_executor, + table_set=table_set, greeting=None, client=client_factory, history=True, @@ -220,10 +226,10 @@ def client_factory(**kwargs): asyncio.run(captured["greeting"]()) - fake_greeter._generate_async_snapshot.assert_called_once_with( + fake_greeter.generate_async.assert_called_once_with( base=fake_greeting_base, tables=["t"], - data_sources={"t": fake_source}, + table_set=table_set, ) @@ -271,8 +277,7 @@ def client_factory(**kwargs): fake_input, MagicMock(), fake_session, - data_sources={"t": fake_source}, - executor=fake_executor, + table_set=fake_table_set(fake_source, fake_executor), greeting=None, client=client_factory, history=True, @@ -329,8 +334,7 @@ def client_factory(**kwargs): fake_input, MagicMock(), fake_session, - data_sources={"t": fake_source}, - executor=fake_executor, + table_set=fake_table_set(fake_source, fake_executor), greeting=None, client=client_factory, history=HistoryOptions(restore_mode="bookmark"), @@ -381,8 +385,7 @@ def client_factory(**kwargs): fake_input, MagicMock(), fake_session, - data_sources={"t": fake_source}, - executor=fake_executor, + table_set=fake_table_set(fake_source, fake_executor), greeting=None, client=client_factory, history=False, # even with history disabled, registration must still happen diff --git a/pkg-py/tests/test_state.py b/pkg-py/tests/test_state.py index 5a2911d98..e404c27a3 100644 --- a/pkg-py/tests/test_state.py +++ b/pkg-py/tests/test_state.py @@ -403,7 +403,7 @@ def client_factory(update_callback, reset_callback): class DummyStateAccessor(StateDictQueryChat[pd.DataFrame]): def __init__(self, qc: QueryChat): - self._data_sources = dict(qc._data_sources) + self._table_set = qc._table_set self._query_executor = qc._require_query_executor("test") self.greeting = None From 3c53e8dd8f9966dda2992b8d6eeaa942ac17dfbc Mon Sep 17 00:00:00 2001 From: Carson Date: Sat, 12 Sep 2026 18:38:15 -0500 Subject: [PATCH 07/35] refactor(r): give each Shiny session a read-only TableSet instead of mutating QueryChat $server(data_source = ) derives a session-local TableSet and closes what it created in onSessionEnded(). No live-session counting or retired-resource deferral is added. Late $add_table() warns and parks the superseded set for $cleanup(); replace/remove after sessions start errors. --- pkg-r/R/QueryChat.R | 399 +++++++++++------- pkg-r/R/QueryChatGreeter.R | 16 +- pkg-r/R/querychat_module.R | 15 +- pkg-r/man/QueryChat.Rd | 29 +- pkg-r/tests/testthat/helper-fixtures.R | 2 +- pkg-r/tests/testthat/test-QueryChat.R | 81 +++- pkg-r/tests/testthat/test-querychat_module.R | 84 +--- .../tests/testthat/test-server_data_source.R | 235 +++++++++++ 8 files changed, 620 insertions(+), 241 deletions(-) create mode 100644 pkg-r/tests/testthat/test-server_data_source.R diff --git a/pkg-r/R/QueryChat.R b/pkg-r/R/QueryChat.R index 8b3b560a3..582a0afe1 100644 --- a/pkg-r/R/QueryChat.R +++ b/pkg-r/R/QueryChat.R @@ -89,13 +89,14 @@ QueryChat <- R6::R6Class( "QueryChat", private = list( - .data_sources = list(), + .table_set = NULL, + # Instance sets swapped out by $add_table()/$add_tables() after a session + # started. A running session may still hold one, so $cleanup() closes them. + .superseded_table_sets = list(), + .sessions_started = FALSE, .deferred_table_name = NULL, - .query_executor = NULL, - .server_initialized = FALSE, .client_spec = NULL, .client_console = NULL, - .system_prompt = NULL, # Store init parameters for deferred system prompt building .prompt_template = NULL, .data_description = NULL, @@ -105,16 +106,44 @@ QueryChat <- R6::R6Class( .data_dicts = list(), .greeter = NULL, - require_initialized = function(method_name) { - if (length(private$.data_sources) == 0) { + data_sources = function() { + if (is.null(private$.table_set)) { + return(list()) + } + private$.table_set$data_sources + }, + + require_table_set = function(method_name) { + if (is.null(private$.table_set)) { cli::cli_abort( "{.arg data_source} must be set before calling {.fn ${method_name}}. Either pass {.arg data_source} to {.fn $new}, or call {.fn $add_table}." ) } + private$.table_set + }, + + require_initialized = function(method_name) { + private$require_table_set(method_name) + invisible(NULL) + }, + + # Non-mutating counterpart of auto_fill_data_description(), for sets built + # on behalf of a session. + resolve_data_description = function(sources) { + if (private$.data_description_mode == "supplied") { + return(private$.data_description) + } + if (length(sources) == 1) { + desc <- sources[[1]]$get_data_description() + if (nzchar(desc %||% "")) { + return(desc) + } + } + NULL }, - auto_fill_data_description = function(sources = private$.data_sources) { + auto_fill_data_description = function(sources = private$data_sources()) { if (length(sources) != 1) { return() } @@ -131,26 +160,68 @@ QueryChat <- R6::R6Class( } }, - build_system_prompt = function(data_sources = NULL) { - sources <- data_sources %||% private$.data_sources + build_table_set = function( + sources, + data_description = private$.data_description + ) { if (length(sources) == 0) { cli::cli_abort("Cannot build system prompt without data sources") } - + validate_source_group_compatibility(sources) prompt_template <- private$.prompt_template %||% system.file("prompts", "prompt.md", package = "querychat") - - private$.system_prompt <- QueryChatSystemPrompt$new( + system_prompt <- QueryChatSystemPrompt$new( prompt_template = prompt_template, data_sources = sources, - data_description = private$.data_description, + data_description = data_description, extra_instructions = private$.extra_instructions, categorical_threshold = private$.categorical_threshold, data_dicts = private$.data_dicts ) + TableSet$new(sources, system_prompt, data_description = data_description) + }, + + check_late_change = function(method_name, destructive) { + if (!private$.sessions_started) { + return(invisible(NULL)) + } + if (destructive) { + cli::cli_abort(c( + "Cannot call {.fn ${method_name}} to replace or remove a table while sessions may be using it.", + "i" = "Configure all tables before calling {.fn $server} or {.fn $app}." + )) + } + cli::cli_warn(c( + "{.fn ${method_name}} called after a session has started.", + "i" = "Sessions that are already running keep the tables they started with; only new sessions will see this change." + )) + invisible(NULL) + }, + + swap_table_set = function(new_set, replaced = list()) { + old_set <- private$.table_set + private$.table_set <- new_set + if (is.null(old_set)) { + return(invisible(NULL)) + } + if (private$.sessions_started) { + # `replaced` is empty here: check_late_change() rejects replacement + # once sessions have started. + private$.superseded_table_sets <- c( + private$.superseded_table_sets, + list(old_set) + ) + return(invisible(NULL)) + } + tryCatch(old_set$cleanup_executor(), error = function(e) NULL) + for (source in replaced) { + source$cleanup() + } + invisible(NULL) }, create_session_client = function( + table_set, client_spec = NULL, tools = NA, handoff_available = FALSE, @@ -168,7 +239,7 @@ QueryChat <- R6::R6Class( tools <- check_viz_deps(tools) chat$set_system_prompt( - private$.system_prompt$render( + table_set$system_prompt$render( tools = tools, handoff_available = handoff_available ) @@ -178,14 +249,9 @@ QueryChat <- R6::R6Class( return(chat) } - # Build executor lazily - if (is.null(private$.query_executor)) { - private$.query_executor <- build_query_executor(private$.data_sources) - } - executor <- private$.query_executor - tbl_names <- names(private$.data_sources) + executor <- table_set$executor() + tbl_names <- table_set$table_names() - # Always register get_schema tool chat$register_tool( tool_get_schema( private$.data_dicts, @@ -381,9 +447,9 @@ QueryChat <- R6::R6Class( } } normalized <- normalize_data_source(data_source, table_name) - private$.data_sources[[normalized$table_name]] <- normalized - private$auto_fill_data_description() - private$build_system_prompt() + sources <- stats::setNames(list(normalized), normalized$table_name) + private$auto_fill_data_description(sources) + private$.table_set <- private$build_table_set(sources) self$greeter$tables <- c(self$greeter$tables, normalized$table_name) self$id <- id %||% sprintf("querychat_%s", normalized$table_name) } else { @@ -417,6 +483,9 @@ QueryChat <- R6::R6Class( #' @description #' Add a table to this QueryChat instance. #' + #' Replacing or removing an existing table after a session has started is + #' an error; adding a new one warns. + #' #' @param data_source A data frame, database connection, or DataSource object. #' @param table_name The SQL table name for this data source. #' @param replace Whether to replace an existing table with this name. @@ -431,16 +500,17 @@ QueryChat <- R6::R6Class( replace = FALSE, include_in_greeting = FALSE ) { - if (private$.server_initialized) { - cli::cli_abort("Cannot add tables after server initialization.") - } check_bool(include_in_greeting) check_sql_table_name(table_name) - if (table_name %in% names(private$.data_sources) && !replace) { + current <- private$data_sources() + exists <- table_name %in% names(current) + if (exists && !replace) { cli::cli_abort( "Table {.val {table_name}} already exists. Use {.code replace = TRUE} to replace." ) } + private$check_late_change("add_table", destructive = exists) + if ( is_data_source(data_source) && !identical(data_source$table_name, table_name) @@ -452,21 +522,13 @@ QueryChat <- R6::R6Class( ) ) } - normalized <- normalize_data_source(data_source, table_name) - - other_sources <- private$.data_sources[ - names(private$.data_sources) != table_name - ] - check_source_compatibility(other_sources, normalized, table_name) - next_sources <- private$.data_sources + normalized <- normalize_data_source(data_source, table_name) + next_sources <- current next_sources[[table_name]] <- normalized - private$auto_fill_data_description(next_sources) - tryCatch( - { - private$build_system_prompt(data_sources = next_sources) - }, + new_set <- tryCatch( + private$build_table_set(next_sources), error = function(e) { if (!inherits(data_source, "DataSource")) { normalized$cleanup() @@ -475,18 +537,17 @@ QueryChat <- R6::R6Class( } ) - old_source <- private$.data_sources[[table_name]] - private$.data_sources <- next_sources - if (!is.null(old_source) && !identical(old_source, normalized)) { - old_source$cleanup() - } - - if (!is.null(private$.query_executor)) { - tryCatch(private$.query_executor$cleanup(), error = function(e) NULL) - private$.query_executor <- NULL + old_source <- current[[table_name]] + replaced <- if ( + !is.null(old_source) && !identical(old_source, normalized) + ) { + list(old_source) + } else { + list() } + private$swap_table_set(new_set, replaced = replaced) - if (isTRUE(include_in_greeting)) { + if (isTRUE(include_in_greeting) && !table_name %in% self$greeter$tables) { self$greeter$tables <- c(self$greeter$tables, table_name) } @@ -500,6 +561,9 @@ QueryChat <- R6::R6Class( #' system prompt exactly once after all tables have been staged, avoiding #' N-1 spurious intermediate rebuilds. #' + #' Replacing or removing an existing table after a session has started is + #' an error; adding a new one warns. + #' #' @param conn A DBI connection. Only DBI connections are supported; pass #' individual data frames or other sources via `$add_table()`. #' @param tables Table names to register. When `NULL`, all tables returned @@ -518,9 +582,6 @@ QueryChat <- R6::R6Class( replace = FALSE, include_in_greeting = FALSE ) { - if (private$.server_initialized) { - cli::cli_abort("Cannot add tables after server initialization.") - } if (!inherits(conn, "DBIConnection")) { cli::cli_abort( "{.fn add_tables} requires a {.cls DBIConnection}, not {.obj_type_friendly {conn}}.", @@ -533,14 +594,20 @@ QueryChat <- R6::R6Class( if (length(tables) == 0) { cli::cli_abort("No tables found in database.") } + current <- private$data_sources() for (table_name in tables) { check_sql_table_name(table_name) - if (table_name %in% names(private$.data_sources) && !replace) { - cli::cli_abort( - "Table {.val {table_name}} already exists. Use {.code replace = TRUE} to replace." - ) - } } + existing <- intersect(tables, names(current)) + if (length(existing) > 0 && !replace) { + cli::cli_abort( + "Table {.val {existing[[1]]}} already exists. Use {.code replace = TRUE} to replace." + ) + } + private$check_late_change( + "add_tables", + destructive = length(existing) > 0 + ) if ( !rlang::is_bool(include_in_greeting) && @@ -562,47 +629,32 @@ QueryChat <- R6::R6Class( lapply(tables, function(tbl) normalize_data_source(conn, tbl)), tables ) - - staged <- list() - for (table_name in tables) { - other_sources <- private$.data_sources[ - names(private$.data_sources) != table_name - ] - check_source_compatibility( - c(other_sources, staged), - normalized[[table_name]], - table_name - ) - staged[[table_name]] <- normalized[[table_name]] - } - - next_sources <- private$.data_sources + next_sources <- current for (table_name in tables) { next_sources[[table_name]] <- normalized[[table_name]] } - private$auto_fill_data_description(next_sources) - private$build_system_prompt(data_sources = next_sources) + new_set <- private$build_table_set(next_sources) + replaced <- list() for (table_name in tables) { - old_source <- private$.data_sources[[table_name]] + old_source <- current[[table_name]] if ( !is.null(old_source) && !identical(old_source, normalized[[table_name]]) ) { - old_source$cleanup() + replaced <- c(replaced, list(old_source)) } } - private$.data_sources <- next_sources - - if (!is.null(private$.query_executor)) { - tryCatch(private$.query_executor$cleanup(), error = function(e) NULL) - private$.query_executor <- NULL - } + private$swap_table_set(new_set, replaced = replaced) - if (length(greeting_tbls) > 0) { - self$greeter$tables <- c(self$greeter$tables, greeting_tbls) + new_greeting <- self$greeter$tables + for (name in greeting_tbls) { + if (!name %in% new_greeting) { + new_greeting <- c(new_greeting, name) + } } + self$greeter$tables <- new_greeting invisible(self) }, @@ -610,44 +662,35 @@ QueryChat <- R6::R6Class( #' @description #' Remove a table from this QueryChat instance. #' + #' Removing an existing table after a session has started is an error. + #' #' @param table_name The name of the table to remove. #' #' @return Invisibly returns `self` for chaining. remove_table = function(table_name) { - if (private$.server_initialized) { - cli::cli_abort("Cannot remove tables after server initialization.") - } - if (!table_name %in% names(private$.data_sources)) { + private$check_late_change("remove_table", destructive = TRUE) + current <- private$data_sources() + if (!table_name %in% names(current)) { cli::cli_abort("Table {.val {table_name}} not found.") } - if (length(private$.data_sources) == 1) { + if (length(current) == 1) { cli::cli_abort( "Cannot remove last table. At least one table is required." ) } - removed <- private$.data_sources[[table_name]] - next_sources <- private$.data_sources[ - names(private$.data_sources) != table_name - ] - private$build_system_prompt(data_sources = next_sources) - private$.data_sources <- next_sources + removed <- current[[table_name]] + next_sources <- current[names(current) != table_name] + new_set <- private$build_table_set(next_sources) + private$swap_table_set(new_set, replaced = list(removed)) if (!is.null(private$.greeter)) { - private$.greeter$tables <- setdiff( - private$.greeter$tables, - table_name - ) - } - if (!is.null(private$.query_executor)) { - tryCatch(private$.query_executor$cleanup(), error = function(e) NULL) - private$.query_executor <- NULL + private$.greeter$tables <- setdiff(private$.greeter$tables, table_name) } - removed$cleanup() invisible(self) }, #' @description #' Return the names of all registered tables. - table_names = function() names(private$.data_sources), + table_names = function() names(private$data_sources()) %||% character(), #' @description #' Create a chat client, complete with registered tools, for the current @@ -687,6 +730,7 @@ QueryChat <- R6::R6Class( } private$create_session_client( + table_set = private$require_table_set("$client"), tools = tools, session = session, update_dashboard = update_dashboard, @@ -755,8 +799,8 @@ QueryChat <- R6::R6Class( ) && identical(resolved_history$restore_mode, "bookmark") - first_table_name <- names(private$.data_sources)[[1]] - table_names <- names(private$.data_sources) + first_table_name <- names(private$data_sources())[[1]] + table_names <- names(private$data_sources()) multi_table <- length(table_names) > 1 ui <- function(req) { @@ -1059,12 +1103,13 @@ QueryChat <- R6::R6Class( #' @description #' Initialize the querychat server logic. #' - #' @param 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 connection scoped to - #' per-user OAuth credentials). Registered under `table_name` if given, - #' otherwise the `table_name` passed to `$new()`, or the first - #' already-registered table. + #' @param data_source Optional data source to register for this session + #' only, for the deferred pattern where the source can't be created + #' until the server function runs (for example a per-user database + #' connection). The instance's own tables are not modified; a + #' same-named instance table is shadowed for this session; any + #' connection querychat created for it is cleaned up when the session + #' ends. #' @param client Optional chat client override for this session. #' @param history Conversation history configuration for this call. Overrides #' the value set on `$new()`. Resolves to `TRUE` when neither this nor the @@ -1105,10 +1150,15 @@ QueryChat <- R6::R6Class( ) } + private$.sessions_started <- TRUE + table_set <- private$.table_set + greeting_tables <- self$greeter$tables + session_source <- NULL + if (!is.null(data_source)) { tbl_name <- table_name %||% private$.deferred_table_name if (is.null(tbl_name)) { - existing_tables <- names(private$.data_sources) + existing_tables <- names(private$data_sources()) if (length(existing_tables) > 0) { tbl_name <- existing_tables[[1]] } @@ -1121,20 +1171,54 @@ QueryChat <- R6::R6Class( ) ) } - self$add_table( - data_source, - tbl_name, - replace = TRUE, - include_in_greeting = TRUE + check_sql_table_name(tbl_name) + if ( + is_data_source(data_source) && + !identical(data_source$table_name, tbl_name) + ) { + cli::cli_abort( + c( + "{.arg data_source}'s own table name ({.val {data_source$table_name}}) does not match the given {.arg table_name} ({.val {tbl_name}}).", + "i" = "Pass a matching {.arg table_name}, or omit it to use {.val {data_source$table_name}}." + ) + ) + } + session_source <- normalize_data_source(data_source, tbl_name) + next_sources <- private$data_sources() + next_sources[[tbl_name]] <- session_source + table_set <- tryCatch( + private$build_table_set( + next_sources, + data_description = private$resolve_data_description(next_sources) + ), + error = function(e) { + if (!inherits(data_source, "DataSource")) { + session_source$cleanup() + } + stop(e) + } ) + if (!tbl_name %in% greeting_tables) { + greeting_tables <- c(greeting_tables, tbl_name) + } } - private$require_initialized("$server") - - private$.server_initialized <- TRUE + if (is.null(table_set)) { + private$require_table_set("$server") + } - if (is.null(private$.query_executor)) { - private$.query_executor <- build_query_executor(private$.data_sources) + if (!is.null(session_source)) { + session_table_set <- table_set + session$onSessionEnded(function() { + warn_on_cleanup_failure( + session_table_set$cleanup_executor(), + "session query executor" + ) + warn_on_cleanup_failure( + session_source$cleanup(), + "session data source" + ) + }) } resolved_client_spec <- client %||% private$.client_spec @@ -1142,6 +1226,7 @@ QueryChat <- R6::R6Class( create_session_client <- function(...) { private$create_session_client( + table_set = table_set, client_spec = base_client, ... ) @@ -1165,18 +1250,17 @@ QueryChat <- R6::R6Class( }) %||% TRUE - result <- mod_server( + mod_server( id %||% self$id, - data_sources = private$.data_sources, - executor = private$.query_executor, + table_set = table_set, greeting = self$greeting, client = create_session_client, tools = self$tools, history = resolved_history, greeter = self$greeter, - greeting_base = base_client + greeting_base = base_client, + greeting_tables = greeting_tables ) - result }, #' @description @@ -1193,15 +1277,26 @@ QueryChat <- R6::R6Class( }, #' @description - #' Clean up resources associated with the data source. + #' Clean up resources this object created. + #' + #' Closes the query executors and data-source connections querychat opened + #' (in-memory DuckDB), including those of table sets superseded by a late + #' `$add_table()`. Connections you passed in are never closed. #' - #' @return Invisibly returns `NULL`. Resources are cleaned up internally. + #' @return Invisibly returns `NULL`. cleanup = function() { - if (!is.null(private$.query_executor)) { - private$.query_executor$cleanup() + for (ts in private$.superseded_table_sets) { + warn_on_cleanup_failure(ts$cleanup_executor(), "query executor") } - for (source in private$.data_sources) { - source$cleanup() + private$.superseded_table_sets <- list() + if (!is.null(private$.table_set)) { + warn_on_cleanup_failure( + private$.table_set$cleanup_executor(), + "query executor" + ) + for (source in private$.table_set$data_sources) { + warn_on_cleanup_failure(source$cleanup(), "data source") + } } invisible(NULL) } @@ -1217,11 +1312,21 @@ QueryChat <- R6::R6Class( return(invisible(value)) } if (is.null(private$.greeter)) { - client_factory <- function(tables, prompt, base = NULL) { + client_factory <- function( + tables, + prompt, + base = NULL, + table_set = NULL + ) { + ts <- table_set %||% private$.table_set sp <- QueryChatSystemPrompt$new( prompt_template = prompt, - data_sources = private$.data_sources, - data_description = private$.data_description, + data_sources = if (is.null(ts)) list() else ts$data_sources, + data_description = if (is.null(ts)) { + private$.data_description + } else { + ts$data_description + }, extra_instructions = NULL, categorical_threshold = private$.categorical_threshold, data_dicts = private$.data_dicts, @@ -1242,8 +1347,9 @@ QueryChat <- R6::R6Class( #' @field system_prompt Get the system prompt. system_prompt = function() { - private$require_initialized("$system_prompt") - private$.system_prompt$render(tools = self$tools) + private$require_table_set("$system_prompt")$system_prompt$render( + tools = self$tools + ) }, #' @field data_source Removed. Use `$add_table()` and `$remove_table()` to manage tables. @@ -1458,6 +1564,17 @@ check_viz_deps <- function(tools) { setdiff(tools, "visualize") } +# Runs one teardown step, warning instead of erroring so the rest still runs. +warn_on_cleanup_failure <- function(expr, what) { + tryCatch( + expr, + error = function(e) { + cli::cli_warn("Failed to clean up {what}: {conditionMessage(e)}") + } + ) + invisible(NULL) +} + normalize_data_source <- function(data_source, table_name) { if (is_data_source(data_source)) { return(data_source) diff --git a/pkg-r/R/QueryChatGreeter.R b/pkg-r/R/QueryChatGreeter.R index 1af88edff..47017e493 100644 --- a/pkg-r/R/QueryChatGreeter.R +++ b/pkg-r/R/QueryChatGreeter.R @@ -13,7 +13,7 @@ QueryChatGreeter <- R6::R6Class( .prompt = NULL ), public = list( - #' @param client_factory function(tables, prompt, base) returning a configured greeting client. + #' @param client_factory function(tables, prompt, base = NULL, table_set = NULL) returning a configured greeting client. initialize = function(client_factory) { private$.client_factory <- client_factory private$.tables <- character() @@ -26,8 +26,18 @@ QueryChatGreeter <- R6::R6Class( #' @description Build a fresh greeting client (no history) configured with the greeting system prompt. #' @param base Optional resolved client to clone (resolve-once base from `$server()`). - build_client = function(base = NULL) { - private$.client_factory(private$.tables, private$.prompt, base) + #' @param tables Table names to describe. Defaults to `$tables`. `$server()` + #' passes the session's snapshot so a lazily generated greeting describes + #' the session that asked for it. + #' @param table_set Optional `TableSet` to build the prompt from. Defaults to + #' the owning QueryChat's instance set. + build_client = function(base = NULL, tables = NULL, table_set = NULL) { + private$.client_factory( + tables %||% private$.tables, + private$.prompt, + base, + table_set = table_set + ) }, #' @description Generate a greeting synchronously and return it as text. diff --git a/pkg-r/R/querychat_module.R b/pkg-r/R/querychat_module.R index 710f21042..6020c2275 100644 --- a/pkg-r/R/querychat_module.R +++ b/pkg-r/R/querychat_module.R @@ -45,16 +45,19 @@ querychat_dependency <- function() { # Main module server function mod_server <- function( id, - data_sources, - executor, + table_set, greeting, client, tools, history, greeter = NULL, - greeting_base = NULL + greeting_base = NULL, + greeting_tables = NULL ) { shiny::moduleServer(id, function(input, output, session) { + data_sources <- table_set$data_sources + executor <- table_set$executor() + current_table_val <- shiny::reactiveVal(NULL, label = "current_table") # Per-table reactive state @@ -130,7 +133,11 @@ mod_server <- function( "i" = "For faster startup, lower cost, and determinism, consider providing a {.arg greeting} to {.fn QueryChat}.", "i" = "You can use your {.help querychat::QueryChat} object's {.fn $generate_greeting} method to generate a greeting." )) - greeting_client <- greeter$build_client(greeting_base) + greeting_client <- greeter$build_client( + greeting_base, + tables = greeting_tables, + table_set = table_set + ) stream <- greeting_client$stream_async(GREETING_PROMPT) shinychat::chat_greeting(stream, persistent = TRUE) } diff --git a/pkg-r/man/QueryChat.Rd b/pkg-r/man/QueryChat.Rd index f2b2837a3..19b4c4a98 100644 --- a/pkg-r/man/QueryChat.Rd +++ b/pkg-r/man/QueryChat.Rd @@ -242,6 +242,9 @@ or \code{FALSE} to never clean up automatically.} \if{latex}{\out{\hypertarget{method-QueryChat-add_table}{}}} \subsection{\code{QueryChat$add_table()}}{ Add a table to this QueryChat instance. + +Replacing or removing an existing table after a session has started is +an error; adding a new one warns. \subsection{Usage}{ \if{html}{\out{
}} \preformatted{QueryChat$add_table( @@ -278,6 +281,9 @@ context. Default is \code{FALSE}.} Unlike calling \verb{$add_table()} repeatedly, this method builds the system prompt exactly once after all tables have been staged, avoiding N-1 spurious intermediate rebuilds. + +Replacing or removing an existing table after a session has started is +an error; adding a new one warns. \subsection{Usage}{ \if{html}{\out{
}} \preformatted{QueryChat$add_tables( @@ -314,6 +320,8 @@ the tables being added). Any other type raises an error.} \if{latex}{\out{\hypertarget{method-QueryChat-remove_table}{}}} \subsection{\code{QueryChat$remove_table()}}{ Remove a table from this QueryChat instance. + +Removing an existing table after a session has started is an error. \subsection{Usage}{ \if{html}{\out{
}} \preformatted{QueryChat$remove_table(table_name)} @@ -571,12 +579,13 @@ and \code{window_title} is omitted, it is also used as the document title.} \subsection{Arguments}{ \if{html}{\out{
}} \describe{ - \item{\code{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 connection scoped to -per-user OAuth credentials). Registered under \code{table_name} if given, -otherwise the \code{table_name} passed to \verb{$new()}, or the first -already-registered table.} + \item{\code{data_source}}{Optional data source to register for this session +only, for the deferred pattern where the source can't be created +until the server function runs (for example a per-user database +connection). The instance's own tables are not modified; a +same-named instance table is shadowed for this session; any +connection querychat created for it is cleaned up when the session +ends.} \item{\code{client}}{Optional chat client override for this session.} \item{\code{history}}{Conversation history configuration for this call. Overrides the value set on \verb{$new()}. Resolves to \code{TRUE} when neither this nor the @@ -628,14 +637,18 @@ or \code{NULL} before any query. \if{html}{\out{}} \if{latex}{\out{\hypertarget{method-QueryChat-cleanup}{}}} \subsection{\code{QueryChat$cleanup()}}{ - Clean up resources associated with the data source. + Clean up resources this object created. + +Closes the query executors and data-source connections querychat opened +(in-memory DuckDB), including those of table sets superseded by a late +\verb{$add_table()}. Connections you passed in are never closed. \subsection{Usage}{ \if{html}{\out{
}} \preformatted{QueryChat$cleanup()} \if{html}{\out{
}} } \subsection{Returns}{ - Invisibly returns \code{NULL}. Resources are cleaned up internally. + Invisibly returns \code{NULL}. } } diff --git a/pkg-r/tests/testthat/helper-fixtures.R b/pkg-r/tests/testthat/helper-fixtures.R index 3fab7caa8..9b254fa07 100644 --- a/pkg-r/tests/testthat/helper-fixtures.R +++ b/pkg-r/tests/testthat/helper-fixtures.R @@ -2,7 +2,7 @@ # Access the internal data source for a named table (test helper only) qc_data_source <- function(qc, table_name) { - qc$.__enclos_env__$private$.data_sources[[table_name]] + qc$.__enclos_env__$private$.table_set$data_sources[[table_name]] } # Simple data frame with id, name, and value columns diff --git a/pkg-r/tests/testthat/test-QueryChat.R b/pkg-r/tests/testthat/test-QueryChat.R index 0b810131c..107bb1e69 100644 --- a/pkg-r/tests/testthat/test-QueryChat.R +++ b/pkg-r/tests/testthat/test-QueryChat.R @@ -740,10 +740,6 @@ test_that("QueryChat$server() resolves history (explicit > constructor > TRUE) a skip_if_no_dataframe_engine() withr::local_envvar(OPENAI_API_KEY = "boop") - ds <- local_data_frame_source(new_test_df()) - executor <- build_query_executor(list(test_table = ds)) - withr::defer(executor$cleanup()) - captured <- NULL local_mocked_bindings( mod_server = function(id, ..., history) { @@ -754,7 +750,6 @@ test_that("QueryChat$server() resolves history (explicit > constructor > TRUE) a ) qc <- local_querychat(history = FALSE) - qc$.__enclos_env__$private$.query_executor <- executor expect_no_warning( shiny::testServer( @@ -774,7 +769,6 @@ test_that("QueryChat$server() resolves history (explicit > constructor > TRUE) a ) qc_no_history <- local_querychat(client = mock_ellmer_chat_client()) - qc_no_history$.__enclos_env__$private$.query_executor <- executor shiny::testServer( function(input, output, session) qc_no_history$server(), { @@ -846,6 +840,7 @@ describe("QueryChat internal client handoff availability", { public = list( internal_client = function(handoff_available = FALSE) { private$create_session_client( + table_set = private$.table_set, tools = NULL, handoff_available = handoff_available ) @@ -1142,7 +1137,7 @@ describe("QueryChat deferred client with $server()", { ) }) - it("$server(data_source=, table_name=) registers under the given name", { + it("$server(data_source=, table_name=) registers under the given name for that session only", { skip_if_no_dataframe_engine() qc <- QueryChat$new( NULL, @@ -1154,9 +1149,11 @@ describe("QueryChat deferred client with $server()", { function(input, output, session) { qc$server(data_source = new_users_df(), table_name = "users") }, - {} + { + expect_equal(session$returned$table_names(), "users") + } ) - expect_equal(qc$table_names(), "users") + expect_equal(qc$table_names(), character()) }) it("id stays fixed across deferred registration (no desync from an already-rendered UI)", { @@ -1340,16 +1337,6 @@ describe("QueryChat$add_tables()", { ) }) - it("calling after server initialization raises error", { - conn <- local_multi_table_conn() - qc <- QueryChat$new(NULL, "placeholder", greeting = "Test") - qc$.__enclos_env__$private$.server_initialized <- TRUE - expect_error( - qc$add_tables(conn), - "after server initialization" - ) - }) - it("system prompt built exactly once for multiple tables", { conn <- local_multi_table_conn() qc <- QueryChat$new(NULL, "placeholder", greeting = "Test") @@ -1366,6 +1353,62 @@ describe("QueryChat$add_tables()", { }) }) +describe("QueryChat table changes after a session has started", { + it("warns when adding a new table and keeps the old set for $cleanup()", { + skip_if_no_dataframe_engine() + qc <- local_querychat(new_users_df(), "users", greeting = "hi") + old_set <- qc$.__enclos_env__$private$.table_set + old_set$executor() + qc$.__enclos_env__$private$.sessions_started <- TRUE + + expect_warning( + qc$add_table(new_test_df(), "other"), + "after a session has started" + ) + + expect_equal(qc$table_names(), c("users", "other")) + expect_identical( + qc$.__enclos_env__$private$.superseded_table_sets[[1]], + old_set + ) + expect_true(old_set$executor_built()) + qc$cleanup() + expect_length(qc$.__enclos_env__$private$.superseded_table_sets, 0) + }) + + it("errors when replacing an existing table", { + skip_if_no_dataframe_engine() + qc <- local_querychat(new_users_df(), "users", greeting = "hi") + qc$.__enclos_env__$private$.sessions_started <- TRUE + + expect_error( + qc$add_table(new_test_df(), "users", replace = TRUE), + "replace or remove" + ) + expect_length(qc$.__enclos_env__$private$.superseded_table_sets, 0) + }) + + it("errors when removing a table", { + skip_if_not_installed("duckdb") + qc <- local_querychat(new_users_df(), "users", greeting = "hi") + qc$add_table(new_test_df(), "other") + qc$.__enclos_env__$private$.sessions_started <- TRUE + + expect_error(qc$remove_table("other"), "replace or remove") + expect_equal(qc$table_names(), c("users", "other")) + }) + + it("closes a replaced source immediately before any session starts", { + skip_if_no_dataframe_engine() + qc <- local_querychat(new_users_df(), "users", greeting = "hi") + old_source <- qc_data_source(qc, "users") + + qc$add_table(new_test_df(), "users", replace = TRUE) + + expect_false(DBI::dbIsValid(old_source$.__enclos_env__$private$conn)) + }) +}) + describe("QueryChatGreeter", { skip_if_no_dataframe_engine() diff --git a/pkg-r/tests/testthat/test-querychat_module.R b/pkg-r/tests/testthat/test-querychat_module.R index 87ef47616..df2f6c03b 100644 --- a/pkg-r/tests/testthat/test-querychat_module.R +++ b/pkg-r/tests/testthat/test-querychat_module.R @@ -20,8 +20,6 @@ test_that("mod_server() return includes table() and table_names() for single-tab skip_if_no_dataframe_engine() ds <- local_data_frame_source(new_test_df()) - executor <- build_query_executor(list(test_table = ds)) - withr::defer(executor$cleanup()) client_factory <- function(...) { structure(list(), class = c("MockChat", "Chat")) @@ -37,8 +35,7 @@ test_that("mod_server() return includes table() and table_names() for single-tab mod_server, args = list( id = "test", - data_sources = list(test_table = ds), - executor = executor, + table_set = local_table_set(list(test_table = ds)), greeting = "Hello", client = client_factory, tools = "query", @@ -79,8 +76,6 @@ test_that("mod_server() return includes table() and table_names() for multi-tabl ds1 <- local_data_frame_source(new_test_df(), table_name = "tbl_a") ds2 <- local_data_frame_source(new_test_df(), table_name = "tbl_b") data_sources <- list(tbl_a = ds1, tbl_b = ds2) - executor <- build_query_executor(data_sources) - withr::defer(executor$cleanup()) result <- NULL client_factory <- function(...) { @@ -98,8 +93,7 @@ test_that("mod_server() return includes table() and table_names() for multi-tabl mod_server, args = list( id = "test", - data_sources = data_sources, - executor = executor, + table_set = local_table_set(data_sources), greeting = "Hello", client = client_factory, tools = "query", @@ -133,8 +127,6 @@ test_that("mod_server() passes visualize callback and tools to client factory", skip_if_no_dataframe_engine() ds <- local_data_frame_source(new_test_df()) - executor <- build_query_executor(list(test_table = ds)) - withr::defer(executor$cleanup()) captured <- NULL client_factory <- function(...) { @@ -152,8 +144,7 @@ test_that("mod_server() passes visualize callback and tools to client factory", mod_server, args = list( id = "test", - data_sources = list(test_table = ds), - executor = executor, + table_set = local_table_set(list(test_table = ds)), greeting = "Hello", client = client_factory, tools = c("query", "visualize"), @@ -173,8 +164,6 @@ test_that("mod_server() exposes current_table() starting as NULL", { skip_if_no_dataframe_engine() ds <- local_data_frame_source(new_test_df()) - executor <- build_query_executor(list(test_table = ds)) - withr::defer(executor$cleanup()) client_factory <- function(...) { structure(list(), class = c("MockChat", "Chat")) @@ -190,8 +179,7 @@ test_that("mod_server() exposes current_table() starting as NULL", { mod_server, args = list( id = "test", - data_sources = list(test_table = ds), - executor = executor, + table_set = local_table_set(list(test_table = ds)), greeting = "Hello", client = client_factory, tools = "query", @@ -210,8 +198,6 @@ test_that("mod_server() current_table() updates on update_dashboard and reset_qu ds1 <- local_data_frame_source(new_test_df(), table_name = "tbl_a") ds2 <- local_data_frame_source(new_test_df(), table_name = "tbl_b") data_sources <- list(tbl_a = ds1, tbl_b = ds2) - executor <- build_query_executor(data_sources) - withr::defer(executor$cleanup()) captured_callbacks <- NULL client_factory <- function(...) { @@ -229,8 +215,7 @@ test_that("mod_server() current_table() updates on update_dashboard and reset_qu mod_server, args = list( id = "test", - data_sources = data_sources, - executor = executor, + table_set = local_table_set(data_sources), greeting = "Hello", client = client_factory, tools = "query", @@ -362,8 +347,8 @@ describe("mod_server() handoff startup", { it("builds a handoff-aware session client and starts after chat_server", { skip_if_no_dataframe_engine() ds <- local_data_frame_source(new_test_df(), engine = "sqlite") - executor <- build_query_executor(list(test_table = ds)) - withr::defer(executor$cleanup()) + table_set <- local_table_set(list(test_table = ds)) + executor <- table_set$executor() events <- character() captured_client_args <- NULL captured_handoff_args <- NULL @@ -396,8 +381,7 @@ describe("mod_server() handoff startup", { mod_server, args = list( id = "test", - data_sources = list(test_table = ds), - executor = executor, + table_set = table_set, greeting = "Hello", client = client_factory, tools = "query", @@ -431,8 +415,6 @@ test_that("restored viz widgets survive a second bookmark cycle", { skip_if_no_dataframe_engine() ds <- local_data_frame_source(new_test_df()) - executor <- build_query_executor(list(test_table = ds)) - withr::defer(executor$cleanup()) callbacks <- NULL bookmark_fn <- NULL restore_fn <- NULL @@ -473,8 +455,7 @@ test_that("restored viz widgets survive a second bookmark cycle", { mod_server, args = list( id = "test", - data_sources = list(test_table = ds), - executor = executor, + table_set = local_table_set(list(test_table = ds)), greeting = "Hello", client = client_factory, tools = c("query", "visualize"), @@ -517,8 +498,6 @@ test_that("onBookmark callback mutates environment-backed state$values", { skip_if_no_dataframe_engine() ds <- local_data_frame_source(new_test_df()) - executor <- build_query_executor(list(test_table = ds)) - withr::defer(executor$cleanup()) bookmark_fn <- NULL client_factory <- function(...) { @@ -542,8 +521,7 @@ test_that("onBookmark callback mutates environment-backed state$values", { mod_server, args = list( id = "test", - data_sources = list(test_table = ds), - executor = executor, + table_set = local_table_set(list(test_table = ds)), greeting = "Hello", client = client_factory, tools = "query", @@ -564,8 +542,6 @@ test_that("mod_server() calls chat_server('chat', ...) with the pre-built client skip_if_no_dataframe_engine() ds <- local_data_frame_source(new_test_df()) - executor <- build_query_executor(list(test_table = ds)) - withr::defer(executor$cleanup()) captured_chat_args <- NULL client_factory <- function(...) { @@ -585,8 +561,7 @@ test_that("mod_server() calls chat_server('chat', ...) with the pre-built client mod_server, args = list( id = "test", - data_sources = list(test_table = ds), - executor = executor, + table_set = local_table_set(list(test_table = ds)), greeting = "Hello", client = client_factory, tools = "query", @@ -603,8 +578,6 @@ test_that("mod_server() calls chat_restore() with the auto-bookmark trigger disa skip_if_no_dataframe_engine() ds <- local_data_frame_source(new_test_df()) - executor <- build_query_executor(list(test_table = ds)) - withr::defer(executor$cleanup()) client_factory <- function(...) { structure(list(), class = c("MockChat", "Chat")) @@ -624,8 +597,7 @@ test_that("mod_server() calls chat_restore() with the auto-bookmark trigger disa mod_server, args = list( id = "test", - data_sources = list(test_table = ds), - executor = executor, + table_set = local_table_set(list(test_table = ds)), greeting = "Hello", client = client_factory, tools = "query", @@ -647,8 +619,6 @@ test_that("mod_server() skips chat_restore() when history is bookmark mode", { skip_if_no_dataframe_engine() ds <- local_data_frame_source(new_test_df()) - executor <- build_query_executor(list(test_table = ds)) - withr::defer(executor$cleanup()) client_factory <- function(...) { structure(list(), class = c("MockChat", "Chat")) @@ -668,8 +638,7 @@ test_that("mod_server() skips chat_restore() when history is bookmark mode", { mod_server, args = list( id = "test", - data_sources = list(test_table = ds), - executor = executor, + table_set = local_table_set(list(test_table = ds)), greeting = "Hello", client = client_factory, tools = "query", @@ -685,8 +654,6 @@ test_that("mod_server() builds the auto-generated greeting from the greeter, not skip_if_no_dataframe_engine() ds <- local_data_frame_source(new_test_df()) - executor <- build_query_executor(list(test_table = ds)) - withr::defer(executor$cleanup()) main_client_calls <- list() client_factory <- function(...) { @@ -704,7 +671,7 @@ test_that("mod_server() builds the auto-generated greeting from the greeter, not build_client_calls <- list() fake_greeter <- list( - build_client = function(base = NULL) { + build_client = function(base = NULL, tables = NULL, table_set = NULL) { build_client_calls[[length(build_client_calls) + 1L]] <<- base fake_greeting_client } @@ -725,8 +692,7 @@ test_that("mod_server() builds the auto-generated greeting from the greeter, not mod_server, args = list( id = "test", - data_sources = list(test_table = ds), - executor = executor, + table_set = local_table_set(list(test_table = ds)), greeting = NULL, client = client_factory, tools = "query", @@ -753,8 +719,6 @@ test_that("mod_server() chat_update input updates table state", { skip_if_no_dataframe_engine() ds <- local_data_frame_source(new_test_df()) - executor <- build_query_executor(list(test_table = ds)) - withr::defer(executor$cleanup()) client_factory <- function(...) { structure(list(), class = c("MockChat", "Chat")) @@ -770,8 +734,7 @@ test_that("mod_server() chat_update input updates table state", { mod_server, args = list( id = "test", - data_sources = list(test_table = ds), - executor = executor, + table_set = local_table_set(list(test_table = ds)), greeting = "Hello", client = client_factory, tools = "query", @@ -801,8 +764,6 @@ test_that("mod_server() registers table/viz state with both bookmark and history skip_if_no_dataframe_engine() ds <- local_data_frame_source(new_test_df()) - executor <- build_query_executor(list(test_table = ds)) - withr::defer(executor$cleanup()) client_factory <- function(...) { structure(list(), class = c("MockChat", "Chat")) @@ -843,8 +804,7 @@ test_that("mod_server() registers table/viz state with both bookmark and history mod_server, args = list( id = "test", - data_sources = list(test_table = ds), - executor = executor, + table_set = local_table_set(list(test_table = ds)), greeting = "Hello", client = client_factory, tools = "query", @@ -863,8 +823,6 @@ test_that("history on_save callback returns merged values (R history contract)", skip_if_no_dataframe_engine() ds <- local_data_frame_source(new_test_df()) - executor <- build_query_executor(list(test_table = ds)) - withr::defer(executor$cleanup()) client_factory <- function(...) { structure(list(), class = c("MockChat", "Chat")) @@ -888,8 +846,7 @@ test_that("history on_save callback returns merged values (R history contract)", mod_server, args = list( id = "test", - data_sources = list(test_table = ds), - executor = executor, + table_set = local_table_set(list(test_table = ds)), greeting = "Hello", client = client_factory, tools = "query", @@ -924,8 +881,6 @@ test_that("history on_save callback works with no active reactive context", { skip_if_no_dataframe_engine() ds <- local_data_frame_source(new_test_df()) - executor <- build_query_executor(list(test_table = ds)) - withr::defer(executor$cleanup()) client_factory <- function(...) { structure(list(), class = c("MockChat", "Chat")) @@ -949,8 +904,7 @@ test_that("history on_save callback works with no active reactive context", { mod_server, args = list( id = "test", - data_sources = list(test_table = ds), - executor = executor, + table_set = local_table_set(list(test_table = ds)), greeting = "Hello", client = client_factory, tools = "query", diff --git a/pkg-r/tests/testthat/test-server_data_source.R b/pkg-r/tests/testthat/test-server_data_source.R new file mode 100644 index 000000000..c12e6b381 --- /dev/null +++ b/pkg-r/tests/testthat/test-server_data_source.R @@ -0,0 +1,235 @@ +# QueryChat$server(data_source = ) registers a table for one session only. +# mod_server() is mocked so each test can inspect what a session received. + +local_captured_mod_server <- function(env = parent.frame()) { + calls <- new.env(parent = emptyenv()) + calls$args <- list() + testthat::local_mocked_bindings( + mod_server = function(id, ...) { + calls$args[[length(calls$args) + 1L]] <- list(...) + list() + }, + .package = "querychat", + .env = env + ) + calls +} + +# Runs qc$server(...) under an explicit MockShinySession and returns that +# session still open, so a test can hold several sessions at once and end +# each with session$close() (which fires onSessionEnded callbacks). +# shiny::testServer() is not used here because it closes its session on exit. +start_server_session <- function(qc, ..., env = parent.frame()) { + args <- list(...) + session <- shiny::MockShinySession$new() + withr::defer(if (!session$isClosed()) session$close(), envir = env) + shiny::withReactiveDomain(session, do.call(qc$server, args)) + session +} + +source_conn_valid <- function(source) { + DBI::dbIsValid(source$.__enclos_env__$private$conn) +} + +describe("QueryChat$server(data_source = ) session isolation", { + it("does not change the instance's tables", { + skip_if_no_dataframe_engine() + withr::local_envvar(OPENAI_API_KEY = "boop") + calls <- local_captured_mod_server() + qc <- QueryChat$new(NULL, table_name = "users", greeting = "hi") + withr::defer(qc$cleanup()) + + start_server_session(qc, data_source = new_users_df()) + + expect_equal(qc$table_names(), character()) + expect_equal(calls$args[[1]]$table_set$table_names(), "users") + }) + + it("gives each session its own source", { + skip_if_no_dataframe_engine() + withr::local_envvar(OPENAI_API_KEY = "boop") + calls <- local_captured_mod_server() + qc <- QueryChat$new(NULL, table_name = "users", greeting = "hi") + withr::defer(qc$cleanup()) + + start_server_session(qc, data_source = new_users_df()) + start_server_session(qc, data_source = new_test_df()) + + first <- calls$args[[1]]$table_set$data_sources$users + second <- calls$args[[2]]$table_set$data_sources$users + expect_false(identical(first, second)) + expect_equal(names(first$get_data()), c("id", "name", "age")) + expect_equal(names(second$get_data()), c("id", "name", "value")) + }) + + it("shadows a same-named instance table for that session only", { + skip_if_no_dataframe_engine() + withr::local_envvar(OPENAI_API_KEY = "boop") + calls <- local_captured_mod_server() + qc <- QueryChat$new(new_users_df(), "users", greeting = "hi") + withr::defer(qc$cleanup()) + config_source <- qc_data_source(qc, "users") + + start_server_session(qc, data_source = new_test_df()) + + expect_identical(qc_data_source(qc, "users"), config_source) + session_source <- calls$args[[1]]$table_set$data_sources$users + expect_false(identical(session_source, config_source)) + expect_equal(names(session_source$get_data()), c("id", "name", "value")) + }) + + it("keeps sessions from seeing each other's tables", { + skip_if_not_installed("duckdb") + withr::local_envvar(OPENAI_API_KEY = "boop") + calls <- local_captured_mod_server() + qc <- QueryChat$new(NULL, greeting = "hi") + qc$add_table(new_users_df(), "orders") + withr::defer(qc$cleanup()) + + start_server_session( + qc, + data_source = new_test_df(), + table_name = "returns" + ) + start_server_session( + qc, + data_source = new_metrics_df(), + table_name = "orders" + ) + + expect_equal( + calls$args[[1]]$table_set$table_names(), + c("orders", "returns") + ) + expect_equal(calls$args[[2]]$table_set$table_names(), "orders") + }) + + it("snapshots greeting tables per session without touching the greeter", { + skip_if_no_dataframe_engine() + withr::local_envvar(OPENAI_API_KEY = "boop") + calls <- local_captured_mod_server() + qc <- QueryChat$new(NULL, table_name = "users", greeting = "hi") + withr::defer(qc$cleanup()) + + start_server_session(qc, data_source = new_users_df()) + + expect_equal(calls$args[[1]]$greeting_tables, "users") + expect_equal(qc$greeter$tables, character()) + }) + + it("passes the session's table set and description to the greeter factory", { + skip_if_no_dataframe_engine() + withr::local_envvar(OPENAI_API_KEY = "boop") + calls <- local_captured_mod_server() + qc <- QueryChat$new(NULL, table_name = "users", greeting = "hi") + withr::defer(qc$cleanup()) + start_server_session(qc, data_source = new_users_df()) + session_set <- calls$args[[1]]$table_set + + seen <- NULL + qc$greeter$.__enclos_env__$private$.client_factory <- function( + tables, + prompt, + base = NULL, + table_set = NULL + ) { + seen <<- list(tables = tables, table_set = table_set) + structure(list(), class = c("MockChat", "Chat")) + } + qc$greeter$build_client(tables = "users", table_set = session_set) + + expect_equal(seen$tables, "users") + expect_identical(seen$table_set, session_set) + }) +}) + +describe("QueryChat$server(data_source = ) session cleanup", { + it("closes the session's own source and executor when the session ends", { + skip_if_no_dataframe_engine() + withr::local_envvar(OPENAI_API_KEY = "boop") + calls <- local_captured_mod_server() + qc <- QueryChat$new(NULL, table_name = "users", greeting = "hi") + withr::defer(qc$cleanup()) + + session <- start_server_session(qc, data_source = new_users_df()) + session_set <- calls$args[[1]]$table_set + session_source <- session_set$data_sources$users + session_set$executor() + expect_true(source_conn_valid(session_source)) + + session$close() + + expect_false(source_conn_valid(session_source)) + }) + + it("does not close another session's source", { + skip_if_no_dataframe_engine() + withr::local_envvar(OPENAI_API_KEY = "boop") + calls <- local_captured_mod_server() + qc <- QueryChat$new(NULL, table_name = "users", greeting = "hi") + withr::defer(qc$cleanup()) + + session_a <- start_server_session(qc, data_source = new_users_df()) + session_b <- start_server_session(qc, data_source = new_test_df()) + source_a <- calls$args[[1]]$table_set$data_sources$users + source_b <- calls$args[[2]]$table_set$data_sources$users + + session_b$close() + expect_false(source_conn_valid(source_b)) + expect_true(source_conn_valid(source_a)) + + session_a$close() + expect_false(source_conn_valid(source_a)) + }) + + it("leaves the instance's config-time source open until $cleanup()", { + skip_if_no_dataframe_engine() + withr::local_envvar(OPENAI_API_KEY = "boop") + local_captured_mod_server() + qc <- QueryChat$new(new_users_df(), "users", greeting = "hi") + config_source <- qc_data_source(qc, "users") + + session <- start_server_session(qc, data_source = new_test_df()) + session$close() + expect_true(source_conn_valid(config_source)) + + qc$cleanup() + expect_false(source_conn_valid(config_source)) + }) + + it("owns nothing when no data_source is passed", { + skip_if_no_dataframe_engine() + withr::local_envvar(OPENAI_API_KEY = "boop") + local_captured_mod_server() + qc <- QueryChat$new(new_users_df(), "users", greeting = "hi") + withr::defer(qc$cleanup()) + config_source <- qc_data_source(qc, "users") + + session <- start_server_session(qc) + session$close() + + expect_true(source_conn_valid(config_source)) + }) + + it("closes the normalized source and leaves the instance untouched when registration fails", { + skip_if_no_dataframe_engine() + skip_if_not_installed("RSQLite") + withr::local_envvar(OPENAI_API_KEY = "boop") + calls <- local_captured_mod_server() + db <- local_sqlite_connection(table_name = "other") + qc <- QueryChat$new(new_users_df(), "users", greeting = "hi") + withr::defer(qc$cleanup()) + + expect_error( + start_server_session(qc, data_source = db$conn, table_name = "other"), + "same type" + ) + + expect_equal(qc$table_names(), "users") + start_server_session(qc) + expect_equal( + calls$args[[length(calls$args)]]$table_set$table_names(), + "users" + ) + }) +}) From 28a0e87225016d4ee5388a1b47a129acaafcea47 Mon Sep 17 00:00:00 2001 From: Carson Date: Sat, 12 Sep 2026 18:49:09 -0500 Subject: [PATCH 08/35] fix(r): don't close a caller-supplied DataSource when its session ends $server(data_source = )'s onSessionEnded() cleanup unconditionally closed session_source, even when it was a DataSource the caller constructed and passed in (rather than a raw connection/data.frame querychat normalized itself). That violated the "querychat closes only what it created" rule and could disconnect a connection the caller still owned/reused. Guard the cleanup with the same !inherits(data_source, "DataSource") check the error-path cleanup already uses. --- pkg-r/R/QueryChat.R | 14 ++++++++++---- pkg-r/tests/testthat/test-server_data_source.R | 14 ++++++++++++++ 2 files changed, 24 insertions(+), 4 deletions(-) diff --git a/pkg-r/R/QueryChat.R b/pkg-r/R/QueryChat.R index 582a0afe1..7a125483c 100644 --- a/pkg-r/R/QueryChat.R +++ b/pkg-r/R/QueryChat.R @@ -1209,15 +1209,21 @@ QueryChat <- R6::R6Class( if (!is.null(session_source)) { session_table_set <- table_set + # Mirrors the tryCatch() guard above: querychat only owns (and so + # only closes) sources it normalized itself from a raw connection or + # data.frame, never a DataSource the caller constructed and passed in. + owns_session_source <- !inherits(data_source, "DataSource") session$onSessionEnded(function() { warn_on_cleanup_failure( session_table_set$cleanup_executor(), "session query executor" ) - warn_on_cleanup_failure( - session_source$cleanup(), - "session data source" - ) + if (owns_session_source) { + warn_on_cleanup_failure( + session_source$cleanup(), + "session data source" + ) + } }) } diff --git a/pkg-r/tests/testthat/test-server_data_source.R b/pkg-r/tests/testthat/test-server_data_source.R index c12e6b381..26fffe8a6 100644 --- a/pkg-r/tests/testthat/test-server_data_source.R +++ b/pkg-r/tests/testthat/test-server_data_source.R @@ -211,6 +211,20 @@ describe("QueryChat$server(data_source = ) session cleanup", { expect_true(source_conn_valid(config_source)) }) + it("does not close a caller-supplied DataSource when the session ends", { + skip_if_no_dataframe_engine() + withr::local_envvar(OPENAI_API_KEY = "boop") + local_captured_mod_server() + qc <- QueryChat$new(NULL, table_name = "users", greeting = "hi") + withr::defer(qc$cleanup()) + caller_source <- local_data_frame_source(new_users_df(), "users") + + session <- start_server_session(qc, data_source = caller_source) + session$close() + + expect_true(source_conn_valid(caller_source)) + }) + it("closes the normalized source and leaves the instance untouched when registration fails", { skip_if_no_dataframe_engine() skip_if_not_installed("RSQLite") From f58f7557798530f3e20cf85d6d494fb9f92315d1 Mon Sep 17 00:00:00 2001 From: Carson Date: Sat, 12 Sep 2026 18:49:14 -0500 Subject: [PATCH 09/35] refactor(py): track one owned base client; session overrides close on session end --- pkg-py/src/querychat/_querychat_base.py | 43 +++--- pkg-py/src/querychat/_shiny.py | 12 +- pkg-py/tests/test_cleanup.py | 170 ++++++++---------------- 3 files changed, 78 insertions(+), 147 deletions(-) diff --git a/pkg-py/src/querychat/_querychat_base.py b/pkg-py/src/querychat/_querychat_base.py index 532ea6bf9..4fc9993cf 100644 --- a/pkg-py/src/querychat/_querychat_base.py +++ b/pkg-py/src/querychat/_querychat_base.py @@ -53,6 +53,8 @@ from pins.boards import BaseBoard from shinychat.types import HistoryOptions + from shiny import Session + from ._data_dict import DataDict from ._viz_tools import VisualizeData @@ -107,17 +109,15 @@ def __init__( self._extra_instructions = extra_instructions self._categorical_threshold = categorical_threshold - # Clients querychat materializes from a spec (constructor string spec, - # deferred default resolution, .server() overrides) are tracked so - # cleanup() can close them; user-supplied Chat instances are never - # tracked -- their lifecycle remains the caller's responsibility. - self._owned_clients: list[chatlas.Chat] = [] + # Owned iff querychat resolved it from a spec; a user-supplied Chat is + # never closed by cleanup(). self._base_client: chatlas.Chat | None if isinstance(client, str): self._base_client = resolve_client(client) - self._owned_clients.append(self._base_client) + self._base_client_owned = True else: self._base_client = client + self._base_client_owned = False self._client_console = None self._greeter: QueryChatGreeter | None = None @@ -232,32 +232,26 @@ def _create_client(self, base: chatlas.Chat | None = None) -> chatlas.Chat: if base is None: if self._base_client is None: self._base_client = resolve_client(None) - self._owned_clients.append(self._base_client) + self._base_client_owned = True base = self._base_client return create_client(base) - def _resolve_override_client( - self, client: str | chatlas.Chat | None - ) -> chatlas.Chat: + def _resolve_session_client( + self, client: str | chatlas.Chat | MISSING_TYPE, session: Session + ) -> chatlas.Chat | None: """ - Resolve a per-call client override (e.g., ``.server(client=...)``). + Resolve a ``.server(client=...)`` override. - Like the constructor's ``client``, a spec-resolved override is - querychat-created and must be closed on cleanup(); a user-supplied - Chat instance is not. + A spec-resolved override is owned by the session and closed on + ``session.on_ended``; a user-supplied Chat is returned untouched. """ + if isinstance(client, MISSING_TYPE): + return None resolved = resolve_client(client) if not isinstance(client, chatlas.Chat): - self._owned_clients.append(resolved) + session.on_ended(resolved.close) return resolved - def _close_owned_client(self, client: chatlas.Chat) -> None: - """Close an owned client and stop tracking it. Idempotent.""" - try: - client.close() - finally: - self._owned_clients[:] = [c for c in self._owned_clients if c is not client] - def _create_session_client( self, table_set: TableSet, @@ -681,9 +675,8 @@ def cleanup(self) -> None: warn_on_failure(self._table_set.cleanup_executor, "query executor") for source in self._table_set.data_sources.values(): warn_on_failure(source.cleanup, "data source") - for client in self._owned_clients: - warn_on_failure(client.close, "chatlas client") - self._owned_clients.clear() + if self._base_client is not None and self._base_client_owned: + warn_on_failure(self._base_client.close, "chatlas client") def normalize_data_source( diff --git a/pkg-py/src/querychat/_shiny.py b/pkg-py/src/querychat/_shiny.py index 89f6e0d02..0dcf5fe4d 100644 --- a/pkg-py/src/querychat/_shiny.py +++ b/pkg-py/src/querychat/_shiny.py @@ -3,7 +3,6 @@ import warnings from typing import TYPE_CHECKING, Any, Literal, Optional, overload -import chatlas from htmltools import TagChild, tags from narwhals.stable.v1.typing import IntoDataFrameT, IntoFrameT, IntoLazyFrameT from shiny.express._stub_session import ExpressStubSession @@ -39,6 +38,7 @@ if TYPE_CHECKING: from pathlib import Path + import chatlas import ibis import narwhals.stable.v1 as nw import sqlalchemy @@ -640,7 +640,7 @@ def page(self, title, *, id: Optional[str] = None, **kwargs): **kwargs, ) - def server( # noqa: PLR0912 + def server( self, *, data_source: IntoFrame | sqlalchemy.Engine | ibis.Table | None = None, @@ -745,13 +745,7 @@ def server( # noqa: PLR0912 if table_set is None: table_set = self._require_table_set("server") - resolved_client: chatlas.Chat | None = ( - None - if isinstance(client, MISSING_TYPE) - else self._resolve_override_client(client) - ) - if resolved_client is not None and not isinstance(client, chatlas.Chat): - session.on_ended(lambda: self._close_owned_client(resolved_client)) + resolved_client = self._resolve_session_client(client, session) if session_source is not None: session_set, owned_source = table_set, session_source diff --git a/pkg-py/tests/test_cleanup.py b/pkg-py/tests/test_cleanup.py index 19c3fe10e..d32101f76 100644 --- a/pkg-py/tests/test_cleanup.py +++ b/pkg-py/tests/test_cleanup.py @@ -22,120 +22,45 @@ def sample_df(): class TestClientOwnership: - """querychat closes the chatlas client only if it created it.""" - - def test_ownership_registration(self, monkeypatch, sample_df): - monkeypatch.setenv("OPENAI_API_KEY", "sk-dummy-key-for-testing") - - # String spec: querychat-created, tracked at construction - qc = QueryChatBase(sample_df, "users", client="openai") - assert qc._owned_clients == [qc._base_client] - # None (deferred default/env): tracked once resolved, not before - assert QueryChatBase(sample_df, "users")._owned_clients == [] - # User-supplied instance: never tracked - assert ( - QueryChatBase(sample_df, "users", client=ChatOpenAI())._owned_clients == [] - ) - - def test_cleanup_closes_owned_string_client(self, monkeypatch, sample_df): + def test_string_spec_client_is_owned(self, monkeypatch, sample_df): monkeypatch.setenv("OPENAI_API_KEY", "sk-dummy-key-for-testing") qc = QueryChatBase(sample_df, "users", client="openai") - assert isinstance(qc._base_client, chatlas.Chat) - qc.cleanup() - assert qc._base_client.provider._client.is_closed() + assert qc._base_client_owned is True - def test_cleanup_closes_deferred_default_client(self, monkeypatch, sample_df): + def test_user_supplied_client_is_not_owned(self, sample_df): + qc = QueryChatBase(sample_df, "users", client=ChatOpenAI(api_key="sk-x")) + assert qc._base_client_owned is False + + def test_deferred_default_client_is_owned(self, monkeypatch, sample_df): monkeypatch.setenv("OPENAI_API_KEY", "sk-dummy-key-for-testing") - monkeypatch.delenv("QUERYCHAT_CLIENT", raising=False) qc = QueryChatBase(sample_df, "users") assert qc._base_client is None - # Trigger deferred resolution (env var / "openai" default) - qc._create_client() - assert isinstance(qc._base_client, chatlas.Chat) + qc.client() + assert qc._base_client is not None + assert qc._base_client_owned is True + + def test_cleanup_closes_owned_client(self, monkeypatch, sample_df): + monkeypatch.setenv("OPENAI_API_KEY", "sk-dummy-key-for-testing") + qc = QueryChatBase(sample_df, "users", client="openai") qc.cleanup() assert qc._base_client.provider._client.is_closed() - def test_cleanup_does_not_close_user_supplied_client(self, monkeypatch, sample_df): - monkeypatch.setenv("OPENAI_API_KEY", "sk-dummy-key-for-testing") - chat = ChatOpenAI() + def test_cleanup_does_not_close_user_supplied_client(self, sample_df): + chat = ChatOpenAI(api_key="sk-x") qc = QueryChatBase(sample_df, "users", client=chat) qc.cleanup() assert not chat.provider._client.is_closed() def test_cleanup_closes_clones_via_shared_provider(self, monkeypatch, sample_df): - """ - Session/console clones share the base provider (deepcopy by - reference), so closing the owned base client covers them. - """ monkeypatch.setenv("OPENAI_API_KEY", "sk-dummy-key-for-testing") qc = QueryChatBase(sample_df, "users", client="openai") - clone = qc._create_client() - assert clone.provider is qc._base_client.provider + clone = qc.client() qc.cleanup() assert clone.provider._client.is_closed() class TestServerClientOverrides: - """ - Clients resolved for .server(client=...) overrides follow the same - ownership rule: spec-resolved overrides are closed, user-supplied ones - are not. - """ - - def test_owned_override_closed_when_base_is_user_supplied( - self, monkeypatch, sample_df - ): - monkeypatch.setenv("OPENAI_API_KEY", "sk-dummy-key-for-testing") - user_chat = ChatOpenAI() - qc = QueryChatBase(sample_df, "users", client=user_chat) - override = qc._resolve_override_client("openai") - qc.cleanup() - assert override.provider._client.is_closed() - assert not user_chat.provider._client.is_closed() - - def test_owned_override_closed_when_base_deferred(self, monkeypatch, sample_df): - monkeypatch.setenv("OPENAI_API_KEY", "sk-dummy-key-for-testing") - qc = QueryChatBase(sample_df, "users") - override = qc._resolve_override_client("openai") - assert qc._base_client is None - qc.cleanup() - assert override.provider._client.is_closed() - - def test_deferred_default_override_closed(self, monkeypatch, sample_df): - monkeypatch.setenv("OPENAI_API_KEY", "sk-dummy-key-for-testing") - monkeypatch.delenv("QUERYCHAT_CLIENT", raising=False) - qc = QueryChatBase(sample_df, "users", client="openai") - override = qc._resolve_override_client(None) - qc.cleanup() - assert override.provider._client.is_closed() - - def test_user_supplied_override_not_closed(self, monkeypatch, sample_df): - monkeypatch.setenv("OPENAI_API_KEY", "sk-dummy-key-for-testing") - qc = QueryChatBase(sample_df, "users", client="openai") - override = ChatOpenAI() - qc._resolve_override_client(override) - qc.cleanup() - assert not override.provider._client.is_closed() - - def test_close_owned_client_closes_and_untracks(self, monkeypatch, sample_df): - monkeypatch.setenv("OPENAI_API_KEY", "sk-dummy-key-for-testing") - qc = QueryChatBase(sample_df, "users", client="openai") - override = qc._resolve_override_client("openai") - qc._close_owned_client(override) - assert override.provider._client.is_closed() - assert all(c is not override for c in qc._owned_clients) - qc.cleanup() # already untracked: no double-close - - def test_close_owned_client_is_idempotent(self, monkeypatch, sample_df): - monkeypatch.setenv("OPENAI_API_KEY", "sk-dummy-key-for-testing") - qc = QueryChatBase(sample_df, "users", client="openai") - override = qc._resolve_override_client("openai") - qc._close_owned_client(override) - qc._close_owned_client(override) # should not raise - - -class TestServerSessionEndClosing: - """.server() closes owned spec-resolved overrides when the session ends.""" + """.server(client=...) overrides belong to the session, not the instance.""" @pytest.fixture def ended_callbacks(self, monkeypatch): @@ -146,25 +71,39 @@ def ended_callbacks(self, monkeypatch): monkeypatch.setattr(shiny_mod, "mod_server", lambda *args, **kwargs: None) return callbacks - def test_owned_override_closed_on_session_end( - self, monkeypatch, sample_df, ended_callbacks + @pytest.fixture + def resolved_clients(self, monkeypatch): + import querychat._querychat_base as base_mod + + created: list[chatlas.Chat] = [] + real = base_mod.resolve_client + + def spy(spec): + chat = real(spec) + created.append(chat) + return chat + + monkeypatch.setattr(base_mod, "resolve_client", spy) + return created + + def test_spec_override_closed_on_session_end( + self, monkeypatch, sample_df, ended_callbacks, resolved_clients ): monkeypatch.setenv("OPENAI_API_KEY", "sk-dummy-key-for-testing") qc = shiny_mod.QueryChat(sample_df, "users") qc.server(client="openai") - (override,) = qc._owned_clients + (override,) = resolved_clients + assert not override.provider._client.is_closed() for cb in ended_callbacks: cb() assert override.provider._client.is_closed() - assert qc._owned_clients == [] def test_user_supplied_override_not_closed_on_session_end( - self, monkeypatch, sample_df, ended_callbacks + self, sample_df, ended_callbacks ): - monkeypatch.setenv("OPENAI_API_KEY", "sk-dummy-key-for-testing") - qc = shiny_mod.QueryChat(sample_df, "users") - chat = ChatOpenAI() + qc = shiny_mod.QueryChat(sample_df, "users", client=ChatOpenAI(api_key="sk-x")) + chat = ChatOpenAI(api_key="sk-x") qc.server(client=chat) for cb in ended_callbacks: @@ -172,17 +111,27 @@ def test_user_supplied_override_not_closed_on_session_end( qc.cleanup() assert not chat.provider._client.is_closed() - def test_owned_override_tracked_and_closed_by_cleanup( - self, monkeypatch, sample_df, ended_callbacks + def test_cleanup_does_not_close_live_session_override( + self, monkeypatch, sample_df, ended_callbacks, resolved_clients ): monkeypatch.setenv("OPENAI_API_KEY", "sk-dummy-key-for-testing") - qc = shiny_mod.QueryChat(sample_df, "users") + qc = shiny_mod.QueryChat(sample_df, "users", client=ChatOpenAI(api_key="sk-x")) qc.server(client="openai") - # Session still alive: override stays tracked and cleanup() closes it - (override,) = qc._owned_clients + (override,) = resolved_clients qc.cleanup() - assert override.provider._client.is_closed() + assert not override.provider._client.is_closed() + + def test_session_override_does_not_leak_into_instance( + self, monkeypatch, sample_df, ended_callbacks + ): + monkeypatch.setenv("OPENAI_API_KEY", "sk-dummy-key-for-testing") + base = ChatOpenAI(api_key="sk-x") + qc = shiny_mod.QueryChat(sample_df, "users", client=base) + qc.server(client="openai") + + assert qc._base_client is base + assert qc._base_client_owned is False class TestCleanupDataSources: @@ -202,16 +151,11 @@ def test_cleanup_is_idempotent(self, monkeypatch, sample_df): qc.cleanup() qc.cleanup() # should not raise - def test_cleanup_closes_remaining_clients_after_close_failure( - self, monkeypatch, sample_df - ): + def test_cleanup_warns_on_client_close_failure(self, monkeypatch, sample_df): monkeypatch.setenv("OPENAI_API_KEY", "sk-dummy-key-for-testing") qc = QueryChatBase(sample_df, "users", client="openai") - override = qc._resolve_override_client("openai") with ( patch.object(qc._base_client, "close", side_effect=RuntimeError("boom")), pytest.warns(UserWarning, match="Failed to clean up chatlas client"), ): qc.cleanup() - assert override.provider._client.is_closed() - assert qc._owned_clients == [] From 9ac54f65c772c77ab28a6d6245d067f25bf4863e Mon Sep 17 00:00:00 2001 From: Carson Date: Sat, 12 Sep 2026 18:51:56 -0500 Subject: [PATCH 10/35] docs(r): describe session-scoped $server(data_source = ) and the cleanup ownership rule --- pkg-r/NEWS.md | 6 ++++++ pkg-r/vignettes/build.Rmd | 2 ++ 2 files changed, 8 insertions(+) diff --git a/pkg-r/NEWS.md b/pkg-r/NEWS.md index 8f7f9f83e..e2fb09f37 100644 --- a/pkg-r/NEWS.md +++ b/pkg-r/NEWS.md @@ -72,6 +72,12 @@ * `$add_table()` no longer rewrites the Shiny module `$id` when the registered table is the only one; the id is now fixed at construction time, matching Python. The rewrite could desync the module namespace from an already-rendered UI when a table was registered between `$ui()` and `$server()` (e.g. via `$server(data_source = )`). (#305) +* `$server(data_source = )` no longer modifies the `QueryChat` instance. The table is registered for that session only: the instance's tables, greeting tables, and system prompt are unchanged, a same-named instance table is shadowed for that session, and any connection querychat created for it is cleaned up when the session ends. A second session's `$server(data_source = )` call therefore no longer errors with "Cannot add tables after server initialization." (#300, #306) + +* `$cleanup()` follows one rule: querychat closes only what it created. `DBISource$cleanup()` and `TblSqlSource$cleanup()` no longer disconnect your connection; disconnect it yourself on shutdown. `DataFrameSource`/`PinSource` DuckDB connections are still closed. + +* Adding a *new* table with `$add_table()`/`$add_tables()` after a session has started now warns instead of erroring; running sessions keep their tables and new sessions see the addition. Replacing or removing an existing table after a session has started still errors. + # querychat 0.3.0 ## New features diff --git a/pkg-r/vignettes/build.Rmd b/pkg-r/vignettes/build.Rmd index f4ef3cdf7..7866cb513 100644 --- a/pkg-r/vignettes/build.Rmd +++ b/pkg-r/vignettes/build.Rmd @@ -158,6 +158,8 @@ server <- function(input, output, session) { shinyApp(ui, server) ``` +A data source passed to `$server()` belongs to that session: other sessions never see it, it does not change the tables registered on the `QueryChat` object, and querychat cleans up any connection it created for it when the session ends. Connections you open yourself are yours to close; `session$onSessionEnded()` is a good place. + If your chat client also depends on session-scoped credentials, you can defer that too by passing it to `qc$server(client = ...)` alongside the `data_source`. This is also a useful pattern when using something like [`{pool}`](https://github.com/rstudio/pool) to efficiently manage a pool of database connections (which we strongly recommend for production apps). From 84f580f2a89a46c1b658274870be4a6686aac819 Mon Sep 17 00:00:00 2001 From: Carson Date: Sat, 12 Sep 2026 18:53:52 -0500 Subject: [PATCH 11/35] docs(py): describe session-scoped server(data_source=) and the cleanup ownership rule --- pkg-py/CHANGELOG.md | 12 ++++++++++++ pkg-py/docs/build.qmd | 2 ++ 2 files changed, 14 insertions(+) diff --git a/pkg-py/CHANGELOG.md b/pkg-py/CHANGELOG.md index ec48860b3..c6c8d1001 100644 --- a/pkg-py/CHANGELOG.md +++ b/pkg-py/CHANGELOG.md @@ -5,6 +5,18 @@ All notable changes to this project will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [Unreleased] + +### Changes + +* `.server(data_source=)` no longer modifies the `QueryChat` instance. The table is registered for that session only: the instance's tables, greeting tables, and system prompt are unchanged, a same-named instance table is shadowed for that session, and the session's data source is cleaned up when the session ends. This removes the concurrent-session edge cases that `0.8.0` patched around (#300, #302, #303, #304, #308). + +* `cleanup()` follows one rule: querychat closes only what it created. `SQLAlchemySource.cleanup()` no longer disposes your engine; dispose it yourself on shutdown. `DataFrameSource`/`PinSource` DuckDB connections and querychat-created chat clients are still closed. + +* Adding a *new* table with `add_table()`/`add_tables()` after a session has started now warns instead of raising; running sessions keep their tables and new sessions see the addition. Replacing or removing an existing table after a session has started still raises. + +* `cleanup()` no longer closes a spec-resolved `.server(client=...)` override while its session is still running; it is closed when the session ends. + ## [0.8.0] - 2026-09-12 ### New features diff --git a/pkg-py/docs/build.qmd b/pkg-py/docs/build.qmd index b61dc6bee..6071718c3 100644 --- a/pkg-py/docs/build.qmd +++ b/pkg-py/docs/build.qmd @@ -187,6 +187,8 @@ app = App(app_ui, server) If your chat client also depends on session-scoped credentials, you can defer that too by passing it to `qc.server(client=...)` alongside the `data_source`. +A data source passed to `qc.server()` belongs to that session: other sessions never see it, it does not change the tables registered on `qc`, and querychat cleans up any connection it created for it when the session ends. If you create the connection yourself (a SQLAlchemy engine, an Ibis backend), closing it is up to you; `session.on_ended` is a good place. + ::: ::: From f221546c7e61912292c639716d78d1b95ed6706a Mon Sep 17 00:00:00 2001 From: Carson Date: Sat, 12 Sep 2026 19:14:38 -0500 Subject: [PATCH 12/35] fix(py): close final-review gaps in session-local table set refactor MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fix a session resource leak where a failed .server(client=...) resolution left the session's own normalized DataSource/TableSet uncleaned (on_ended is now registered before client resolution can raise), and make the on_ended cleanup honor the "only close what querychat created" rule for caller-supplied DataSource objects. Also fixes a couple of misleading error messages (replace-by-non-last-key ordering, remove_table's not-found vs late-change error precedence), documents cleanup()'s session-resource boundary, drops a no-longer-inexpressible pyright suppression on mod_server(), asserts the _swap_table_set() no-replace invariant, and tidies up dead/weak tests. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude --- pkg-py/src/querychat/_querychat_base.py | 24 ++++++++++-- pkg-py/src/querychat/_shiny.py | 10 +++-- pkg-py/src/querychat/_shiny_module.py | 17 ++++---- pkg-py/tests/test_server_data_source.py | 52 ++++++++++++++++--------- 4 files changed, 69 insertions(+), 34 deletions(-) diff --git a/pkg-py/src/querychat/_querychat_base.py b/pkg-py/src/querychat/_querychat_base.py index 4fc9993cf..2b3bcd5c1 100644 --- a/pkg-py/src/querychat/_querychat_base.py +++ b/pkg-py/src/querychat/_querychat_base.py @@ -211,6 +211,11 @@ def _swap_table_set(self, new_set: TableSet, *, replaced: list[DataSource]) -> N # A running session may still query through old_set's executor; # cleanup() closes it. `replaced` is always empty here because # _check_late_change() rejects replacement once sessions started. + if replaced: + raise AssertionError( + "_swap_table_set() received replaced sources after sessions " + "started; _check_late_change() should have rejected this." + ) self._superseded_table_sets.append(old_set) return with contextlib.suppress(Exception): @@ -481,7 +486,11 @@ def add_table( normalized = normalize_data_source(data_source, table_name) try: self._warn_if_prompt_rebuilt_with_history() - new_set = self._build_table_set({**self._data_sources, table_name: normalized}) + merged = { + **{k: v for k, v in self._data_sources.items() if k != table_name}, + table_name: normalized, + } + new_set = self._build_table_set(merged) except Exception: if normalized is not data_source: normalized.cleanup() @@ -601,7 +610,11 @@ def normalized_builder(name: str) -> DataSource: normalized = {name: normalized_builder(name) for name in tables} self._warn_if_prompt_rebuilt_with_history() - new_set = self._build_table_set({**self._data_sources, **normalized}) + merged = { + **{k: v for k, v in self._data_sources.items() if k not in normalized}, + **normalized, + } + new_set = self._build_table_set(merged) replaced = [ old @@ -634,8 +647,6 @@ def remove_table(self, table_name: str) -> None: If called to replace or remove an existing table after a session has started. """ - self._check_late_change("remove_table", destructive=True) - if table_name not in self._data_sources: available = ", ".join(self._data_sources.keys()) raise ValueError(f"Table '{table_name}' not found. Available: {available}") @@ -645,6 +656,8 @@ def remove_table(self, table_name: str) -> None: "Cannot remove last table. At least one table is required." ) + self._check_late_change("remove_table", destructive=True) + removed = self._data_sources[table_name] remaining = {n: s for n, s in self._data_sources.items() if n != table_name} self._warn_if_prompt_rebuilt_with_history() @@ -665,6 +678,9 @@ def cleanup(self) -> None: if querychat created it from a spec (``client=None`` or a string such as ``"openai/gpt-4o"``); a ``chatlas.Chat`` you supplied is left open. + Resources a session registers via ``.server(data_source=...)`` are + released when that session ends, not here. + Safe to call multiple times. In long-lived applications, call this when the app shuts down (e.g., via ``atexit``). """ diff --git a/pkg-py/src/querychat/_shiny.py b/pkg-py/src/querychat/_shiny.py index 0dcf5fe4d..158857556 100644 --- a/pkg-py/src/querychat/_shiny.py +++ b/pkg-py/src/querychat/_shiny.py @@ -745,17 +745,19 @@ def server( if table_set is None: table_set = self._require_table_set("server") - resolved_client = self._resolve_session_client(client, session) - if session_source is not None: - session_set, owned_source = table_set, session_source + session_set = table_set + owned_source = session_source if session_source is not data_source else None def cleanup_session() -> None: warn_on_failure(session_set.cleanup_executor, "session query executor") - warn_on_failure(owned_source.cleanup, "session data source") + if owned_source is not None: + warn_on_failure(owned_source.cleanup, "session data source") session.on_ended(cleanup_session) + resolved_client = self._resolve_session_client(client, session) + def create_session_client(**kwargs) -> chatlas.Chat: return self._create_session_client(table_set, base=resolved_client, **kwargs) diff --git a/pkg-py/src/querychat/_shiny_module.py b/pkg-py/src/querychat/_shiny_module.py index 74db9d6be..3e80a588f 100644 --- a/pkg-py/src/querychat/_shiny_module.py +++ b/pkg-py/src/querychat/_shiny_module.py @@ -250,23 +250,24 @@ def mod_server( greeting_base: chatlas.Chat | None = None, greeting_tables: list[str] | None = None, # TableSet erases each DataSource's frame type for uniform executor handling, - # so IntoFrameT can't be bound from a parameter here -- only from callers' - # annotated return-type usage (e.g. QueryChat.server() -> ServerValues[IntoFrameT]). -) -> ServerValues[IntoFrameT]: # pyright: ignore[reportInvalidTypeVarUse] + # so IntoFrameT can't be bound from any parameter here. There's no way to + # express the real return type, so callers (e.g. QueryChat.server()) must + # cast/annotate the result as ServerValues[IntoFrameT] themselves. +) -> ServerValues[Any]: if not callable(client): raise TypeError("mod_server() requires a callable client factory.") - table_states: dict[str, TableState[IntoFrameT]] = {} + table_states: dict[str, TableState[Any]] = {} _current_table: ReactiveStringOrNone = ReactiveStringOrNone(None) def _make_table_state( - source: DataSource[IntoFrameT], exec: QueryExecutor - ) -> TableState[IntoFrameT]: + source: DataSource[Any], exec: QueryExecutor + ) -> TableState[Any]: table_sql = ReactiveStringOrNone(None) table_title = ReactiveStringOrNone(None) @reactive.calc - def filtered_df() -> IntoFrameT: + def filtered_df() -> Any: query = table_sql.get() if query: return exec.execute_query(query) @@ -477,7 +478,7 @@ def _on_history_restore(values: dict[str, Any]) -> None: df_warned = False @reactive.calc - def _multi_table_df() -> IntoFrameT: + def _multi_table_df() -> Any: nonlocal df_warned if not df_warned: df_warned = True diff --git a/pkg-py/tests/test_server_data_source.py b/pkg-py/tests/test_server_data_source.py index 3a197f459..49ba875de 100644 --- a/pkg-py/tests/test_server_data_source.py +++ b/pkg-py/tests/test_server_data_source.py @@ -50,24 +50,6 @@ def end(self): cb() -@pytest.fixture -def fake_sessions(monkeypatch): - """Patch mod_server/get_current_session; each server() call gets a new session.""" - sessions: list[FakeSession] = [] - - def fake_mod_server(*args, **kwargs): - return MagicMock() - - def next_session(): - session = FakeSession() - sessions.append(session) - return session - - monkeypatch.setattr(shiny_mod, "mod_server", fake_mod_server) - monkeypatch.setattr(shiny_mod, "get_current_session", next_session) - return sessions - - @pytest.fixture def session_runs(monkeypatch): """Each server() call gets a fresh FakeSession; mod_server kwargs are captured.""" @@ -259,6 +241,8 @@ def test_session_without_data_source_owns_nothing(self, users_df, session_runs): qc.server() config_source = qc._data_sources["users"] + assert sessions[0]._ended_callbacks == [] + with ( patch.object(config_source, "cleanup") as cleanup, patch.object(qc._table_set, "cleanup_executor") as cleanup_executor, @@ -310,3 +294,35 @@ def spy(data_source, table_name): qc.server() assert calls[-1]["table_set"].table_names == ["users"] + + def test_client_resolution_failure_still_cleans_up_session_source( + self, users_df, session_runs, monkeypatch + ): + """ + A `.server(client=...)` override that fails to resolve must not leak + the session's own normalized data source: on_ended cleanup must already + be registered by the time client resolution can raise. + """ + import duckdb + + sessions, _calls = session_runs + qc = shiny_mod.QueryChat(None, table_name="users") + created = [] + real_normalize = shiny_mod.normalize_data_source + + def spy(data_source, table_name): + source = real_normalize(data_source, table_name) + created.append(source) + return source + + monkeypatch.setattr(shiny_mod, "normalize_data_source", spy) + + with pytest.raises(ValueError, match="not a known chatlas provider"): + qc.server(data_source=users_df, client="not-a-real-provider") + + (session_source,) = created + assert sessions[0]._ended_callbacks != [] + + sessions[0].end() + with pytest.raises(duckdb.ConnectionException): + session_source.execute_query("SELECT 1") From 3112b36cce2a03bff18c053437fc2923e5aa35cd Mon Sep 17 00:00:00 2001 From: Carson Date: Sat, 12 Sep 2026 19:21:15 -0500 Subject: [PATCH 13/35] fix(r): PinSource$cleanup() no longer leaks its own DuckDB connection PinSource always opens its own connection (duckdb or SQLite), unlike DBISource which wraps a caller-supplied one. Now that DBISource$cleanup() is a no-op, PinSource needs its own override to actually close what it opened. --- pkg-r/R/PinSource.R | 19 +++++++++++++++++++ pkg-r/man/PinSource.Rd | 21 ++++++++++++++++++++- pkg-r/tests/testthat/test-PinSource.R | 13 +++++++++++++ 3 files changed, 52 insertions(+), 1 deletion(-) diff --git a/pkg-r/R/PinSource.R b/pkg-r/R/PinSource.R index 84bbde013..148efddc5 100644 --- a/pkg-r/R/PinSource.R +++ b/pkg-r/R/PinSource.R @@ -174,6 +174,25 @@ PinSource <- R6::R6Class( parts <- c(parts, paste("Tags:", paste(meta$tags, collapse = ", "))) } paste(parts, collapse = "\n\n") + }, + + #' @description + #' Disconnect the DuckDB or SQLite connection this PinSource opened, and + #' shut down the DuckDB instance if used. + #' + #' Unlike [DBISource]'s `cleanup()`, this isn't a no-op: PinSource always + #' opens its own connection (never a caller-supplied one), so it owns it. + #' + #' @return `NULL` (invisibly) + cleanup = function() { + if (!is.null(private$conn) && DBI::dbIsValid(private$conn)) { + if (inherits(private$conn, "duckdb_connection")) { + DBI::dbDisconnect(private$conn, shutdown = TRUE) + } else { + DBI::dbDisconnect(private$conn) + } + } + invisible(NULL) } ), private = list( diff --git a/pkg-r/man/PinSource.Rd b/pkg-r/man/PinSource.Rd index 34eacf4a9..59289dc52 100644 --- a/pkg-r/man/PinSource.Rd +++ b/pkg-r/man/PinSource.Rd @@ -67,12 +67,12 @@ if (rlang::is_installed(c("pins", "duckdb"))) { \itemize{ \item \href{#method-PinSource-initialize}{\code{PinSource$new()}} \item \href{#method-PinSource-get_data_description}{\code{PinSource$get_data_description()}} + \item \href{#method-PinSource-cleanup}{\code{PinSource$cleanup()}} \item \href{#method-PinSource-clone}{\code{PinSource$clone()}} } } \if{html}{\out{
Inherited methods
    -
  • DBISource$cleanup()
  • DBISource$execute_query()
  • DBISource$get_data()
  • DBISource$get_db_type()
    • @@ -140,6 +140,25 @@ string if none are set. } } +\if{html}{\out{
    }} +\if{html}{\out{}} +\if{latex}{\out{\hypertarget{method-PinSource-cleanup}{}}} +\subsection{\code{PinSource$cleanup()}}{ + Disconnect the DuckDB or SQLite connection this PinSource opened, and +shut down the DuckDB instance if used. + +Unlike \link{DBISource}'s \code{cleanup()}, this isn't a no-op: PinSource always +opens its own connection (never a caller-supplied one), so it owns it. + \subsection{Usage}{ + \if{html}{\out{
    }} + \preformatted{PinSource$cleanup()} + \if{html}{\out{
    }} + } + \subsection{Returns}{ + \code{NULL} (invisibly) + } +} + \if{html}{\out{
    }} \if{html}{\out{}} \if{latex}{\out{\hypertarget{method-PinSource-clone}{}}} diff --git a/pkg-r/tests/testthat/test-PinSource.R b/pkg-r/tests/testthat/test-PinSource.R index 2230dc31f..e4b9d6d7e 100644 --- a/pkg-r/tests/testthat/test-PinSource.R +++ b/pkg-r/tests/testthat/test-PinSource.R @@ -398,3 +398,16 @@ describe("QueryChat + PinSource integration", { qc$cleanup() }) }) + +test_that("PinSource$cleanup() disconnects the connection it opened", { + skip_if_not_installed("pins") + skip_if_not_installed("duckdb") + skip_if_not_installed("nanoparquet") + + ps <- local_pin_source(type = "parquet") + conn <- ps$.__enclos_env__$private$conn + + ps$cleanup() + + expect_false(DBI::dbIsValid(conn)) +}) From a0329b9bcf35bd71d785e56e56c75604d43628d0 Mon Sep 17 00:00:00 2001 From: Carson Date: Sat, 12 Sep 2026 19:21:33 -0500 Subject: [PATCH 14/35] fix(r): keep $add_table() and $server(data_source=) description inference in sync, tighten late-change ordering resolve_data_description() (session path) and auto_fill_data_description() (instance path) diverged for multi-table sets: an identical table set could get a different inferred data_description depending on which path built it. auto_fill_data_description() now delegates the value computation to resolve_data_description(), the single source of truth, and only owns mode bookkeeping. Also, in $server() and $add_table(), late-change bookkeeping (locking the instance, warning about late changes) now happens only after the can-fail validation/build steps succeed, so a rejected or failed call doesn't leave the instance in an inconsistent state or warn about a change that never took effect. --- pkg-r/R/QueryChat.R | 73 +++++++++++++++-------- pkg-r/tests/testthat/test-QueryChat.R | 83 +++++++++++++++++++++++++++ 2 files changed, 131 insertions(+), 25 deletions(-) diff --git a/pkg-r/R/QueryChat.R b/pkg-r/R/QueryChat.R index 7a125483c..848a248f2 100644 --- a/pkg-r/R/QueryChat.R +++ b/pkg-r/R/QueryChat.R @@ -128,36 +128,46 @@ QueryChat <- R6::R6Class( invisible(NULL) }, - # Non-mutating counterpart of auto_fill_data_description(), for sets built - # on behalf of a session. + # Non-mutating counterpart of auto_fill_data_description(): the single + # source of truth for "what description should this set of sources get", + # given the instance's current mode/description. Used both to compute the + # value auto_fill_data_description() mutates in, and to build sets on + # behalf of a session (which must never mutate the instance). + # + # When sources isn't a single table and the mode isn't "supplied", this + # falls back to whatever description is already stored (bug-for-bug with + # auto_fill_data_description()'s historical early return): a multi-table + # set doesn't clear a stale single-source "inferred" description. resolve_data_description = function(sources) { if (private$.data_description_mode == "supplied") { return(private$.data_description) } - if (length(sources) == 1) { - desc <- sources[[1]]$get_data_description() - if (nzchar(desc %||% "")) { - return(desc) - } + if (length(sources) != 1) { + return(private$.data_description) + } + desc <- sources[[1]]$get_data_description() + if (nzchar(desc %||% "")) { + return(desc) } NULL }, auto_fill_data_description = function(sources = private$data_sources()) { - if (length(sources) != 1) { - return() + if (private$.data_description_mode == "supplied") { + return(invisible(NULL)) } - if (private$.data_description_mode == "inferred") { - private$.data_description <- NULL - private$.data_description_mode <- "empty" + if (length(sources) != 1) { + return(invisible(NULL)) } - if (private$.data_description_mode == "empty") { - desc <- sources[[1]]$get_data_description() - if (nzchar(desc %||% "")) { - private$.data_description <- desc - private$.data_description_mode <- "inferred" - } + private$.data_description <- private$resolve_data_description(sources) + private$.data_description_mode <- if ( + is.null(private$.data_description) + ) { + "empty" + } else { + "inferred" } + invisible(NULL) }, build_table_set = function( @@ -213,9 +223,9 @@ QueryChat <- R6::R6Class( ) return(invisible(NULL)) } - tryCatch(old_set$cleanup_executor(), error = function(e) NULL) + warn_on_cleanup_failure(old_set$cleanup_executor(), "query executor") for (source in replaced) { - source$cleanup() + warn_on_cleanup_failure(source$cleanup(), "data source") } invisible(NULL) }, @@ -509,7 +519,6 @@ QueryChat <- R6::R6Class( "Table {.val {table_name}} already exists. Use {.code replace = TRUE} to replace." ) } - private$check_late_change("add_table", destructive = exists) if ( is_data_source(data_source) && @@ -524,15 +533,29 @@ QueryChat <- R6::R6Class( } normalized <- normalize_data_source(data_source, table_name) + cleanup_normalized <- function() { + if (!inherits(data_source, "DataSource")) { + normalized$cleanup() + } + } next_sources <- current next_sources[[table_name]] <- normalized private$auto_fill_data_description(next_sources) new_set <- tryCatch( private$build_table_set(next_sources), error = function(e) { - if (!inherits(data_source, "DataSource")) { - normalized$cleanup() - } + cleanup_normalized() + stop(e) + } + ) + + # Only after the change is known to be valid do we check whether it's + # too late to apply it, so a rejected/failed add_table() doesn't warn. + tryCatch( + private$check_late_change("add_table", destructive = exists), + error = function(e) { + new_set$cleanup_executor() + cleanup_normalized() stop(e) } ) @@ -1150,7 +1173,6 @@ QueryChat <- R6::R6Class( ) } - private$.sessions_started <- TRUE table_set <- private$.table_set greeting_tables <- self$greeter$tables session_source <- NULL @@ -1256,6 +1278,7 @@ QueryChat <- R6::R6Class( }) %||% TRUE + private$.sessions_started <- TRUE mod_server( id %||% self$id, table_set = table_set, diff --git a/pkg-r/tests/testthat/test-QueryChat.R b/pkg-r/tests/testthat/test-QueryChat.R index 107bb1e69..b893155d9 100644 --- a/pkg-r/tests/testthat/test-QueryChat.R +++ b/pkg-r/tests/testthat/test-QueryChat.R @@ -1409,6 +1409,89 @@ describe("QueryChat table changes after a session has started", { }) }) +describe("auto_fill_data_description()/resolve_data_description() parity", { + # A DataFrameSource subclass that infers a description the way PinSource + # does (from its own metadata), without depending on the pins package. All + # tables in a set must share the same source class, so both the + # description-bearing table and the second table use this same class. + DescribedDataFrameSource <- R6::R6Class( + "DescribedDataFrameSource", + inherit = DataFrameSource, + public = list( + description = "", + get_data_description = function() self$description + ) + ) + + local_described_source <- function( + description = "", + data = new_test_df(), + table_name = "primary", + env = parent.frame() + ) { + source <- DescribedDataFrameSource$new(data, table_name) + source$description <- description + withr::defer(source$cleanup(), envir = env) + source + } + + it("keeps a single source's inferred description once a second table is added, identically via $add_table() and $server(data_source=)", { + skip_if_no_dataframe_engine() + + # Instance path: $add_table() adds a second table to a single-source + # instance that had inferred a description from its only source. + qc <- local_querychat( + local_described_source("Motor Trend Cars"), + "primary", + greeting = "hi" + ) + instance_desc_before <- qc$.__enclos_env__$private$.table_set$data_description + expect_equal(instance_desc_before, "Motor Trend Cars") + + qc$add_table( + local_described_source(data = new_metrics_df(), table_name = "secondary"), + "secondary" + ) + instance_desc_after <- qc$.__enclos_env__$private$.table_set$data_description + expect_equal(instance_desc_after, "Motor Trend Cars") + + # Session path: $server(data_source=) builds an equivalent two-table set + # for one session, starting from the same single inferred-description + # instance state. Before the fix, this produced NULL instead of the + # stale "Motor Trend Cars" description the instance path kept. + qc2 <- local_querychat( + local_described_source("Motor Trend Cars"), + "primary", + greeting = "hi", + client = mock_ellmer_chat_client() + ) + calls <- new.env(parent = emptyenv()) + calls$args <- list() + testthat::local_mocked_bindings( + mod_server = function(id, ...) { + calls$args[[length(calls$args) + 1L]] <- list(...) + list() + }, + .package = "querychat" + ) + session <- shiny::MockShinySession$new() + withr::defer(if (!session$isClosed()) session$close()) + shiny::withReactiveDomain( + session, + qc2$server( + data_source = local_described_source( + data = new_metrics_df(), + table_name = "secondary" + ), + table_name = "secondary" + ) + ) + + session_table_set <- calls$args[[1]]$table_set + expect_equal(session_table_set$data_description, instance_desc_after) + }) +}) + describe("QueryChatGreeter", { skip_if_no_dataframe_engine() From ce4e9796ea598fe90274d9cd88b1753d6655b87e Mon Sep 17 00:00:00 2001 From: Carson Date: Sat, 12 Sep 2026 19:21:47 -0500 Subject: [PATCH 15/35] fix(r): TableSet$cleanup_executor() resets its executor so it can rebuild Previously, after cleanup_executor() closed the executor, executor_built() still reported TRUE and a subsequent $executor() call would hand back the already-closed executor instead of building a fresh one. --- pkg-r/R/TableSet.R | 7 ++++++- pkg-r/tests/testthat/test-TableSet.R | 11 +++++++++++ 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/pkg-r/R/TableSet.R b/pkg-r/R/TableSet.R index 50334ce49..2fe8d4f98 100644 --- a/pkg-r/R/TableSet.R +++ b/pkg-r/R/TableSet.R @@ -48,7 +48,12 @@ TableSet <- R6::R6Class( # Closes the executor if it was ever built. Never touches data sources. cleanup_executor = function() { if (!is.null(private$.executor)) { - private$.executor$cleanup() + tryCatch( + private$.executor$cleanup(), + finally = { + private$.executor <- NULL + } + ) } invisible(NULL) } diff --git a/pkg-r/tests/testthat/test-TableSet.R b/pkg-r/tests/testthat/test-TableSet.R index cbe25549a..4011062c3 100644 --- a/pkg-r/tests/testthat/test-TableSet.R +++ b/pkg-r/tests/testthat/test-TableSet.R @@ -68,6 +68,17 @@ test_that("TableSet$cleanup_executor() closes a built DuckDB executor", { ex <- ts$executor() ts$cleanup_executor() expect_error(ex$execute_query("SELECT 1")) + expect_false(ts$executor_built()) +}) + +test_that("TableSet$cleanup_executor() is safe to call twice", { + skip_if_not_installed("duckdb") + a <- local_data_frame_source(new_test_df(), "a") + b <- local_data_frame_source(new_users_df(), "b") + ts <- local_table_set(list(a = a, b = b)) + ts$executor() + ts$cleanup_executor() + expect_no_error(ts$cleanup_executor()) }) test_that("TableSet fields are read-only", { From f7e18f9b095eb493a08df60dac1a2c7c4947b770 Mon Sep 17 00:00:00 2001 From: Carson Date: Sat, 12 Sep 2026 19:21:52 -0500 Subject: [PATCH 16/35] test(r): rename a test that no longer claims to close anything The exercised source is a DBISource wrapping a caller's connection, whose cleanup() is correctly a no-op, so nothing is actually "closed" in this scenario; the test only verifies the instance is left untouched. --- pkg-r/tests/testthat/test-server_data_source.R | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pkg-r/tests/testthat/test-server_data_source.R b/pkg-r/tests/testthat/test-server_data_source.R index 26fffe8a6..f7a4cf076 100644 --- a/pkg-r/tests/testthat/test-server_data_source.R +++ b/pkg-r/tests/testthat/test-server_data_source.R @@ -225,7 +225,7 @@ describe("QueryChat$server(data_source = ) session cleanup", { expect_true(source_conn_valid(caller_source)) }) - it("closes the normalized source and leaves the instance untouched when registration fails", { + it("leaves the instance untouched when registration fails", { skip_if_no_dataframe_engine() skip_if_not_installed("RSQLite") withr::local_envvar(OPENAI_API_KEY = "boop") From 30073d28edf49386e5bb3efc2e2185bcc1332cb6 Mon Sep 17 00:00:00 2001 From: Carson Date: Sat, 12 Sep 2026 19:39:48 -0500 Subject: [PATCH 17/35] fix(r): don't let a rejected add_table()/add_tables() call corrupt the data description check_late_change() (or build_table_set()) can abort after auto_fill_data_description() already mutated private$.data_description / .data_description_mode, leaving those fields reflecting a change that never actually took effect while private$.table_set correctly stayed unchanged. Split auto_fill_data_description() into a non-mutating pending_data_description() and a commit_data_description(), and only commit in add_table()/add_tables() once the can-fail steps have all succeeded. --- pkg-r/R/QueryChat.R | 62 ++++++++++++++++++++------- pkg-r/tests/testthat/test-QueryChat.R | 31 ++++++++++++++ 2 files changed, 78 insertions(+), 15 deletions(-) diff --git a/pkg-r/R/QueryChat.R b/pkg-r/R/QueryChat.R index 848a248f2..d765442bb 100644 --- a/pkg-r/R/QueryChat.R +++ b/pkg-r/R/QueryChat.R @@ -152,24 +152,37 @@ QueryChat <- R6::R6Class( NULL }, - auto_fill_data_description = function(sources = private$data_sources()) { + # Non-mutating: computes what auto_fill_data_description() would set, but + # leaves the caller to decide when (or whether) to actually commit it via + # commit_data_description(). Returns NULL when there's nothing to update, + # matching auto_fill_data_description()'s early-return conditions. + pending_data_description = function(sources) { if (private$.data_description_mode == "supplied") { - return(invisible(NULL)) + return(NULL) } if (length(sources) != 1) { - return(invisible(NULL)) + return(NULL) } - private$.data_description <- private$resolve_data_description(sources) - private$.data_description_mode <- if ( - is.null(private$.data_description) - ) { - "empty" - } else { - "inferred" + desc <- private$resolve_data_description(sources) + list( + description = desc, + mode = if (is.null(desc)) "empty" else "inferred" + ) + }, + + commit_data_description = function(pending) { + if (is.null(pending)) { + return(invisible(NULL)) } + private$.data_description <- pending$description + private$.data_description_mode <- pending$mode invisible(NULL) }, + auto_fill_data_description = function(sources = private$data_sources()) { + private$commit_data_description(private$pending_data_description(sources)) + }, + build_table_set = function( sources, data_description = private$.data_description @@ -540,9 +553,17 @@ QueryChat <- R6::R6Class( } next_sources <- current next_sources[[table_name]] <- normalized - private$auto_fill_data_description(next_sources) + pending_description <- private$pending_data_description(next_sources) + candidate_description <- if (is.null(pending_description)) { + private$.data_description + } else { + pending_description$description + } new_set <- tryCatch( - private$build_table_set(next_sources), + private$build_table_set( + next_sources, + data_description = candidate_description + ), error = function(e) { cleanup_normalized() stop(e) @@ -550,7 +571,8 @@ QueryChat <- R6::R6Class( ) # Only after the change is known to be valid do we check whether it's - # too late to apply it, so a rejected/failed add_table() doesn't warn. + # too late to apply it, so a rejected/failed add_table() doesn't warn, + # and doesn't commit the data description either. tryCatch( private$check_late_change("add_table", destructive = exists), error = function(e) { @@ -559,6 +581,7 @@ QueryChat <- R6::R6Class( stop(e) } ) + private$commit_data_description(pending_description) old_source <- current[[table_name]] replaced <- if ( @@ -656,8 +679,17 @@ QueryChat <- R6::R6Class( for (table_name in tables) { next_sources[[table_name]] <- normalized[[table_name]] } - private$auto_fill_data_description(next_sources) - new_set <- private$build_table_set(next_sources) + pending_description <- private$pending_data_description(next_sources) + candidate_description <- if (is.null(pending_description)) { + private$.data_description + } else { + pending_description$description + } + new_set <- private$build_table_set( + next_sources, + data_description = candidate_description + ) + private$commit_data_description(pending_description) replaced <- list() for (table_name in tables) { diff --git a/pkg-r/tests/testthat/test-QueryChat.R b/pkg-r/tests/testthat/test-QueryChat.R index b893155d9..72da53043 100644 --- a/pkg-r/tests/testthat/test-QueryChat.R +++ b/pkg-r/tests/testthat/test-QueryChat.R @@ -1490,6 +1490,37 @@ describe("auto_fill_data_description()/resolve_data_description() parity", { session_table_set <- calls$args[[1]]$table_set expect_equal(session_table_set$data_description, instance_desc_after) }) + + it("rejected destructive $add_table() after a session started leaves the data description untouched", { + skip_if_no_dataframe_engine() + + qc <- local_querychat( + local_described_source("Motor Trend Cars", table_name = "primary"), + "primary", + greeting = "hi" + ) + desc_before <- qc$.__enclos_env__$private$.data_description + mode_before <- qc$.__enclos_env__$private$.data_description_mode + expect_equal(desc_before, "Motor Trend Cars") + expect_equal(mode_before, "inferred") + + qc$.__enclos_env__$private$.sessions_started <- TRUE + + expect_error( + qc$add_table( + local_described_source(table_name = "primary"), + "primary", + replace = TRUE + ), + "replace or remove" + ) + + expect_equal(qc$.__enclos_env__$private$.data_description, desc_before) + expect_equal( + qc$.__enclos_env__$private$.data_description_mode, + mode_before + ) + }) }) describe("QueryChatGreeter", { From e36bca33cbe1486d13366813356c94f9deab1640 Mon Sep 17 00:00:00 2001 From: Carson Date: Sat, 12 Sep 2026 20:56:26 -0500 Subject: [PATCH 18/35] fix(py): preserve table registration order on replace, warn on failed executor cleanup add_table()/add_tables() rebuilt the merged dict by removing the replaced key and re-adding it, which moved a replaced table to the end of table_names() instead of keeping its original position -- changing the primary/first table an app selects. Also swap a silently-suppressed executor cleanup for warn_on_failure(), matching the rest of the teardown code's "warn, don't swallow" pattern. Found via Copilot review on #310. --- pkg-py/src/querychat/_querychat_base.py | 15 +++++---------- pkg-py/tests/test_base.py | 7 +++++++ pkg-py/tests/test_multi_table.py | 7 +++++++ 3 files changed, 19 insertions(+), 10 deletions(-) diff --git a/pkg-py/src/querychat/_querychat_base.py b/pkg-py/src/querychat/_querychat_base.py index 2b3bcd5c1..0fd9b56b2 100644 --- a/pkg-py/src/querychat/_querychat_base.py +++ b/pkg-py/src/querychat/_querychat_base.py @@ -218,8 +218,7 @@ def _swap_table_set(self, new_set: TableSet, *, replaced: list[DataSource]) -> N ) self._superseded_table_sets.append(old_set) return - with contextlib.suppress(Exception): - old_set.cleanup_executor() + warn_on_failure(old_set.cleanup_executor, "query executor") for source in replaced: source.cleanup() @@ -486,10 +485,8 @@ def add_table( normalized = normalize_data_source(data_source, table_name) try: self._warn_if_prompt_rebuilt_with_history() - merged = { - **{k: v for k, v in self._data_sources.items() if k != table_name}, - table_name: normalized, - } + merged = dict(self._data_sources) + merged[table_name] = normalized new_set = self._build_table_set(merged) except Exception: if normalized is not data_source: @@ -610,10 +607,8 @@ def normalized_builder(name: str) -> DataSource: normalized = {name: normalized_builder(name) for name in tables} self._warn_if_prompt_rebuilt_with_history() - merged = { - **{k: v for k, v in self._data_sources.items() if k not in normalized}, - **normalized, - } + merged = dict(self._data_sources) + merged.update(normalized) new_set = self._build_table_set(merged) replaced = [ diff --git a/pkg-py/tests/test_base.py b/pkg-py/tests/test_base.py index 59b0f6372..248fcacaa 100644 --- a/pkg-py/tests/test_base.py +++ b/pkg-py/tests/test_base.py @@ -362,6 +362,13 @@ def test_replace_true_succeeds(self, multi_table_engine): qc.add_tables(multi_table_engine, ["orders"], replace=True) assert "orders" in qc.table_names() + def test_replace_preserves_table_order(self, multi_table_engine): + """Replacing a non-last table keeps its original registration order.""" + qc = QueryChatBase() + qc.add_tables(multi_table_engine, ["orders", "customers"]) + qc.add_tables(multi_table_engine, ["orders"], replace=True) + assert qc.table_names() == ["orders", "customers"] + def test_non_engine_raises_type_error(self, sample_df): qc = QueryChatBase() with pytest.raises(TypeError, match=r"sqlalchemy\.Engine or ibis SQLBackend"): diff --git a/pkg-py/tests/test_multi_table.py b/pkg-py/tests/test_multi_table.py index 96f24bbf2..e6433ec36 100644 --- a/pkg-py/tests/test_multi_table.py +++ b/pkg-py/tests/test_multi_table.py @@ -222,6 +222,13 @@ def test_add_table_after_server_raises(self, orders_df, customers_df): with pytest.raises(RuntimeError, match="replace or remove"): qc.add_table(orders_df, "orders", replace=True) + def test_add_table_replace_preserves_position(self, orders_df, customers_df): + """Replacing a non-last table keeps its original registration order.""" + qc = QueryChat(orders_df, "orders", greeting="Hello!") + qc.add_table(customers_df, "customers") + qc.add_table(orders_df, "orders", replace=True) + assert qc.table_names() == ["orders", "customers"] + class TestRemoveTable: """Tests for remove_table() method.""" From 3c214f9b4932f0eaa86ac2abe427f09446599da6 Mon Sep 17 00:00:00 2001 From: Carson Date: Sat, 12 Sep 2026 20:56:33 -0500 Subject: [PATCH 19/35] fix: only mark sessions as started once server registration succeeds $server()/.server()/app_server() flipped sessions_started/_sessions_started before the module server actually finished setting up. If that setup raised, the instance was left permanently locked: later add_table() calls would warn and replace/remove would raise, even though no session had actually started. Found via Copilot review on #310 (R side); mirrored the same fix into Python's two synchronous server-registration paths (Express's _ensure_server_started intentionally keeps its own ordering, since it's designed to never retry mod_server() regardless of outcome). --- pkg-py/src/querychat/_shiny.py | 7 ++++--- pkg-py/tests/test_server_data_source.py | 22 ++++++++++++++++++++++ pkg-r/R/QueryChat.R | 5 +++-- pkg-r/tests/testthat/test-QueryChat.R | 20 ++++++++++++++++++++ 4 files changed, 49 insertions(+), 5 deletions(-) diff --git a/pkg-py/src/querychat/_shiny.py b/pkg-py/src/querychat/_shiny.py index 158857556..9f7bec3af 100644 --- a/pkg-py/src/querychat/_shiny.py +++ b/pkg-py/src/querychat/_shiny.py @@ -418,7 +418,6 @@ def app_ui(request): def app_server(input: Inputs, output: Outputs, session: Session): if enable_bookmarking: session.bookmark.exclude.extend(["reset_query", "sql_editor"]) - self._sessions_started = True table_set = self._require_table_set("app") vals = mod_server( self.id, @@ -431,6 +430,7 @@ def app_server(input: Inputs, output: Outputs, session: Session): greeting_base=None, greeting_tables=list(self.greeter.tables), ) + self._sessions_started = True @reactive.calc def active_table_name() -> str: @@ -710,7 +710,6 @@ def server( ".server() must be called within an active Shiny session (i.e., within the server function). " ) - self._sessions_started = True table_set: TableSet | None = self._table_set greeting_tables = list(self.greeter.tables) session_source: DataSource | None = None @@ -785,7 +784,7 @@ def create_session_client(**kwargs) -> chatlas.Chat: ) ) - return mod_server( + result = mod_server( id or self.id, table_set=table_set, greeting=self.greeting, @@ -796,6 +795,8 @@ def create_session_client(**kwargs) -> chatlas.Chat: greeting_base=resolved_client, greeting_tables=greeting_tables, ) + self._sessions_started = True + return result class QueryChatExpress(QueryChatBase[IntoFrameT]): diff --git a/pkg-py/tests/test_server_data_source.py b/pkg-py/tests/test_server_data_source.py index 49ba875de..b1e7714f5 100644 --- a/pkg-py/tests/test_server_data_source.py +++ b/pkg-py/tests/test_server_data_source.py @@ -1,5 +1,6 @@ """Tests for QueryChat.server(data_source=...) parity with R (#300).""" +import warnings from unittest.mock import MagicMock, patch import pandas as pd @@ -326,3 +327,24 @@ def spy(data_source, table_name): sessions[0].end() with pytest.raises(duckdb.ConnectionException): session_source.execute_query("SELECT 1") + + def test_mod_server_failure_does_not_mark_sessions_started( + self, users_df, monkeypatch + ): + """A failed mod_server() call must not lock out later add_table()/remove_table().""" + monkeypatch.setattr(shiny_mod, "get_current_session", lambda: MagicMock()) + + def failing_mod_server(*args, **kwargs): + raise RuntimeError("boom") + + monkeypatch.setattr(shiny_mod, "mod_server", failing_mod_server) + + qc = shiny_mod.QueryChat(users_df, "users") + with pytest.raises(RuntimeError, match="boom"): + qc.server() + + assert qc._sessions_started is False + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + qc.add_table(pd.DataFrame({"id": [1]}), "other") + assert not any("session has started" in str(w.message) for w in caught) diff --git a/pkg-r/R/QueryChat.R b/pkg-r/R/QueryChat.R index d765442bb..1553dff45 100644 --- a/pkg-r/R/QueryChat.R +++ b/pkg-r/R/QueryChat.R @@ -1310,8 +1310,7 @@ QueryChat <- R6::R6Class( }) %||% TRUE - private$.sessions_started <- TRUE - mod_server( + result <- mod_server( id %||% self$id, table_set = table_set, greeting = self$greeting, @@ -1322,6 +1321,8 @@ QueryChat <- R6::R6Class( greeting_base = base_client, greeting_tables = greeting_tables ) + private$.sessions_started <- TRUE + result }, #' @description diff --git a/pkg-r/tests/testthat/test-QueryChat.R b/pkg-r/tests/testthat/test-QueryChat.R index 72da53043..861b3c85e 100644 --- a/pkg-r/tests/testthat/test-QueryChat.R +++ b/pkg-r/tests/testthat/test-QueryChat.R @@ -1407,6 +1407,26 @@ describe("QueryChat table changes after a session has started", { expect_false(DBI::dbIsValid(old_source$.__enclos_env__$private$conn)) }) + + it("does not mark sessions as started when $server() fails", { + skip_if_no_dataframe_engine() + qc <- local_querychat(new_users_df(), "users", greeting = "hi") + local_mocked_bindings( + mod_server = function(...) stop("boom"), + .package = "querychat" + ) + + expect_error( + shiny::testServer(function(input, output, session) qc$server(), {}), + "boom" + ) + + expect_false(qc$.__enclos_env__$private$.sessions_started) + expect_no_warning( + qc$add_table(new_test_df(), "other"), + message = "after a session has started" + ) + }) }) describe("auto_fill_data_description()/resolve_data_description() parity", { From dd276a0b60cae74586e98cfc7f1d4430ecfb3f7f Mon Sep 17 00:00:00 2001 From: Carson Date: Sat, 12 Sep 2026 21:15:10 -0500 Subject: [PATCH 20/35] fix(py): make remaining teardown paths warn instead of raise, fix Express session-start ordering - _swap_table_set() now wraps each replaced source's cleanup() in warn_on_failure(), matching the executor cleanup immediately above it and R's equivalent -- a failing cleanup no longer aborts the replace/remove call after the new state is already committed, and no longer skips cleaning up the remaining replaced sources. - A spec-resolved .server(client=...) override's close callback is now wrapped the same way, so a broken client can't raise out of session.on_ended. - QueryChatExpress._ensure_server_started() only marks _sessions_started after mod_server() succeeds, mirroring QueryChat.server()/app_server(). _server_attempted still flips before the call, since Express is intentionally never supposed to retry mod_server() regardless of outcome. Found via Copilot review on #310. --- pkg-py/src/querychat/_querychat_base.py | 4 ++-- pkg-py/src/querychat/_shiny.py | 2 +- pkg-py/tests/test_multi_table.py | 13 +++++++++++ pkg-py/tests/test_server_data_source.py | 19 +++++++++++++++ pkg-py/tests/test_shiny.py | 31 +++++++++++++++++++++++++ 5 files changed, 66 insertions(+), 3 deletions(-) diff --git a/pkg-py/src/querychat/_querychat_base.py b/pkg-py/src/querychat/_querychat_base.py index 0fd9b56b2..35ae46ba6 100644 --- a/pkg-py/src/querychat/_querychat_base.py +++ b/pkg-py/src/querychat/_querychat_base.py @@ -220,7 +220,7 @@ def _swap_table_set(self, new_set: TableSet, *, replaced: list[DataSource]) -> N return warn_on_failure(old_set.cleanup_executor, "query executor") for source in replaced: - source.cleanup() + warn_on_failure(source.cleanup, "data source") def _require_single_table(self, method_name: str) -> None: """Raise if multiple tables are registered, directing to per-table API.""" @@ -253,7 +253,7 @@ def _resolve_session_client( return None resolved = resolve_client(client) if not isinstance(client, chatlas.Chat): - session.on_ended(resolved.close) + session.on_ended(lambda: warn_on_failure(resolved.close, "chatlas client")) return resolved def _create_session_client( diff --git a/pkg-py/src/querychat/_shiny.py b/pkg-py/src/querychat/_shiny.py index 9f7bec3af..d28adb392 100644 --- a/pkg-py/src/querychat/_shiny.py +++ b/pkg-py/src/querychat/_shiny.py @@ -1066,7 +1066,6 @@ def _ensure_server_started(self) -> None: if self._table_set is None: return self._server_attempted = True - self._sessions_started = True table_set = self._table_set resolved_history: bool | HistoryOptions = ( self.history @@ -1088,6 +1087,7 @@ def _ensure_server_started(self) -> None: greeting_base=None, greeting_tables=list(self.greeter.tables), ) + self._sessions_started = True def sidebar( self, diff --git a/pkg-py/tests/test_multi_table.py b/pkg-py/tests/test_multi_table.py index e6433ec36..b5296d14d 100644 --- a/pkg-py/tests/test_multi_table.py +++ b/pkg-py/tests/test_multi_table.py @@ -314,6 +314,19 @@ def test_cleanup_all_sources(self, orders_df, customers_df): # Connections should be closed after cleanup # (DuckDB connections don't have is_closed, but they're closed) + def test_replace_warns_instead_of_raising_on_replaced_source_cleanup_failure( + self, orders_df + ): + """A failing cleanup() on a replaced source warns rather than raising.""" + qc = QueryChat(orders_df, "orders", greeting="Hello!") + old_source = qc._data_sources["orders"] + with ( + patch.object(old_source, "cleanup", side_effect=RuntimeError("boom")), + pytest.warns(UserWarning, match="Failed to clean up data source"), + ): + qc.add_table(orders_df, "orders", replace=True) + assert qc.table_names() == ["orders"] + @pytest.fixture def orders_qc(orders_df): diff --git a/pkg-py/tests/test_server_data_source.py b/pkg-py/tests/test_server_data_source.py index b1e7714f5..560b706c7 100644 --- a/pkg-py/tests/test_server_data_source.py +++ b/pkg-py/tests/test_server_data_source.py @@ -5,6 +5,7 @@ import pandas as pd import pytest +import querychat._querychat_base as base_mod import querychat._shiny as shiny_mod @@ -348,3 +349,21 @@ def failing_mod_server(*args, **kwargs): warnings.simplefilter("always") qc.add_table(pd.DataFrame({"id": [1]}), "other") assert not any("session has started" in str(w.message) for w in caught) + + def test_client_override_close_failure_warns_on_session_end( + self, users_df, session_runs, monkeypatch + ): + """ + A spec-resolved `.server(client=...)` override that fails to close on + session end warns instead of raising. + """ + sessions, _calls = session_runs + fake_client = MagicMock() + fake_client.close.side_effect = RuntimeError("boom") + monkeypatch.setattr(base_mod, "resolve_client", lambda client: fake_client) + + qc = shiny_mod.QueryChat(users_df, "users") + qc.server(client="openai") + + with pytest.warns(UserWarning, match="Failed to clean up chatlas client"): + sessions[0].end() diff --git a/pkg-py/tests/test_shiny.py b/pkg-py/tests/test_shiny.py index e68763a23..599a22bcb 100644 --- a/pkg-py/tests/test_shiny.py +++ b/pkg-py/tests/test_shiny.py @@ -186,6 +186,37 @@ def fake_mod_server(*args, **kwargs): assert captured["history"].restore_mode == "bookmark" +def test_express_mod_server_failure_does_not_mark_sessions_started(): + """ + A failed mod_server() call must not lock the instance out of later + add_table()/remove_table() calls (mirrors QueryChat.server()'s ordering + fix). _server_attempted still flips immediately so Express never retries + mod_server() -- that invariant is intentionally unaffected. + """ + from unittest.mock import MagicMock, patch + + import pandas as pd + from querychat._shiny import QueryChatExpress + from shiny._namespaces import Root + from shiny.session import session_context + + def failing_mod_server(*args, **kwargs): + raise RuntimeError("boom") + + mock_session = MagicMock() + mock_session.ns = Root + with session_context(mock_session): + qc = QueryChatExpress(pd.DataFrame({"a": [1, 2, 3]}), "a_table") + with ( + patch("querychat._shiny.mod_server", failing_mod_server), + pytest.raises(RuntimeError, match="boom"), + ): + qc._ensure_server_started() + + assert qc._server_attempted is True + assert qc._sessions_started is False + + def test_express_explicit_enable_bookmarking_warns(): from unittest.mock import MagicMock, patch From 87796fb8732e3b1c6fabb41bb77f8e3b7a095b8a Mon Sep 17 00:00:00 2001 From: Carson Date: Sat, 12 Sep 2026 21:15:16 -0500 Subject: [PATCH 21/35] docs(r): fix DBISource/TblSqlSource examples still telling users to skip cleanup() Both examples predated the change that made cleanup() a no-op for caller-owned connections; they still said cleanup() would disconnect the connection and told users to skip calling it to keep the connection open. Show the correct pattern instead: call cleanup() (a no-op) and disconnect the connection yourself. Found via Copilot review on #310. --- pkg-r/R/DBISource.R | 4 ++-- pkg-r/R/TblSqlSource.R | 3 ++- pkg-r/man/DBISource.Rd | 4 ++-- pkg-r/man/TblSqlSource.Rd | 3 ++- 4 files changed, 8 insertions(+), 6 deletions(-) diff --git a/pkg-r/R/DBISource.R b/pkg-r/R/DBISource.R index fedbf3a15..610d77e93 100644 --- a/pkg-r/R/DBISource.R +++ b/pkg-r/R/DBISource.R @@ -18,9 +18,9 @@ #' # Execute a query #' result <- db_source$execute_query("SELECT * FROM mtcars WHERE mpg > 25") #' -#' # Note: cleanup() will disconnect the connection -#' # If you want to keep the connection open, don't call cleanup() +#' # cleanup() is a no-op: you own `con`, so disconnect it yourself #' db_source$cleanup() +#' DBI::dbDisconnect(con) #' #' @export DBISource <- R6::R6Class( diff --git a/pkg-r/R/TblSqlSource.R b/pkg-r/R/TblSqlSource.R index 5c9af2124..bec97515b 100644 --- a/pkg-r/R/TblSqlSource.R +++ b/pkg-r/R/TblSqlSource.R @@ -22,8 +22,9 @@ #' # Or collect the entire data frame into local memory #' dplyr::collect(result) #' -#' # Finally, clean up when done with the database (closes the DB connection) +#' # cleanup() is a no-op: you own `con`, so disconnect it yourself when done #' mtcars_source$cleanup() +#' DBI::dbDisconnect(con, shutdown = TRUE) #' #' @export TblSqlSource <- R6::R6Class( diff --git a/pkg-r/man/DBISource.Rd b/pkg-r/man/DBISource.Rd index 9c1e622be..bbb896529 100644 --- a/pkg-r/man/DBISource.Rd +++ b/pkg-r/man/DBISource.Rd @@ -23,9 +23,9 @@ db_source$get_db_type() # Returns "SQLite" # Execute a query result <- db_source$execute_query("SELECT * FROM mtcars WHERE mpg > 25") -# Note: cleanup() will disconnect the connection -# If you want to keep the connection open, don't call cleanup() +# cleanup() is a no-op: you own `con`, so disconnect it yourself db_source$cleanup() +DBI::dbDisconnect(con) \dontshow{\}) # examplesIf} } \section{Super class}{ diff --git a/pkg-r/man/TblSqlSource.Rd b/pkg-r/man/TblSqlSource.Rd index fad878df3..d15887a03 100644 --- a/pkg-r/man/TblSqlSource.Rd +++ b/pkg-r/man/TblSqlSource.Rd @@ -26,8 +26,9 @@ dplyr::count(result, cyl, gear) # Or collect the entire data frame into local memory dplyr::collect(result) -# Finally, clean up when done with the database (closes the DB connection) +# cleanup() is a no-op: you own `con`, so disconnect it yourself when done mtcars_source$cleanup() +DBI::dbDisconnect(con, shutdown = TRUE) \dontshow{\}) # examplesIf} } \section{Super classes}{ From e8a5eecce2f2c7e04380365d2cfbaf3555830918 Mon Sep 17 00:00:00 2001 From: Carson Date: Sat, 12 Sep 2026 21:35:19 -0500 Subject: [PATCH 22/35] fix: rollback cleanup failures no longer mask the original build/registration error add_table()'s and .server(data_source=)'s failure paths cleaned up a newly normalized source directly. If that cleanup itself raised, the cleanup exception replaced the original build/compatibility error the caller was supposed to see. Route both through the existing warn_on_failure()/warn_on_cleanup_failure() helpers so a broken rollback cleanup becomes a warning and the original error still propagates. Found via Copilot review on #310. --- pkg-py/src/querychat/_querychat_base.py | 2 +- pkg-py/src/querychat/_shiny.py | 6 ++-- pkg-py/tests/test_multi_table.py | 22 +++++++++++++ pkg-py/tests/test_server_data_source.py | 17 ++++++++++ pkg-r/R/QueryChat.R | 9 ++++-- pkg-r/tests/testthat/test-QueryChat.R | 25 +++++++++++++++ .../tests/testthat/test-server_data_source.R | 32 +++++++++++++++++++ 7 files changed, 107 insertions(+), 6 deletions(-) diff --git a/pkg-py/src/querychat/_querychat_base.py b/pkg-py/src/querychat/_querychat_base.py index 35ae46ba6..b79073254 100644 --- a/pkg-py/src/querychat/_querychat_base.py +++ b/pkg-py/src/querychat/_querychat_base.py @@ -490,7 +490,7 @@ def add_table( new_set = self._build_table_set(merged) except Exception: if normalized is not data_source: - normalized.cleanup() + warn_on_failure(normalized.cleanup, "data source") raise old_source = self._data_sources.get(table_name) diff --git a/pkg-py/src/querychat/_shiny.py b/pkg-py/src/querychat/_shiny.py index d28adb392..bf29b3357 100644 --- a/pkg-py/src/querychat/_shiny.py +++ b/pkg-py/src/querychat/_shiny.py @@ -736,7 +736,7 @@ def server( ) except Exception: if session_source is not data_source: - session_source.cleanup() + warn_on_failure(session_source.cleanup, "session data source") raise if resolved_table_name not in greeting_tables: greeting_tables.append(resolved_table_name) @@ -758,7 +758,9 @@ def cleanup_session() -> None: resolved_client = self._resolve_session_client(client, session) def create_session_client(**kwargs) -> chatlas.Chat: - return self._create_session_client(table_set, base=resolved_client, **kwargs) + return self._create_session_client( + table_set, base=resolved_client, **kwargs + ) if enable_bookmarking is not None: warnings.warn( diff --git a/pkg-py/tests/test_multi_table.py b/pkg-py/tests/test_multi_table.py index b5296d14d..709f65e71 100644 --- a/pkg-py/tests/test_multi_table.py +++ b/pkg-py/tests/test_multi_table.py @@ -520,6 +520,28 @@ def fail_compat(*a, **kw): staged_source.execute_query("SELECT 1") assert orders_qc.table_names() == original_table_names + def test_add_table_failure_warns_but_preserves_original_error_if_rollback_cleanup_fails( + self, orders_qc, customers_df, monkeypatch + ): + """A rollback cleanup() failure must not mask the original build error.""" + + def fail_compat(*a, **kw): + raise ValueError("compat check failed") + + monkeypatch.setattr( + "querychat._query_executor.check_source_compatibility", + fail_compat, + ) + + with ( + patch.object( + DataFrameSource, "cleanup", side_effect=RuntimeError("cleanup boom") + ), + pytest.warns(UserWarning, match="Failed to clean up data source"), + pytest.raises(ValueError, match="compat check failed"), + ): + orders_qc.add_table(customers_df, "customers") + def test_add_table_replace_failure_cleans_staged_source_and_preserves_state( self, orders_qc, customers_df, monkeypatch ): diff --git a/pkg-py/tests/test_server_data_source.py b/pkg-py/tests/test_server_data_source.py index 560b706c7..eedb81d56 100644 --- a/pkg-py/tests/test_server_data_source.py +++ b/pkg-py/tests/test_server_data_source.py @@ -297,6 +297,23 @@ def spy(data_source, table_name): qc.server() assert calls[-1]["table_set"].table_names == ["users"] + def test_failed_registration_warns_but_preserves_original_error_if_rollback_cleanup_fails( + self, users_df, session_runs, monkeypatch + ): + """A rollback cleanup() failure must not mask the original registration error.""" + import polars as pl + from querychat._datasource import DataFrameSource + + _sessions, _calls = session_runs + qc = shiny_mod.QueryChat(users_df, "users") + + with ( + patch.object(DataFrameSource, "cleanup", side_effect=RuntimeError("boom")), + pytest.warns(UserWarning, match="Failed to clean up session data source"), + pytest.raises(ValueError, match="same DataFrame backend"), + ): + qc.server(data_source=pl.DataFrame({"id": [1]}), table_name="other") + def test_client_resolution_failure_still_cleans_up_session_source( self, users_df, session_runs, monkeypatch ): diff --git a/pkg-r/R/QueryChat.R b/pkg-r/R/QueryChat.R index 1553dff45..dd964768c 100644 --- a/pkg-r/R/QueryChat.R +++ b/pkg-r/R/QueryChat.R @@ -548,7 +548,7 @@ QueryChat <- R6::R6Class( normalized <- normalize_data_source(data_source, table_name) cleanup_normalized <- function() { if (!inherits(data_source, "DataSource")) { - normalized$cleanup() + warn_on_cleanup_failure(normalized$cleanup(), "data source") } } next_sources <- current @@ -576,7 +576,7 @@ QueryChat <- R6::R6Class( tryCatch( private$check_late_change("add_table", destructive = exists), error = function(e) { - new_set$cleanup_executor() + warn_on_cleanup_failure(new_set$cleanup_executor(), "query executor") cleanup_normalized() stop(e) } @@ -1247,7 +1247,10 @@ QueryChat <- R6::R6Class( ), error = function(e) { if (!inherits(data_source, "DataSource")) { - session_source$cleanup() + warn_on_cleanup_failure( + session_source$cleanup(), + "session data source" + ) } stop(e) } diff --git a/pkg-r/tests/testthat/test-QueryChat.R b/pkg-r/tests/testthat/test-QueryChat.R index 861b3c85e..2e25dfb00 100644 --- a/pkg-r/tests/testthat/test-QueryChat.R +++ b/pkg-r/tests/testthat/test-QueryChat.R @@ -1259,6 +1259,31 @@ describe("QueryChat$add_table()", { ) expect_equal(length(qc$table_names()), 0L) }) + + it("warns but preserves the original error if rollback cleanup fails", { + skip_if_no_dataframe_engine() + qc <- local_querychat(new_users_df(), "users", greeting = "hi") + + attempt_add_table <- function() { + testthat::local_mocked_bindings( + dbDisconnect = function(...) stop("boom"), + .package = "DBI" + ) + testthat::local_mocked_bindings( + check_source_compatibility = function(...) { + cli::cli_abort("compat check failed") + }, + .package = "querychat" + ) + qc$add_table(new_test_df(), "other") + } + + expect_warning( + expect_error(attempt_add_table(), "compat check failed"), + "Failed to clean up data source" + ) + expect_equal(qc$table_names(), "users") + }) }) describe("QueryChat$add_tables()", { diff --git a/pkg-r/tests/testthat/test-server_data_source.R b/pkg-r/tests/testthat/test-server_data_source.R index f7a4cf076..a9434085f 100644 --- a/pkg-r/tests/testthat/test-server_data_source.R +++ b/pkg-r/tests/testthat/test-server_data_source.R @@ -246,4 +246,36 @@ describe("QueryChat$server(data_source = ) session cleanup", { "users" ) }) + + it("warns but preserves the original error if rollback cleanup fails", { + skip_if_no_dataframe_engine() + withr::local_envvar(OPENAI_API_KEY = "boop") + local_captured_mod_server() + qc <- QueryChat$new(new_users_df(), "users", greeting = "hi") + withr::defer(qc$cleanup()) + + attempt_registration <- function() { + testthat::local_mocked_bindings( + dbDisconnect = function(...) stop("boom"), + .package = "DBI" + ) + testthat::local_mocked_bindings( + check_source_compatibility = function(...) { + cli::cli_abort("compat check failed") + }, + .package = "querychat" + ) + start_server_session( + qc, + data_source = new_test_df(), + table_name = "other" + ) + } + + expect_warning( + expect_error(attempt_registration(), "compat check failed"), + "Failed to clean up session data source" + ) + expect_equal(qc$table_names(), "users") + }) }) From 2c480770eec7ebc63c4d684a772a0c925109c744 Mon Sep 17 00:00:00 2001 From: Carson Date: Sat, 12 Sep 2026 21:35:26 -0500 Subject: [PATCH 23/35] fix(py): make TableSet.data_sources/system_prompt read-only TableSet is documented as immutable after construction, and R's TableSet already enforces that (active bindings reject assignment). Python's version stored them as plain attributes, so ts.data_sources = ... and ts.system_prompt = ... both silently succeeded -- a running session's prompt or table mapping could diverge from its cached executor. Back them with private fields and expose read-only properties instead. Found via Copilot review on #310. --- pkg-py/src/querychat/_table_set.py | 12 ++++++++++-- pkg-py/tests/test_table_set.py | 12 ++++++++++++ 2 files changed, 22 insertions(+), 2 deletions(-) diff --git a/pkg-py/src/querychat/_table_set.py b/pkg-py/src/querychat/_table_set.py index b33930c81..978828ded 100644 --- a/pkg-py/src/querychat/_table_set.py +++ b/pkg-py/src/querychat/_table_set.py @@ -32,10 +32,18 @@ def __init__( ) -> None: if not data_sources: raise ValueError("TableSet requires at least one data source") - self.data_sources: Mapping[str, DataSource] = MappingProxyType( + self._data_sources: Mapping[str, DataSource] = MappingProxyType( dict(data_sources) ) - self.system_prompt = system_prompt + self._system_prompt = system_prompt + + @property + def data_sources(self) -> Mapping[str, DataSource]: + return self._data_sources + + @property + def system_prompt(self) -> QueryChatSystemPrompt: + return self._system_prompt @cached_property def executor(self) -> QueryExecutor: diff --git a/pkg-py/tests/test_table_set.py b/pkg-py/tests/test_table_set.py index 118ef7e9a..2a9c01855 100644 --- a/pkg-py/tests/test_table_set.py +++ b/pkg-py/tests/test_table_set.py @@ -37,6 +37,18 @@ def test_data_sources_is_read_only(users): ts.data_sources["other"] = ts.data_sources["users"] # type: ignore[index] +def test_data_sources_attribute_is_read_only(users): + ts = make_table_set(users=users) + with pytest.raises(AttributeError): + ts.data_sources = {} # type: ignore[misc] + + +def test_system_prompt_attribute_is_read_only(users): + ts = make_table_set(users=users) + with pytest.raises(AttributeError): + ts.system_prompt = None # type: ignore[misc] + + def test_table_names_preserve_insertion_order(users, orders): with pytest.warns(UserWarning, match="without a data_dict"): ts = make_table_set(users=users, orders=orders) From 94df2073b5f8817b318a245f2b038174280f36e4 Mon Sep 17 00:00:00 2001 From: Carson Date: Sun, 13 Sep 2026 10:30:01 -0500 Subject: [PATCH 24/35] refactor(py): make TableSet generic so mod_server keeps IntoFrameT typing mod_server() previously bound IntoFrameT from its data_sources parameter; taking a non-generic TableSet erased it and forced ServerValues[Any]. TableSet is now Generic[IntoFrameT], and @module.server's ParamSpec typing preserves the TypeVar, so ServerValues[IntoFrameT] is inferred at call sites again. Annotation-only; no runtime behavior change. --- pkg-py/src/querychat/_querychat_base.py | 20 ++++++++++++-------- pkg-py/src/querychat/_shiny.py | 2 +- pkg-py/src/querychat/_shiny_module.py | 18 +++++++----------- pkg-py/src/querychat/_table_set.py | 12 +++++++----- 4 files changed, 27 insertions(+), 25 deletions(-) diff --git a/pkg-py/src/querychat/_querychat_base.py b/pkg-py/src/querychat/_querychat_base.py index b79073254..5bfa3e289 100644 --- a/pkg-py/src/querychat/_querychat_base.py +++ b/pkg-py/src/querychat/_querychat_base.py @@ -92,10 +92,10 @@ def __init__( history: Optional[bool | HistoryOptions] = None, ): self._data_dicts: list[DataDict] = _normalize_data_dicts(data_dict) - self._table_set: TableSet | None = None + self._table_set: TableSet[IntoFrameT] | None = None # Instance sets swapped out by add_table()/add_tables() after a session # started. A running session may still hold one, so cleanup() closes them. - self._superseded_table_sets: list[TableSet] = [] + self._superseded_table_sets: list[TableSet[IntoFrameT]] = [] self._sessions_started = False self._deferred_table_name: str | None = None @@ -139,13 +139,13 @@ def __init__( self._deferred_table_name = table_name @property - def _data_sources(self) -> Mapping[str, DataSource]: + def _data_sources(self) -> Mapping[str, DataSource[IntoFrameT]]: """Read-only view of the instance's registered tables.""" if self._table_set is None: return {} return self._table_set.data_sources - def _require_table_set(self, method_name: str) -> TableSet: + def _require_table_set(self, method_name: str) -> TableSet[IntoFrameT]: if self._table_set is None: raise RuntimeError( f"At least one data source must be set before calling {method_name}(). " @@ -159,7 +159,9 @@ def _require_initialized(self, method_name: str) -> None: def _require_query_executor(self, method_name: str) -> QueryExecutor: return self._require_table_set(method_name).executor - def _build_table_set(self, sources: dict[str, DataSource]) -> TableSet: + def _build_table_set( + self, sources: dict[str, DataSource[IntoFrameT]] + ) -> TableSet[IntoFrameT]: validate_source_group_compatibility(sources) prompt = QueryChatSystemPrompt( prompt_template=self._prompt_template, @@ -203,7 +205,9 @@ def _check_late_change(self, method_name: str, *, destructive: bool) -> None: stacklevel=3, ) - def _swap_table_set(self, new_set: TableSet, *, replaced: list[DataSource]) -> None: + def _swap_table_set( + self, new_set: TableSet[IntoFrameT], *, replaced: list[DataSource] + ) -> None: old_set, self._table_set = self._table_set, new_set if old_set is None: return @@ -258,7 +262,7 @@ def _resolve_session_client( def _create_session_client( self, - table_set: TableSet, + table_set: TableSet[IntoFrameT], *, base: chatlas.Chat | None = None, tools: TOOL_GROUPS | tuple[TOOL_GROUPS, ...] | MISSING_TYPE | None = MISSING, @@ -368,7 +372,7 @@ def client_factory( prompt: str | Path, base: chatlas.Chat | None = None, *, - table_set: TableSet | None = None, + table_set: TableSet[IntoFrameT] | None = None, ) -> chatlas.Chat: resolved = table_set if table_set is not None else self._table_set sp = QueryChatSystemPrompt( diff --git a/pkg-py/src/querychat/_shiny.py b/pkg-py/src/querychat/_shiny.py index bf29b3357..a49379740 100644 --- a/pkg-py/src/querychat/_shiny.py +++ b/pkg-py/src/querychat/_shiny.py @@ -710,7 +710,7 @@ def server( ".server() must be called within an active Shiny session (i.e., within the server function). " ) - table_set: TableSet | None = self._table_set + table_set: TableSet[IntoFrameT] | None = self._table_set greeting_tables = list(self.greeter.tables) session_source: DataSource | None = None diff --git a/pkg-py/src/querychat/_shiny_module.py b/pkg-py/src/querychat/_shiny_module.py index 3e80a588f..3d7656d0a 100644 --- a/pkg-py/src/querychat/_shiny_module.py +++ b/pkg-py/src/querychat/_shiny_module.py @@ -241,7 +241,7 @@ def mod_server( output: Outputs, session: Session, *, - table_set: TableSet | None, + table_set: TableSet[IntoFrameT] | None, greeting: str | None, client: Callable[..., chatlas.Chat], history: bool | HistoryOptions, @@ -249,25 +249,21 @@ def mod_server( greeter: QueryChatGreeter, greeting_base: chatlas.Chat | None = None, greeting_tables: list[str] | None = None, - # TableSet erases each DataSource's frame type for uniform executor handling, - # so IntoFrameT can't be bound from any parameter here. There's no way to - # express the real return type, so callers (e.g. QueryChat.server()) must - # cast/annotate the result as ServerValues[IntoFrameT] themselves. -) -> ServerValues[Any]: +) -> ServerValues[IntoFrameT]: if not callable(client): raise TypeError("mod_server() requires a callable client factory.") - table_states: dict[str, TableState[Any]] = {} + table_states: dict[str, TableState[IntoFrameT]] = {} _current_table: ReactiveStringOrNone = ReactiveStringOrNone(None) def _make_table_state( - source: DataSource[Any], exec: QueryExecutor - ) -> TableState[Any]: + source: DataSource[IntoFrameT], exec: QueryExecutor + ) -> TableState[IntoFrameT]: table_sql = ReactiveStringOrNone(None) table_title = ReactiveStringOrNone(None) @reactive.calc - def filtered_df() -> Any: + def filtered_df() -> IntoFrameT: query = table_sql.get() if query: return exec.execute_query(query) @@ -478,7 +474,7 @@ def _on_history_restore(values: dict[str, Any]) -> None: df_warned = False @reactive.calc - def _multi_table_df() -> Any: + def _multi_table_df() -> IntoFrameT: nonlocal df_warned if not df_warned: df_warned = True diff --git a/pkg-py/src/querychat/_table_set.py b/pkg-py/src/querychat/_table_set.py index 978828ded..69ce4e343 100644 --- a/pkg-py/src/querychat/_table_set.py +++ b/pkg-py/src/querychat/_table_set.py @@ -4,7 +4,9 @@ from functools import cached_property from types import MappingProxyType -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Generic + +from narwhals.stable.v1.typing import IntoFrameT from ._query_executor import QueryExecutor, build_query_executor @@ -15,7 +17,7 @@ from ._system_prompt import QueryChatSystemPrompt -class TableSet: +class TableSet(Generic[IntoFrameT]): """ The tables a chat can query, plus the prompt and executor built from them. @@ -27,18 +29,18 @@ class TableSet: def __init__( self, - data_sources: Mapping[str, DataSource], + data_sources: Mapping[str, DataSource[IntoFrameT]], system_prompt: QueryChatSystemPrompt, ) -> None: if not data_sources: raise ValueError("TableSet requires at least one data source") - self._data_sources: Mapping[str, DataSource] = MappingProxyType( + self._data_sources: Mapping[str, DataSource[IntoFrameT]] = MappingProxyType( dict(data_sources) ) self._system_prompt = system_prompt @property - def data_sources(self) -> Mapping[str, DataSource]: + def data_sources(self) -> Mapping[str, DataSource[IntoFrameT]]: return self._data_sources @property From 5aa0a3e01c39e00d8b97c10e9c94b2abe7922b08 Mon Sep 17 00:00:00 2001 From: Carson Date: Sun, 13 Sep 2026 10:30:01 -0500 Subject: [PATCH 25/35] chore(r): remove dead in_shiny_session() helper Defined in utils-shiny.R but never called anywhere in the package. --- pkg-r/R/utils-shiny.R | 3 --- 1 file changed, 3 deletions(-) delete mode 100644 pkg-r/R/utils-shiny.R diff --git a/pkg-r/R/utils-shiny.R b/pkg-r/R/utils-shiny.R deleted file mode 100644 index 893a57a15..000000000 --- a/pkg-r/R/utils-shiny.R +++ /dev/null @@ -1,3 +0,0 @@ -in_shiny_session <- function() { - !is.null(shiny::getDefaultReactiveDomain()) # nocov -} From 8562f4f55aa7c5c1ab846ef3cb1e53452f23a85c Mon Sep 17 00:00:00 2001 From: Carson Date: Sun, 13 Sep 2026 10:30:01 -0500 Subject: [PATCH 26/35] docs(r): note cleanup=NA auto-close no longer disconnects caller connections The shiny::onStop() hook registered when QueryChat is created inside a running app still fires, but cleanup() now only closes what querychat created, so caller-supplied DBI connections stay open at app stop. --- pkg-r/NEWS.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pkg-r/NEWS.md b/pkg-r/NEWS.md index f76bd203a..f0be8d051 100644 --- a/pkg-r/NEWS.md +++ b/pkg-r/NEWS.md @@ -65,6 +65,8 @@ * `$cleanup()` follows one rule: querychat closes only what it created. `DBISource$cleanup()` and `TblSqlSource$cleanup()` no longer disconnect your connection; disconnect it yourself on shutdown. `DataFrameSource`/`PinSource` DuckDB connections are still closed. +* The automatic `$cleanup()` registered when `QueryChat` is created while a Shiny app is running (`cleanup = NA`, the default) no longer disconnects caller-supplied DBI connections at app stop, for the same reason. If you relied on that to close a connection you passed to `QueryChat$new()`, register your own `shiny::onStop(function() DBI::dbDisconnect(con))` (or disconnect when the session ends). Data frames are unaffected: the in-memory DuckDB connection querychat creates for them is still closed automatically. + * Adding a *new* table with `$add_table()`/`$add_tables()` after a session has started now warns instead of erroring; running sessions keep their tables and new sessions see the addition. Replacing or removing an existing table after a session has started still errors. # querychat 0.3.0 From 2bf21cb71eea705269e8842027331ccada06c6e2 Mon Sep 17 00:00:00 2001 From: Carson Date: Sun, 13 Sep 2026 10:50:55 -0500 Subject: [PATCH 27/35] fix(py): reject a second pin table at registration time Two PinSources used to pass validation but fail at query time with a misleading DuckDB CatalogException, since each pin queries through its own private connection and DataSourceExecutor delegates all queries to the first one. check_source_compatibility() now rejects the second pin up front, with a pointer to the shared-DuckDB-connection pattern. Drive-bys: build_query_executor() drops its dead isinstance-filtering comprehensions (validation already guarantees homogeneity), and TableSet.cleanup_executor() now resets the cached executor so a later access rebuilds it, matching R's TableSet. --- pkg-py/CHANGELOG.md | 2 ++ pkg-py/src/querychat/_query_executor.py | 24 ++++++++++++++------ pkg-py/src/querychat/_table_set.py | 11 +++++++-- pkg-py/tests/test_pin_source.py | 30 +++++++++++++++++++++++++ 4 files changed, 58 insertions(+), 9 deletions(-) diff --git a/pkg-py/CHANGELOG.md b/pkg-py/CHANGELOG.md index c6c8d1001..88428c9a5 100644 --- a/pkg-py/CHANGELOG.md +++ b/pkg-py/CHANGELOG.md @@ -17,6 +17,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 * `cleanup()` no longer closes a spec-resolved `.server(client=...)` override while its session is still running; it is closed when the session ends. +* Registering a second pins table with `add_table()` now raises a clear error at registration time instead of failing at query time: each pin queries through its own DuckDB connection, so cross-pin queries can't run. To combine a pin with other tables, register them in a shared DuckDB connection and pass that instead. + ## [0.8.0] - 2026-09-12 ### New features diff --git a/pkg-py/src/querychat/_query_executor.py b/pkg-py/src/querychat/_query_executor.py index a3f302bdd..b9460ae6e 100644 --- a/pkg-py/src/querychat/_query_executor.py +++ b/pkg-py/src/querychat/_query_executor.py @@ -3,7 +3,7 @@ from __future__ import annotations from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, cast import duckdb import narwhals.stable.v1 as nw @@ -268,6 +268,7 @@ def check_source_compatibility( IbisSource, SQLAlchemySource, ) + from ._pin_source import PinSource first_source = next(iter(existing.values())) @@ -278,6 +279,18 @@ def check_source_compatibility( f"Existing tables use {type(first_source).__name__}." ) + # Reached only when the existing sources are also PinSources: a second pin + # would validate here but fail at query time, since each pin queries + # through its own private connection and DataSourceExecutor delegates all + # queries to the first one. + if isinstance(new_source, PinSource): + raise ValueError( + f"Cannot add pin '{new_name}': only one pin table is supported per " + "chat. Each pin queries through its own DuckDB connection, so " + "cross-pin queries can't run. To combine a pin with other tables, " + "register them in a shared DuckDB connection and pass that instead." + ) + if isinstance(new_source, DataFrameSource) and isinstance( first_source, DataFrameSource ): @@ -322,6 +335,7 @@ def build_query_executor(sources: Mapping[str, DataSource]) -> QueryExecutor: """Pick the executor for a compatible group of sources.""" from ._datasource import DataFrameSource, PolarsLazySource + # After validation, every source has the same type as the first one. validate_source_group_compatibility(dict(sources)) if len(sources) == 1: @@ -330,12 +344,8 @@ def build_query_executor(sources: Mapping[str, DataSource]) -> QueryExecutor: first_source = next(iter(sources.values())) if isinstance(first_source, DataFrameSource): - return DuckDBExecutor( - {n: s for n, s in sources.items() if isinstance(s, DataFrameSource)} - ) + return DuckDBExecutor(cast("dict[str, DataFrameSource]", dict(sources))) if isinstance(first_source, PolarsLazySource): - return PolarsSQLExecutor( - {n: s for n, s in sources.items() if isinstance(s, PolarsLazySource)} - ) + return PolarsSQLExecutor(cast("dict[str, PolarsLazySource]", dict(sources))) return DataSourceExecutor(dict(sources)) diff --git a/pkg-py/src/querychat/_table_set.py b/pkg-py/src/querychat/_table_set.py index 69ce4e343..abaa65a89 100644 --- a/pkg-py/src/querychat/_table_set.py +++ b/pkg-py/src/querychat/_table_set.py @@ -60,6 +60,13 @@ def table_names(self) -> list[str]: return list(self.data_sources) def cleanup_executor(self) -> None: - """Close the executor if it was ever built. Never touches data sources.""" + """Close the executor if it was ever built. Never touches data sources. + + The cached executor is reset (matching R's ``TableSet``), so a later + access rebuilds it. + """ if self.executor_built: - self.executor.cleanup() + try: + self.executor.cleanup() + finally: + del self.__dict__["executor"] diff --git a/pkg-py/tests/test_pin_source.py b/pkg-py/tests/test_pin_source.py index 704cef77c..56c61262a 100644 --- a/pkg-py/tests/test_pin_source.py +++ b/pkg-py/tests/test_pin_source.py @@ -299,3 +299,33 @@ def test_clears_auto_description_on_source_change(self, board, sample_df): assert "Motor Trend Cars" not in prompt_after finally: qc.cleanup() + + +class TestMultiplePins: + """A second pin must fail at registration, not at query time.""" + + def test_second_pin_rejected_by_compatibility_check(self, board, sample_df): + from querychat._query_executor import check_source_compatibility + + board.pin_write(sample_df, "pin_a", type="parquet") + board.pin_write(sample_df, "pin_b", type="parquet") + first = PinSource(board, "pin_a") + second = PinSource(board, "pin_b") + try: + with pytest.raises(ValueError, match="only one pin"): + check_source_compatibility({"pin_a": first}, second, "pin_b") + finally: + first.cleanup() + second.cleanup() + + def test_add_table_rejects_second_pin(self, board, sample_df): + from querychat import QueryChat + + board.pin_write(sample_df, "pin_a", type="parquet") + board.pin_write(sample_df, "pin_b", type="parquet") + qc = QueryChat(board, "pin_a") + try: + with pytest.raises(ValueError, match="only one pin"): + qc.add_table(board, "pin_b") + finally: + qc.cleanup() From 2d520657f5f05012d6566247ebb4c5bdb6f2d192 Mon Sep 17 00:00:00 2001 From: Carson Date: Sun, 13 Sep 2026 10:50:56 -0500 Subject: [PATCH 28/35] fix(r): reject multiple pins and tables from different DBI connections check_source_compatibility() only compared class identity, so two pins (or two DBISources on different connections) validated but then failed at query time: all queries execute against one shared connection. A second pin, or a table from a different connection, now errors at registration time. This matches the Python package's same-engine / same-backend checks. DataFrameSource is exempt from the same-connection rule: it inherits DBISource but opens its own in-memory connection by design, and multi-table data frames are served by a shared DuckDBExecutor. To support the connection-identity check, the DBISource family gains a read-only public $conn active binding (private$conn renamed to private$.conn throughout DBISource, DataFrameSource, TblSqlSource, and PinSource). --- pkg-r/NEWS.md | 4 +++ pkg-r/R/DBISource.R | 35 ++++++++++++------- pkg-r/R/DataFrameSource.R | 10 +++--- pkg-r/R/PinSource.R | 8 ++--- pkg-r/R/QueryExecutor.R | 30 ++++++++++++++++ pkg-r/R/TblSqlSource.R | 12 +++---- pkg-r/man/DBISource.Rd | 7 ++++ pkg-r/tests/testthat/test-PinSource.R | 33 ++++++++++++++++- pkg-r/tests/testthat/test-QueryChat.R | 2 +- pkg-r/tests/testthat/test-QueryExecutor.R | 20 +++++++++++ .../tests/testthat/test-server_data_source.R | 2 +- 11 files changed, 132 insertions(+), 31 deletions(-) diff --git a/pkg-r/NEWS.md b/pkg-r/NEWS.md index f0be8d051..503ba8374 100644 --- a/pkg-r/NEWS.md +++ b/pkg-r/NEWS.md @@ -69,6 +69,10 @@ * Adding a *new* table with `$add_table()`/`$add_tables()` after a session has started now warns instead of erroring; running sessions keep their tables and new sessions see the addition. Replacing or removing an existing table after a session has started still errors. +* Registering a second pins table now fails at `$add_table()` time with a clear error instead of failing at query time: each pin queries through its own connection, so cross-pin queries can't run. To combine a pin with other tables, register them in a shared DuckDB connection and pass that instead. + +* `$add_table()`/`$add_tables()` now reject database tables that use a different DBI connection than the existing tables, matching the Python package. Mixed connections used to validate but then fail at query time, since all queries execute against one shared connection. + # querychat 0.3.0 ## New features diff --git a/pkg-r/R/DBISource.R b/pkg-r/R/DBISource.R index 610d77e93..723434d40 100644 --- a/pkg-r/R/DBISource.R +++ b/pkg-r/R/DBISource.R @@ -27,7 +27,7 @@ DBISource <- R6::R6Class( "DBISource", inherit = DataSource, private = list( - conn = NULL + .conn = NULL ), public = list( #' @description @@ -66,7 +66,7 @@ DBISource <- R6::R6Class( ) } - private$conn <- conn + private$.conn <- conn self$table_name <- table_name # Store original column names for validation @@ -86,15 +86,15 @@ DBISource <- R6::R6Class( #' @return A string identifying the database type get_db_type = function() { # Special handling for known database types - if (inherits(private$conn, "duckdb_connection")) { + if (inherits(private$.conn, "duckdb_connection")) { return("DuckDB") } - if (inherits(private$conn, "SQLiteConnection")) { + if (inherits(private$.conn, "SQLiteConnection")) { return("SQLite") } # Default to 'POSIX' if dbms name not found - conn_info <- DBI::dbGetInfo(private$conn) + conn_info <- DBI::dbGetInfo(private$.conn) dbms_name <- getElement(conn_info, "dbms.name") %||% "POSIX" # Remove ' SQL', if exists (SQL is already in the prompt) @@ -110,7 +110,7 @@ DBISource <- R6::R6Class( get_schema = function(categorical_threshold = 20, table_spec = NULL) { check_number_whole(categorical_threshold, min = 1) get_schema_impl( - private$conn, + private$.conn, self$table_name, categorical_threshold, table_spec = table_spec @@ -123,14 +123,14 @@ DBISource <- R6::R6Class( ) { check_number_whole(categorical_threshold, min = 1) details <- build_column_details_impl( - private$conn, + private$.conn, self$table_name, categorical_threshold, table_spec = table_spec ) list( text = format_schema_from_details( - as.character(DBI::dbQuoteIdentifier(private$conn, self$table_name)), + as.character(DBI::dbQuoteIdentifier(private$.conn, self$table_name)), details ), columns = details @@ -141,10 +141,10 @@ DBISource <- R6::R6Class( #' Get information about semantic views (if any) for the system prompt. #' @return A string with semantic view information, or empty string if none get_semantic_views_description = function() { - if (!is_snowflake_connection(private$conn)) { + if (!is_snowflake_connection(private$.conn)) { return("") } - views <- discover_semantic_views_impl(private$conn) + views <- discover_semantic_views_impl(private$.conn) if (length(views) == 0) { return("") } @@ -161,12 +161,12 @@ DBISource <- R6::R6Class( if (is.null(query) || !nzchar(query)) { query <- paste0( "SELECT * FROM ", - DBI::dbQuoteIdentifier(private$conn, self$table_name) + DBI::dbQuoteIdentifier(private$.conn, self$table_name) ) } check_query(query) - DBI::dbGetQuery(private$conn, query) + DBI::dbGetQuery(private$.conn, query) }, #' @description @@ -181,7 +181,7 @@ DBISource <- R6::R6Class( check_bool(require_all_columns) check_query(query) - rs <- DBI::dbSendQuery(private$conn, query) + rs <- DBI::dbSendQuery(private$.conn, query) df <- DBI::dbFetch(rs, n = 1) DBI::dbClearResult(rs) @@ -220,6 +220,15 @@ DBISource <- R6::R6Class( cleanup = function() { invisible(NULL) } + ), + active = list( + #' @field conn The DBI connection backing this source (read-only). + conn = function(value) { + if (!missing(value)) { + cli::cli_abort("{.field conn} is read-only.") + } + private$.conn + } ) ) diff --git a/pkg-r/R/DataFrameSource.R b/pkg-r/R/DataFrameSource.R index edded59f8..92eb13aa0 100644 --- a/pkg-r/R/DataFrameSource.R +++ b/pkg-r/R/DataFrameSource.R @@ -69,7 +69,7 @@ DataFrameSource <- R6::R6Class( self$table_name <- table_name private$colnames <- colnames(df) - private$conn <- new_dataframe_connection(df, table_name, engine) + private$.conn <- new_dataframe_connection(df, table_name, engine) }, #' @description @@ -77,11 +77,11 @@ DataFrameSource <- R6::R6Class( #' #' @return NULL (invisibly) cleanup = function() { - if (!is.null(private$conn) && DBI::dbIsValid(private$conn)) { - if (inherits(private$conn, "duckdb_connection")) { - DBI::dbDisconnect(private$conn, shutdown = TRUE) + if (!is.null(private$.conn) && DBI::dbIsValid(private$.conn)) { + if (inherits(private$.conn, "duckdb_connection")) { + DBI::dbDisconnect(private$.conn, shutdown = TRUE) } else { - DBI::dbDisconnect(private$conn) + DBI::dbDisconnect(private$.conn) } } invisible(NULL) diff --git a/pkg-r/R/PinSource.R b/pkg-r/R/PinSource.R index 148efddc5..062eb7d11 100644 --- a/pkg-r/R/PinSource.R +++ b/pkg-r/R/PinSource.R @@ -185,11 +185,11 @@ PinSource <- R6::R6Class( #' #' @return `NULL` (invisibly) cleanup = function() { - if (!is.null(private$conn) && DBI::dbIsValid(private$conn)) { - if (inherits(private$conn, "duckdb_connection")) { - DBI::dbDisconnect(private$conn, shutdown = TRUE) + if (!is.null(private$.conn) && DBI::dbIsValid(private$.conn)) { + if (inherits(private$.conn, "duckdb_connection")) { + DBI::dbDisconnect(private$.conn, shutdown = TRUE) } else { - DBI::dbDisconnect(private$conn) + DBI::dbDisconnect(private$.conn) } } invisible(NULL) diff --git a/pkg-r/R/QueryExecutor.R b/pkg-r/R/QueryExecutor.R index 440e81a88..a6492c846 100644 --- a/pkg-r/R/QueryExecutor.R +++ b/pkg-r/R/QueryExecutor.R @@ -255,6 +255,36 @@ check_source_compatibility <- function(existing_sources, new_source, new_name) { ) } + # Reached only when the existing sources are also PinSources: a second pin + # would validate here but fail at query time, since each pin queries + # through its own private connection and DataSourceExecutor delegates all + # queries to the first one. + if (inherits(new_source, "PinSource")) { + cli::cli_abort( + c( + "Cannot add pin {.val {new_name}}: only one pin table is supported per chat.", + "i" = "Each pin queries through its own connection, so cross-pin queries can't run.", + "i" = "To combine a pin with other tables, register them in a shared DuckDB connection and pass that instead." + ) + ) + } + + # DataFrameSource inherits DBISource but is exempt: each instance opens its + # own in-memory connection by design, and multi-table data frames are served + # by a shared DuckDBExecutor instead. + if ( + inherits(new_source, "DBISource") && + !inherits(new_source, "DataFrameSource") && + !identical(new_source$conn, first_source$conn) + ) { + cli::cli_abort( + c( + "Cannot add table {.val {new_name}}: all database tables must share the same connection.", + "i" = "Use {.fn $add_tables} to register tables from a single connection." + ) + ) + } + invisible(NULL) } diff --git a/pkg-r/R/TblSqlSource.R b/pkg-r/R/TblSqlSource.R index bec97515b..63846bf5e 100644 --- a/pkg-r/R/TblSqlSource.R +++ b/pkg-r/R/TblSqlSource.R @@ -56,7 +56,7 @@ TblSqlSource <- R6::R6Class( ) } - private$conn <- dbplyr::remote_con(tbl) + private$.conn <- dbplyr::remote_con(tbl) private$tbl <- tbl # Collect various signals to infer the table name @@ -106,7 +106,7 @@ TblSqlSource <- R6::R6Class( #' @return A string containing schema information formatted for LLM prompts get_schema = function(categorical_threshold = 20, table_spec = NULL) { get_schema_impl( - private$conn, + private$.conn, self$table_name, categorical_threshold, columns = colnames(private$tbl), @@ -120,7 +120,7 @@ TblSqlSource <- R6::R6Class( table_spec = NULL ) { details <- build_column_details_impl( - private$conn, + private$.conn, self$table_name, categorical_threshold, columns = colnames(private$tbl), @@ -129,7 +129,7 @@ TblSqlSource <- R6::R6Class( ) list( text = format_schema_from_details( - as.character(DBI::dbQuoteIdentifier(private$conn, self$table_name)), + as.character(DBI::dbQuoteIdentifier(private$.conn, self$table_name)), details ), columns = details @@ -143,7 +143,7 @@ TblSqlSource <- R6::R6Class( #' @return A data frame containing query results execute_query = function(query) { sql_query <- self$prep_query(query) - dplyr::tbl(private$conn, dplyr::sql(sql_query)) + dplyr::tbl(private$.conn, dplyr::sql(sql_query)) }, #' @description @@ -174,7 +174,7 @@ TblSqlSource <- R6::R6Class( sprintf( "WITH %s AS (\n%s\n)\n%s", - DBI::dbQuoteIdentifier(private$conn, self$table_name), + DBI::dbQuoteIdentifier(private$.conn, self$table_name), private$tbl_cte, query ) diff --git a/pkg-r/man/DBISource.Rd b/pkg-r/man/DBISource.Rd index bbb896529..c7d605083 100644 --- a/pkg-r/man/DBISource.Rd +++ b/pkg-r/man/DBISource.Rd @@ -31,6 +31,13 @@ DBI::dbDisconnect(con) \section{Super class}{ \code{\link[querychat:DataSource]{DataSource}} -> \code{DBISource} } +\section{Active bindings}{ + \if{html}{\out{
    }} + \describe{ + \item{\code{conn}}{The DBI connection backing this source (read-only).} + } + \if{html}{\out{
    }} +} \section{Methods}{ \subsection{Public methods}{ \itemize{ diff --git a/pkg-r/tests/testthat/test-PinSource.R b/pkg-r/tests/testthat/test-PinSource.R index e4b9d6d7e..f871db943 100644 --- a/pkg-r/tests/testthat/test-PinSource.R +++ b/pkg-r/tests/testthat/test-PinSource.R @@ -405,9 +405,40 @@ test_that("PinSource$cleanup() disconnects the connection it opened", { skip_if_not_installed("nanoparquet") ps <- local_pin_source(type = "parquet") - conn <- ps$.__enclos_env__$private$conn + conn <- ps$conn ps$cleanup() expect_false(DBI::dbIsValid(conn)) }) + +describe("PinSource multi-table registration", { + skip_if_not_installed("pins") + skip_if_not_installed("duckdb") + skip_if_not_installed("nanoparquet") + + it("rejects a second pin via check_source_compatibility", { + ps1 <- local_pin_source(name = "pin_a", type = "parquet") + ps2 <- local_pin_source(name = "pin_b", type = "parquet") + + expect_error( + check_source_compatibility(list(pin_a = ps1), ps2, "pin_b"), + "only one pin" + ) + }) + + it("QueryChat$add_table() rejects a second pin", { + board <- pins::board_temp() + suppressMessages( + pins::pin_write(board, mtcars[1:5, ], "pin_a", type = "parquet") + ) + suppressMessages( + pins::pin_write(board, mtcars[1:5, ], "pin_b", type = "parquet") + ) + + qc <- QueryChat$new(board, "pin_a") + withr::defer(qc$cleanup()) + + expect_error(qc$add_table(board, "pin_b"), "only one pin") + }) +}) diff --git a/pkg-r/tests/testthat/test-QueryChat.R b/pkg-r/tests/testthat/test-QueryChat.R index 2e25dfb00..d4c95c89a 100644 --- a/pkg-r/tests/testthat/test-QueryChat.R +++ b/pkg-r/tests/testthat/test-QueryChat.R @@ -1430,7 +1430,7 @@ describe("QueryChat table changes after a session has started", { qc$add_table(new_test_df(), "users", replace = TRUE) - expect_false(DBI::dbIsValid(old_source$.__enclos_env__$private$conn)) + expect_false(DBI::dbIsValid(old_source$conn)) }) it("does not mark sessions as started when $server() fails", { diff --git a/pkg-r/tests/testthat/test-QueryExecutor.R b/pkg-r/tests/testthat/test-QueryExecutor.R index ca922c74f..8e39e0729 100644 --- a/pkg-r/tests/testthat/test-QueryExecutor.R +++ b/pkg-r/tests/testthat/test-QueryExecutor.R @@ -185,4 +185,24 @@ describe("check_source_compatibility()", { check_source_compatibility(existing, dbi_source, "test_table") ) }) + + it("rejects DBISources on different connections", { + skip_if_not_installed("RSQLite") + + conn1 <- DBI::dbConnect(RSQLite::SQLite(), ":memory:") + withr::defer(DBI::dbDisconnect(conn1)) + conn2 <- DBI::dbConnect(RSQLite::SQLite(), ":memory:") + withr::defer(DBI::dbDisconnect(conn2)) + + DBI::dbWriteTable(conn1, "users", new_users_df()) + DBI::dbWriteTable(conn2, "test_table", new_test_df()) + + source1 <- DBISource$new(conn1, "users") + source2 <- DBISource$new(conn2, "test_table") + + expect_error( + check_source_compatibility(list(users = source1), source2, "test_table"), + "same connection" + ) + }) }) diff --git a/pkg-r/tests/testthat/test-server_data_source.R b/pkg-r/tests/testthat/test-server_data_source.R index a9434085f..78de4aaa7 100644 --- a/pkg-r/tests/testthat/test-server_data_source.R +++ b/pkg-r/tests/testthat/test-server_data_source.R @@ -28,7 +28,7 @@ start_server_session <- function(qc, ..., env = parent.frame()) { } source_conn_valid <- function(source) { - DBI::dbIsValid(source$.__enclos_env__$private$conn) + DBI::dbIsValid(source$conn) } describe("QueryChat$server(data_source = ) session isolation", { From fbe94a2a3da1f8cdd4b780685a80f4410ea988a4 Mon Sep 17 00:00:00 2001 From: Carson Date: Sun, 13 Sep 2026 10:50:56 -0500 Subject: [PATCH 29/35] docs(r): clarify cleanup= auto-runs at app stop, not session stop The hook is shiny::onStop() gated by shiny::isRunning(), so it fires when the app stops, not when a session ends, and applies whenever QueryChat is created while an app is running (e.g. top level of app.R). --- pkg-r/R/QueryChat.R | 7 ++++--- pkg-r/man/QueryChat.Rd | 7 ++++--- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/pkg-r/R/QueryChat.R b/pkg-r/R/QueryChat.R index dd964768c..dbb345658 100644 --- a/pkg-r/R/QueryChat.R +++ b/pkg-r/R/QueryChat.R @@ -388,9 +388,10 @@ QueryChat <- R6::R6Class( #' @param data_dict Optional data dictionary. A path to a YAML file, or a #' list of YAML file paths. See [read_data_dict()] for the expected format. #' @param cleanup Whether or not to automatically run `$cleanup()` when the - #' Shiny session/app stops. By default, cleanup only occurs if `QueryChat` - #' gets created within a Shiny session. Set to `TRUE` to always clean up, - #' or `FALSE` to never clean up automatically. + #' Shiny app stops. By default, cleanup only occurs if `QueryChat` gets + #' created while a Shiny app is running (for example, at the top level of + #' `app.R`). Set to `TRUE` to always clean up, or `FALSE` to never clean + #' up automatically. #' #' @return A new `QueryChat` object. initialize = function( diff --git a/pkg-r/man/QueryChat.Rd b/pkg-r/man/QueryChat.Rd index 19b4c4a98..22c511e24 100644 --- a/pkg-r/man/QueryChat.Rd +++ b/pkg-r/man/QueryChat.Rd @@ -226,9 +226,10 @@ format.} \item{\code{data_dict}}{Optional data dictionary. A path to a YAML file, or a list of YAML file paths. See \code{\link[=read_data_dict]{read_data_dict()}} for the expected format.} \item{\code{cleanup}}{Whether or not to automatically run \verb{$cleanup()} when the -Shiny session/app stops. By default, cleanup only occurs if \code{QueryChat} -gets created within a Shiny session. Set to \code{TRUE} to always clean up, -or \code{FALSE} to never clean up automatically.} +Shiny app stops. By default, cleanup only occurs if \code{QueryChat} gets +created while a Shiny app is running (for example, at the top level of +\code{app.R}). Set to \code{TRUE} to always clean up, or \code{FALSE} to never clean +up automatically.} } \if{html}{\out{
}} } From e3c0cde083257ab01bcb00659bc4b10e214e5a79 Mon Sep 17 00:00:00 2001 From: Carson Date: Sun, 13 Sep 2026 11:05:20 -0500 Subject: [PATCH 30/35] style(py): satisfy ruff TRY004/D213 in executor and TableSet The pin-group rejection stays a ValueError (it's a group constraint violation, not a wrong-argument-type error) with a noqa, matching the neighboring compatibility checks. --- pkg-py/src/querychat/_query_executor.py | 4 +++- pkg-py/src/querychat/_table_set.py | 3 ++- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/pkg-py/src/querychat/_query_executor.py b/pkg-py/src/querychat/_query_executor.py index b9460ae6e..b4874c077 100644 --- a/pkg-py/src/querychat/_query_executor.py +++ b/pkg-py/src/querychat/_query_executor.py @@ -284,7 +284,9 @@ def check_source_compatibility( # through its own private connection and DataSourceExecutor delegates all # queries to the first one. if isinstance(new_source, PinSource): - raise ValueError( + # ValueError like the neighboring checks: this is a group constraint + # violation, not a wrong-argument-type error (contra TRY004). + raise ValueError( # noqa: TRY004 f"Cannot add pin '{new_name}': only one pin table is supported per " "chat. Each pin queries through its own DuckDB connection, so " "cross-pin queries can't run. To combine a pin with other tables, " diff --git a/pkg-py/src/querychat/_table_set.py b/pkg-py/src/querychat/_table_set.py index abaa65a89..5484ab603 100644 --- a/pkg-py/src/querychat/_table_set.py +++ b/pkg-py/src/querychat/_table_set.py @@ -60,7 +60,8 @@ def table_names(self) -> list[str]: return list(self.data_sources) def cleanup_executor(self) -> None: - """Close the executor if it was ever built. Never touches data sources. + """ + Close the executor if it was ever built. Never touches data sources. The cached executor is reset (matching R's ``TableSet``), so a later access rebuilds it. From a249b05a26246631ce99ce745ba932d08b9f3cac Mon Sep 17 00:00:00 2001 From: Carson Date: Sun, 13 Sep 2026 11:05:20 -0500 Subject: [PATCH 31/35] docs(r): correct cleanup= wording - onStop() is session-scoped inside a session shiny::onStop() registers on the current session when one exists, so cleanup runs at session end for QueryChat objects created in the server function, and at app stop only for those created outside a session (e.g. top level of app.R). The previous commit's 'app stop' wording was wrong for the in-session case. --- pkg-r/R/QueryChat.R | 12 +++++++----- pkg-r/man/QueryChat.Rd | 12 +++++++----- 2 files changed, 14 insertions(+), 10 deletions(-) diff --git a/pkg-r/R/QueryChat.R b/pkg-r/R/QueryChat.R index dbb345658..482d96d2f 100644 --- a/pkg-r/R/QueryChat.R +++ b/pkg-r/R/QueryChat.R @@ -387,11 +387,13 @@ QueryChat <- R6::R6Class( #' format. #' @param data_dict Optional data dictionary. A path to a YAML file, or a #' list of YAML file paths. See [read_data_dict()] for the expected format. - #' @param cleanup Whether or not to automatically run `$cleanup()` when the - #' Shiny app stops. By default, cleanup only occurs if `QueryChat` gets - #' created while a Shiny app is running (for example, at the top level of - #' `app.R`). Set to `TRUE` to always clean up, or `FALSE` to never clean - #' up automatically. + #' @param cleanup Whether or not to automatically run `$cleanup()`. By + #' default, cleanup only occurs if `QueryChat` gets created while a Shiny + #' app is running: when created inside a session (e.g., in the server + #' function), cleanup runs when that session ends; when created outside + #' a session (e.g., at the top level of `app.R`), it runs when the app + #' stops. Set to `TRUE` to always clean up, or `FALSE` to never clean up + #' automatically. #' #' @return A new `QueryChat` object. initialize = function( diff --git a/pkg-r/man/QueryChat.Rd b/pkg-r/man/QueryChat.Rd index 22c511e24..3c114f7d3 100644 --- a/pkg-r/man/QueryChat.Rd +++ b/pkg-r/man/QueryChat.Rd @@ -225,11 +225,13 @@ used. See the package prompts directory for the default template format.} \item{\code{data_dict}}{Optional data dictionary. A path to a YAML file, or a list of YAML file paths. See \code{\link[=read_data_dict]{read_data_dict()}} for the expected format.} - \item{\code{cleanup}}{Whether or not to automatically run \verb{$cleanup()} when the -Shiny app stops. By default, cleanup only occurs if \code{QueryChat} gets -created while a Shiny app is running (for example, at the top level of -\code{app.R}). Set to \code{TRUE} to always clean up, or \code{FALSE} to never clean -up automatically.} + \item{\code{cleanup}}{Whether or not to automatically run \verb{$cleanup()}. By +default, cleanup only occurs if \code{QueryChat} gets created while a Shiny +app is running: when created inside a session (e.g., in the server +function), cleanup runs when that session ends; when created outside +a session (e.g., at the top level of \code{app.R}), it runs when the app +stops. Set to \code{TRUE} to always clean up, or \code{FALSE} to never clean up +automatically.} } \if{html}{\out{
}} } From 38e784dbde9338bc034ce8f2c698eca9f4e1ae0b Mon Sep 17 00:00:00 2001 From: Carson Date: Sun, 13 Sep 2026 11:45:04 -0500 Subject: [PATCH 32/35] docs(r): drop NEWS bullets describing fixes to unreleased behavior Multi-table and pins have not shipped in R yet, so bullets contrasting registration-time rejection with the old query-time failure describe a before-state no released version had. Also align the cleanup=NA bullet with the corrected session-vs-app stop wording. --- pkg-r/NEWS.md | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/pkg-r/NEWS.md b/pkg-r/NEWS.md index 503ba8374..1b8cb3379 100644 --- a/pkg-r/NEWS.md +++ b/pkg-r/NEWS.md @@ -65,13 +65,10 @@ * `$cleanup()` follows one rule: querychat closes only what it created. `DBISource$cleanup()` and `TblSqlSource$cleanup()` no longer disconnect your connection; disconnect it yourself on shutdown. `DataFrameSource`/`PinSource` DuckDB connections are still closed. -* The automatic `$cleanup()` registered when `QueryChat` is created while a Shiny app is running (`cleanup = NA`, the default) no longer disconnects caller-supplied DBI connections at app stop, for the same reason. If you relied on that to close a connection you passed to `QueryChat$new()`, register your own `shiny::onStop(function() DBI::dbDisconnect(con))` (or disconnect when the session ends). Data frames are unaffected: the in-memory DuckDB connection querychat creates for them is still closed automatically. +* The automatic `$cleanup()` registered when `QueryChat` is created while a Shiny app is running (`cleanup = NA`, the default) no longer disconnects caller-supplied DBI connections when the session or app stops, for the same reason. If you relied on that to close a connection you passed to `QueryChat$new()`, register your own `shiny::onStop(function() DBI::dbDisconnect(con))` (or disconnect when the session ends). Data frames are unaffected: the in-memory DuckDB connection querychat creates for them is still closed automatically. * Adding a *new* table with `$add_table()`/`$add_tables()` after a session has started now warns instead of erroring; running sessions keep their tables and new sessions see the addition. Replacing or removing an existing table after a session has started still errors. -* Registering a second pins table now fails at `$add_table()` time with a clear error instead of failing at query time: each pin queries through its own connection, so cross-pin queries can't run. To combine a pin with other tables, register them in a shared DuckDB connection and pass that instead. - -* `$add_table()`/`$add_tables()` now reject database tables that use a different DBI connection than the existing tables, matching the Python package. Mixed connections used to validate but then fail at query time, since all queries execute against one shared connection. # querychat 0.3.0 From 49c387499d3a85f0bb45d8244c0e0c28a7d2e1b1 Mon Sep 17 00:00:00 2001 From: Carson Date: Sun, 13 Sep 2026 12:12:53 -0500 Subject: [PATCH 33/35] docs: pin rejection message states the real limitation precisely With two pins, queries against the first pin worked fine; what failed was any execution touching the second pin (single-table included, not just joins), while schema and test_query still succeeded for it. Say 'only the first pin's table would be queryable' rather than the narrower 'cross-pin queries can't run'. --- pkg-py/CHANGELOG.md | 2 +- pkg-py/src/querychat/_query_executor.py | 7 ++++--- pkg-r/R/QueryExecutor.R | 2 +- 3 files changed, 6 insertions(+), 5 deletions(-) diff --git a/pkg-py/CHANGELOG.md b/pkg-py/CHANGELOG.md index 88428c9a5..171bd5a22 100644 --- a/pkg-py/CHANGELOG.md +++ b/pkg-py/CHANGELOG.md @@ -17,7 +17,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 * `cleanup()` no longer closes a spec-resolved `.server(client=...)` override while its session is still running; it is closed when the session ends. -* Registering a second pins table with `add_table()` now raises a clear error at registration time instead of failing at query time: each pin queries through its own DuckDB connection, so cross-pin queries can't run. To combine a pin with other tables, register them in a shared DuckDB connection and pass that instead. +* Registering a second pins table with `add_table()` now raises a clear error at registration time instead of failing at query time: each pin queries through its own DuckDB connection, so only the first pin's table would be queryable. To combine a pin with other tables, register them in a shared DuckDB connection and pass that instead. ## [0.8.0] - 2026-09-12 diff --git a/pkg-py/src/querychat/_query_executor.py b/pkg-py/src/querychat/_query_executor.py index b4874c077..d889647ea 100644 --- a/pkg-py/src/querychat/_query_executor.py +++ b/pkg-py/src/querychat/_query_executor.py @@ -288,9 +288,10 @@ def check_source_compatibility( # violation, not a wrong-argument-type error (contra TRY004). raise ValueError( # noqa: TRY004 f"Cannot add pin '{new_name}': only one pin table is supported per " - "chat. Each pin queries through its own DuckDB connection, so " - "cross-pin queries can't run. To combine a pin with other tables, " - "register them in a shared DuckDB connection and pass that instead." + "chat. Each pin queries through its own DuckDB connection, so only " + "the first pin's table would be queryable. To combine a pin with " + "other tables, register them in a shared DuckDB connection and " + "pass that instead." ) if isinstance(new_source, DataFrameSource) and isinstance( diff --git a/pkg-r/R/QueryExecutor.R b/pkg-r/R/QueryExecutor.R index a6492c846..6acca1f34 100644 --- a/pkg-r/R/QueryExecutor.R +++ b/pkg-r/R/QueryExecutor.R @@ -263,7 +263,7 @@ check_source_compatibility <- function(existing_sources, new_source, new_name) { cli::cli_abort( c( "Cannot add pin {.val {new_name}}: only one pin table is supported per chat.", - "i" = "Each pin queries through its own connection, so cross-pin queries can't run.", + "i" = "Each pin queries through its own connection, so only the first pin's table would be queryable.", "i" = "To combine a pin with other tables, register them in a shared DuckDB connection and pass that instead." ) ) From 6bd4845f37e3fb1610e304944f2ae0b7b73dae22 Mon Sep 17 00:00:00 2001 From: Carson Date: Sun, 13 Sep 2026 12:38:35 -0500 Subject: [PATCH 34/35] Run late-change check after the build step in add_table()/add_tables() (#311) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A rejected or failed add_table()/add_tables() call after a session has started no longer warns about — or locks the instance for — a change that never took effect. Mirrors the ordering R's add_table() already had at the remaining three call sites (R's add_tables(), Python's add_table() and add_tables()): build the new table set first (cleaning up newly-normalized sources on failure), then run the late-change check (cleaning up the new set's executor on rejection), and only then commit the change. Python's prompt-rebuild-with-history warning also moves after the check so a rejected call stays silent. Regression tests added per call site; NEWS/CHANGELOG updated. Fixes #311 --- pkg-py/CHANGELOG.md | 2 ++ pkg-py/src/querychat/_querychat_base.py | 33 +++++++++++++++++---- pkg-py/tests/test_base.py | 39 +++++++++++++++++++++++++ pkg-r/NEWS.md | 2 ++ pkg-r/R/QueryChat.R | 36 ++++++++++++++++++----- pkg-r/tests/testthat/test-QueryChat.R | 35 ++++++++++++++++++++++ 6 files changed, 135 insertions(+), 12 deletions(-) diff --git a/pkg-py/CHANGELOG.md b/pkg-py/CHANGELOG.md index 171bd5a22..035c6b1a6 100644 --- a/pkg-py/CHANGELOG.md +++ b/pkg-py/CHANGELOG.md @@ -15,6 +15,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 * Adding a *new* table with `add_table()`/`add_tables()` after a session has started now warns instead of raising; running sessions keep their tables and new sessions see the addition. Replacing or removing an existing table after a session has started still raises. +* A rejected or failed `add_table()`/`add_tables()` call (e.g. an incompatible source type) after a session has started no longer warns about the late change or otherwise affects the instance, since the change never took effect. (#311) + * `cleanup()` no longer closes a spec-resolved `.server(client=...)` override while its session is still running; it is closed when the session ends. * Registering a second pins table with `add_table()` now raises a clear error at registration time instead of failing at query time: each pin queries through its own DuckDB connection, so only the first pin's table would be queryable. To combine a pin with other tables, register them in a shared DuckDB connection and pass that instead. diff --git a/pkg-py/src/querychat/_querychat_base.py b/pkg-py/src/querychat/_querychat_base.py index 5bfa3e289..6477ed639 100644 --- a/pkg-py/src/querychat/_querychat_base.py +++ b/pkg-py/src/querychat/_querychat_base.py @@ -484,11 +484,9 @@ def add_table( exists = table_name in self._data_sources if exists and not replace: raise ValueError(f"Table '{table_name}' already exists") - self._check_late_change("add_table", destructive=exists) normalized = normalize_data_source(data_source, table_name) try: - self._warn_if_prompt_rebuilt_with_history() merged = dict(self._data_sources) merged[table_name] = normalized new_set = self._build_table_set(merged) @@ -497,6 +495,17 @@ def add_table( warn_on_failure(normalized.cleanup, "data source") raise + # Only after the change is known to be valid do we check whether it's + # too late to apply it, so a rejected/failed add_table() doesn't warn. + try: + self._check_late_change("add_table", destructive=exists) + except Exception: + warn_on_failure(new_set.cleanup_executor, "query executor") + if normalized is not data_source: + warn_on_failure(normalized.cleanup, "data source") + raise + + self._warn_if_prompt_rebuilt_with_history() old_source = self._data_sources.get(table_name) replaced = ( [old_source] @@ -595,7 +604,6 @@ def normalized_builder(name: str) -> DataSource: existing = [name for name in tables if name in self._data_sources] if existing and not replace: raise ValueError(f"Table '{existing[0]}' already exists") - self._check_late_change("add_tables", destructive=bool(existing)) if isinstance(include_in_greeting, bool): greeting_names = list(tables) if include_in_greeting else [] @@ -610,11 +618,26 @@ def normalized_builder(name: str) -> DataSource: ) normalized = {name: normalized_builder(name) for name in tables} - self._warn_if_prompt_rebuilt_with_history() merged = dict(self._data_sources) merged.update(normalized) - new_set = self._build_table_set(merged) + try: + new_set = self._build_table_set(merged) + except Exception: + for source in normalized.values(): + warn_on_failure(source.cleanup, "data source") + raise + # Only after the change is known to be valid do we check whether it's + # too late to apply it, so a rejected/failed add_tables() doesn't warn. + try: + self._check_late_change("add_tables", destructive=bool(existing)) + except Exception: + warn_on_failure(new_set.cleanup_executor, "query executor") + for source in normalized.values(): + warn_on_failure(source.cleanup, "data source") + raise + + self._warn_if_prompt_rebuilt_with_history() replaced = [ old for name in tables diff --git a/pkg-py/tests/test_base.py b/pkg-py/tests/test_base.py index 248fcacaa..436e0b752 100644 --- a/pkg-py/tests/test_base.py +++ b/pkg-py/tests/test_base.py @@ -516,6 +516,45 @@ def test_replace_after_sessions_started_raises(self, sample_df): qc.add_table(sample_df, "users", replace=True) assert qc._superseded_table_sets == [] + def test_failed_add_table_after_sessions_started_does_not_warn( + self, sample_df, multi_table_engine + ): + qc = QueryChatBase(sample_df, "users") + qc._sessions_started = True + + # Mixing a SQLAlchemy source with the existing DataFrame source fails + # validation; a change that never took effect must not warn about it. + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + with pytest.raises(ValueError, match="same type"): + qc.add_table(multi_table_engine, "orders") + assert not any("after a session has started" in str(x.message) for x in w) + assert qc.table_names() == ["users"] + + def test_rejected_replace_after_sessions_started_leaves_set_untouched( + self, sample_df + ): + qc = QueryChatBase(sample_df, "users") + old_set = qc._table_set + qc._sessions_started = True + + with pytest.raises(RuntimeError, match="replace or remove"): + qc.add_table(sample_df, "users", replace=True) + assert qc._table_set is old_set + + def test_failed_add_tables_after_sessions_started_does_not_warn( + self, sample_df, multi_table_engine + ): + qc = QueryChatBase(sample_df, "users") + qc._sessions_started = True + + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + with pytest.raises(ValueError, match="same type"): + qc.add_tables(multi_table_engine) + assert not any("after a session has started" in str(x.message) for x in w) + assert qc.table_names() == ["users"] + def test_remove_after_sessions_started_raises(self, sample_df): qc = QueryChatBase(sample_df, "users") qc.add_table(sample_df, "other") diff --git a/pkg-r/NEWS.md b/pkg-r/NEWS.md index 1b8cb3379..871c3e373 100644 --- a/pkg-r/NEWS.md +++ b/pkg-r/NEWS.md @@ -69,6 +69,8 @@ * Adding a *new* table with `$add_table()`/`$add_tables()` after a session has started now warns instead of erroring; running sessions keep their tables and new sessions see the addition. Replacing or removing an existing table after a session has started still errors. +* A rejected or failed `$add_table()`/`$add_tables()` call (e.g. an incompatible source type) after a session has started no longer warns about the late change or otherwise affects the instance, since the change never took effect. (#311) + # querychat 0.3.0 diff --git a/pkg-r/R/QueryChat.R b/pkg-r/R/QueryChat.R index 482d96d2f..2bb36df48 100644 --- a/pkg-r/R/QueryChat.R +++ b/pkg-r/R/QueryChat.R @@ -653,10 +653,6 @@ QueryChat <- R6::R6Class( "Table {.val {existing[[1]]}} already exists. Use {.code replace = TRUE} to replace." ) } - private$check_late_change( - "add_tables", - destructive = length(existing) > 0 - ) if ( !rlang::is_bool(include_in_greeting) && @@ -678,6 +674,11 @@ QueryChat <- R6::R6Class( lapply(tables, function(tbl) normalize_data_source(conn, tbl)), tables ) + cleanup_normalized <- function() { + for (source in normalized) { + warn_on_cleanup_failure(source$cleanup(), "data source") + } + } next_sources <- current for (table_name in tables) { next_sources[[table_name]] <- normalized[[table_name]] @@ -688,9 +689,30 @@ QueryChat <- R6::R6Class( } else { pending_description$description } - new_set <- private$build_table_set( - next_sources, - data_description = candidate_description + new_set <- tryCatch( + private$build_table_set( + next_sources, + data_description = candidate_description + ), + error = function(e) { + cleanup_normalized() + stop(e) + } + ) + + # Only after the change is known to be valid do we check whether it's + # too late to apply it, so a rejected/failed add_tables() doesn't warn, + # and doesn't commit the data description either. + tryCatch( + private$check_late_change( + "add_tables", + destructive = length(existing) > 0 + ), + error = function(e) { + warn_on_cleanup_failure(new_set$cleanup_executor(), "query executor") + cleanup_normalized() + stop(e) + } ) private$commit_data_description(pending_description) diff --git a/pkg-r/tests/testthat/test-QueryChat.R b/pkg-r/tests/testthat/test-QueryChat.R index d4c95c89a..3458429eb 100644 --- a/pkg-r/tests/testthat/test-QueryChat.R +++ b/pkg-r/tests/testthat/test-QueryChat.R @@ -1376,6 +1376,41 @@ describe("QueryChat$add_tables()", { multi_table_warns <- warns[grepl("Multiple tables", warns)] expect_length(multi_table_warns, 1L) }) + + it("failed $add_tables() after a session started does not warn about the late change", { + skip_if_no_dataframe_engine() + conn <- local_multi_table_conn() + qc <- local_querychat(new_users_df(), "users", greeting = "hi") + qc$.__enclos_env__$private$.sessions_started <- TRUE + + # Mixing DBI sources with the existing data-frame source fails validation; + # a change that never took effect must not warn about sessions missing it. + warns <- character(0) + withCallingHandlers( + expect_error(qc$add_tables(conn), "same type"), + warning = function(w) { + warns <<- c(warns, conditionMessage(w)) + invokeRestart("muffleWarning") + } + ) + expect_false(any(grepl("after a session has started", warns))) + expect_equal(qc$table_names(), "users") + }) + + it("rejected destructive $add_tables() after a session started leaves the instance untouched", { + conn <- local_multi_table_conn() + qc <- QueryChat$new(NULL, "placeholder", greeting = "Test") + suppressWarnings(qc$add_tables(conn)) + old_set <- qc$.__enclos_env__$private$.table_set + qc$.__enclos_env__$private$.sessions_started <- TRUE + + expect_error( + qc$add_tables(conn, tables = "orders", replace = TRUE), + "replace or remove" + ) + expect_identical(qc$.__enclos_env__$private$.table_set, old_set) + expect_setequal(qc$table_names(), c("orders", "customers")) + }) }) describe("QueryChat table changes after a session has started", { From 347bc2758470ab15d739b51c64aeeed8e9454c4b Mon Sep 17 00:00:00 2001 From: Carson Date: Sun, 13 Sep 2026 12:42:00 -0500 Subject: [PATCH 35/35] Reject DataSource/registration name mismatches in Python add_table()/.server() A DataSource registers its table under its own table_name; accepting a different registration name stored the source under an alias its underlying connection didn't have, so generated queries failed. Both Python entry points now reject the mismatch, matching R's long-standing checks (fifth Copilot review round on #310). --- pkg-py/CHANGELOG.md | 2 ++ pkg-py/src/querychat/_querychat_base.py | 9 +++++++++ pkg-py/src/querychat/_shiny.py | 14 ++++++++++++-- pkg-py/tests/test_base.py | 11 +++++++++++ pkg-py/tests/test_server_data_source.py | 12 ++++++++++++ 5 files changed, 46 insertions(+), 2 deletions(-) diff --git a/pkg-py/CHANGELOG.md b/pkg-py/CHANGELOG.md index 035c6b1a6..085ff509a 100644 --- a/pkg-py/CHANGELOG.md +++ b/pkg-py/CHANGELOG.md @@ -17,6 +17,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 * A rejected or failed `add_table()`/`add_tables()` call (e.g. an incompatible source type) after a session has started no longer warns about the late change or otherwise affects the instance, since the change never took effect. (#311) +* `add_table()` and `server(data_source=, table_name=)` now reject a `DataSource` whose own `table_name` differs from the registration name, matching R — previously the table was stored under an alias its underlying connection didn't have, so generated queries failed. + * `cleanup()` no longer closes a spec-resolved `.server(client=...)` override while its session is still running; it is closed when the session ends. * Registering a second pins table with `add_table()` now raises a clear error at registration time instead of failing at query time: each pin queries through its own DuckDB connection, so only the first pin's table would be queryable. To combine a pin with other tables, register them in a shared DuckDB connection and pass that instead. diff --git a/pkg-py/src/querychat/_querychat_base.py b/pkg-py/src/querychat/_querychat_base.py index 6477ed639..0d2535b78 100644 --- a/pkg-py/src/querychat/_querychat_base.py +++ b/pkg-py/src/querychat/_querychat_base.py @@ -484,6 +484,15 @@ def add_table( exists = table_name in self._data_sources if exists and not replace: raise ValueError(f"Table '{table_name}' already exists") + if ( + isinstance(data_source, DataSource) + and data_source.table_name != table_name + ): + raise ValueError( + f"data_source's own table name ('{data_source.table_name}') " + f"does not match the given table_name ('{table_name}'). " + "Pass a matching table_name." + ) normalized = normalize_data_source(data_source, table_name) try: diff --git a/pkg-py/src/querychat/_shiny.py b/pkg-py/src/querychat/_shiny.py index a49379740..b2f4cd18b 100644 --- a/pkg-py/src/querychat/_shiny.py +++ b/pkg-py/src/querychat/_shiny.py @@ -12,6 +12,7 @@ from shiny import App, Inputs, Outputs, Session, reactive, render, req, ui +from ._datasource import DataSource from ._icons import bs_icon from ._querychat_base import ( DEFAULT_TOOLS, @@ -45,7 +46,6 @@ from narwhals.stable.v1.typing import IntoFrame from ._data_dict import DataDict - from ._datasource import DataSource from ._table_accessor import TableAccessor from ._table_set import TableSet @@ -640,7 +640,7 @@ def page(self, title, *, id: Optional[str] = None, **kwargs): **kwargs, ) - def server( + def server( # noqa: PLR0912 self, *, data_source: IntoFrame | sqlalchemy.Engine | ibis.Table | None = None, @@ -729,6 +729,16 @@ def server( "table first with add_table()." ) check_table_name(resolved_table_name, data_source=data_source) + if ( + isinstance(data_source, DataSource) + and data_source.table_name != resolved_table_name + ): + raise ValueError( + f"data_source's own table name ('{data_source.table_name}') " + f"does not match the given table_name ('{resolved_table_name}'). " + "Pass a matching table_name, or omit it to use " + f"'{data_source.table_name}'." + ) session_source = normalize_data_source(data_source, resolved_table_name) try: table_set = self._build_table_set( diff --git a/pkg-py/tests/test_base.py b/pkg-py/tests/test_base.py index 436e0b752..11925a6f0 100644 --- a/pkg-py/tests/test_base.py +++ b/pkg-py/tests/test_base.py @@ -463,6 +463,17 @@ def test_system_prompt_built_exactly_once(self, ibis_backend_with_tables): assert len(multi_table_warns) == 1 +def test_add_table_rejects_data_source_name_mismatch(sample_df): + """Reject registering a DataSource under a name its connection doesn't have.""" + qc = QueryChatBase() + mismatched = DataFrameSource(sample_df, "orders") + + with pytest.raises(ValueError, match="does not match"): + qc.add_table(mismatched, "users") + mismatched.cleanup() + assert qc.table_names() == [] + + def test_history_stored_verbatim_no_default_substitution(): """ QueryChatBase stores history exactly as given -- including None -- so callers diff --git a/pkg-py/tests/test_server_data_source.py b/pkg-py/tests/test_server_data_source.py index eedb81d56..9e8fa02dc 100644 --- a/pkg-py/tests/test_server_data_source.py +++ b/pkg-py/tests/test_server_data_source.py @@ -106,6 +106,18 @@ def test_missing_table_name_raises(self, users_df, captured_mod_server): with pytest.raises(ValueError, match="table_name"): qc.server(data_source=users_df) + def test_data_source_name_mismatch_raises(self, users_df, captured_mod_server): + import narwhals.stable.v1 as nw + from querychat._datasource import DataFrameSource + + qc = shiny_mod.QueryChat() + mismatched = DataFrameSource(nw.from_native(users_df), "orders") + with pytest.raises(ValueError, match="does not match"): + qc.server(data_source=mismatched, table_name="users") + mismatched.cleanup() + assert qc.table_names() == [] + assert captured_mod_server == [] + def test_invalid_deferred_table_name_raises_at_construction(self): """A bad deferred name must fail fast, not at .server() registration.""" with pytest.raises(ValueError, match="must begin with a letter"):