diff --git a/pkg-py/CHANGELOG.md b/pkg-py/CHANGELOG.md index ec48860b3..085ff509a 100644 --- a/pkg-py/CHANGELOG.md +++ b/pkg-py/CHANGELOG.md @@ -5,6 +5,24 @@ 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. + +* 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. + ## [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. + ::: ::: 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/src/querychat/_query_executor.py b/pkg-py/src/querychat/_query_executor.py index f1c4e2287..d889647ea 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 @@ -19,6 +19,8 @@ from ._utils import check_query if TYPE_CHECKING: + from collections.abc import Mapping + from ._datasource import DataFrameSource, DataSource, PolarsLazySource @@ -266,6 +268,7 @@ def check_source_compatibility( IbisSource, SQLAlchemySource, ) + from ._pin_source import PinSource first_source = next(iter(existing.values())) @@ -276,6 +279,21 @@ 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): + # 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 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( first_source, DataFrameSource ): @@ -314,3 +332,23 @@ 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 + + # After validation, every source has the same type as the first one. + 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(cast("dict[str, DataFrameSource]", dict(sources))) + if isinstance(first_source, PolarsLazySource): + return PolarsSQLExecutor(cast("dict[str, PolarsLazySource]", dict(sources))) + + return DataSourceExecutor(dict(sources)) diff --git a/pkg-py/src/querychat/_querychat_base.py b/pkg-py/src/querychat/_querychat_base.py index cda17bc18..0d2535b78 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,18 +46,21 @@ ) 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 from pins.boards import BaseBoard from shinychat.types import HistoryOptions + from shiny import Session + from ._data_dict import DataDict from ._viz_tools import VisualizeData 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 +92,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[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[IntoFrameT]] = [] + self._sessions_started = False self._deferred_table_name: str | None = None self.tools = normalize_tools(tools, default=DEFAULT_TOOLS) @@ -123,20 +109,17 @@ 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._system_prompt: QueryChatSystemPrompt | None = None self._greeter: QueryChatGreeter | None = None if data_source is not None: @@ -151,32 +134,52 @@ 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[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[IntoFrameT]: + 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) - client_has_history = ( + 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[IntoFrameT]] + ) -> TableSet[IntoFrameT]: + 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 +188,43 @@ 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[IntoFrameT], *, 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. + 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 + warn_on_failure(old_set.cleanup_executor, "query executor") + for source in replaced: + 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.""" @@ -235,50 +235,34 @@ 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: 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(lambda: warn_on_failure(resolved.close, "chatlas client")) 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[IntoFrameT], *, base: chatlas.Chat | None = None, tools: TOOL_GROUPS | tuple[TOOL_GROUPS, ...] | MISSING_TYPE | None = MISSING, @@ -287,28 +271,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 +299,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 +347,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 +372,12 @@ def client_factory( prompt: str | Path, base: chatlas.Chat | None = None, *, - data_sources: dict[str, DataSource] | 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( 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 +410,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 +471,57 @@ 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 + 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( - "Table name must begin with a letter and contain only " - "letters, numbers, and underscores" + 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." ) - if table_name in self._data_sources and not replace: - raise ValueError(f"Table '{table_name}' already 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) + merged = dict(self._data_sources) + merged[table_name] = normalized + new_set = self._build_table_set(merged) + except Exception: + if normalized is not data_source: + 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: - cleanup_failed_staged_source(data_source, normalized) + 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) - 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 +566,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 +586,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 +609,10 @@ 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") if isinstance(include_in_greeting, bool): greeting_names = list(tables) if include_in_greeting else [] @@ -701,26 +627,33 @@ def normalized_builder(name: str) -> DataSource: ) normalized = {name: normalized_builder(name) for name in tables} + merged = dict(self._data_sources) + merged.update(normalized) + try: + new_set = self._build_table_set(merged) + except Exception: + for source in normalized.values(): + warn_on_failure(source.cleanup, "data source") + raise - 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() + # 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._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() + 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,15 +675,9 @@ 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()." - ) - if table_name not in self._data_sources: available = ", ".join(self._data_sources.keys()) raise ValueError(f"Table '{table_name}' not found. Available: {available}") @@ -760,79 +687,43 @@ 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] + 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() + 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. + + 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. - 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. + 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`). + 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 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) - self._owned_clients.clear() + 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") + 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( @@ -879,22 +770,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..b2f4cd18b 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 @@ -13,8 +12,16 @@ 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, 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, @@ -32,6 +39,7 @@ if TYPE_CHECKING: from pathlib import Path + import chatlas import ibis import narwhals.stable.v1 as nw import sqlalchemy @@ -39,6 +47,7 @@ from ._data_dict import DataDict 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,21 +416,21 @@ 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"]) + 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, greeting_base=None, greeting_tables=list(self.greeter.tables), ) + self._sessions_started = True @reactive.calc def active_table_name() -> str: @@ -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,10 @@ def server( ".server() must be called within an active Shiny session (i.e., within the server function). " ) + table_set: TableSet[IntoFrameT] | 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 +728,49 @@ 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) + 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( + {**self._data_sources, resolved_table_name: session_source} + ) + except Exception: + if session_source is not data_source: + warn_on_failure(session_source.cleanup, "session data source") + raise + if resolved_table_name not in greeting_tables: + greeting_tables.append(resolved_table_name) - 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 table_set is None: + table_set = self._require_table_set("server") + + if session_source is not None: + 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") + 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(base=resolved_client, **kwargs) + return self._create_session_client( + table_set, base=resolved_client, **kwargs + ) if enable_bookmarking is not None: warnings.warn( @@ -760,19 +796,19 @@ def create_session_client(**kwargs) -> chatlas.Chat: ) ) - self._mark_server_initialized(session) - return mod_server( + result = 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, ) + self._sessions_started = True + return result class QueryChatExpress(QueryChatBase[IntoFrameT]): @@ -1020,26 +1056,29 @@ 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 + table_set = self._table_set resolved_history: bool | HistoryOptions = ( self.history if self.history is not None @@ -1051,16 +1090,16 @@ 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, greeting_base=None, greeting_tables=list(self.greeter.tables), ) + self._sessions_started = True def sidebar( self, diff --git a/pkg-py/src/querychat/_shiny_module.py b/pkg-py/src/querychat/_shiny_module.py index b238b4191..3d7656d0a 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[IntoFrameT] | None, greeting: str | None, client: Callable[..., chatlas.Chat], history: bool | HistoryOptions, @@ -299,14 +299,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 +315,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 +347,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/src/querychat/_table_set.py b/pkg-py/src/querychat/_table_set.py new file mode 100644 index 000000000..5484ab603 --- /dev/null +++ b/pkg-py/src/querychat/_table_set.py @@ -0,0 +1,73 @@ +"""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, Generic + +from narwhals.stable.v1.typing import IntoFrameT + +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(Generic[IntoFrameT]): + """ + 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[IntoFrameT]], + system_prompt: QueryChatSystemPrompt, + ) -> None: + if not data_sources: + raise ValueError("TableSet requires at least one data source") + self._data_sources: Mapping[str, DataSource[IntoFrameT]] = MappingProxyType( + dict(data_sources) + ) + self._system_prompt = system_prompt + + @property + def data_sources(self) -> Mapping[str, DataSource[IntoFrameT]]: + return self._data_sources + + @property + def system_prompt(self) -> QueryChatSystemPrompt: + return self._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. + + The cached executor is reset (matching R's ``TableSet``), so a later + access rebuilds it. + """ + if self.executor_built: + try: + self.executor.cleanup() + finally: + del self.__dict__["executor"] diff --git a/pkg-py/tests/test_base.py b/pkg-py/tests/test_base.py index 2e2d425e7..11925a6f0 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, ) @@ -360,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"): @@ -370,12 +379,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: @@ -460,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 @@ -475,3 +489,116 @@ 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_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") + 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..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): + def test_string_spec_client_is_owned(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 == [] - ) + assert qc._base_client_owned is True - def test_cleanup_closes_owned_string_client(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() + 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_cleanup_closes_deferred_default_client(self, monkeypatch, sample_df): + 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,85 +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() - - -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 == [] + assert not override.provider._client.is_closed() - def test_replacement_without_live_session_cleans_immediately( - self, sample_df, ended_callbacks + def test_session_override_does_not_leak_into_instance( + self, monkeypatch, 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"] + 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") - 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 == [] + assert qc._base_client is base + assert qc._base_client_owned is False class TestCleanupDataSources: @@ -270,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 close chatlas client"), + pytest.warns(UserWarning, match="Failed to clean up chatlas client"), ): qc.cleanup() - assert override.provider._client.is_closed() - assert qc._owned_clients == [] 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 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..709f65e71 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,17 @@ 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) + + 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: @@ -251,9 +260,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") @@ -305,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): @@ -408,7 +430,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 +440,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 +471,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 +506,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, ) @@ -522,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 ): @@ -546,7 +566,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 +945,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 +974,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 +1013,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 +1054,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 +1155,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..56c61262a 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,41 @@ 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() + + +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() 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..9e8fa02dc 100644 --- a/pkg-py/tests/test_server_data_source.py +++ b/pkg-py/tests/test_server_data_source.py @@ -1,9 +1,11 @@ """Tests for QueryChat.server(data_source=...) parity with R (#300).""" +import warnings from unittest.mock import MagicMock, patch import pandas as pd import pytest +import querychat._querychat_base as base_mod import querychat._shiny as shiny_mod @@ -51,11 +53,13 @@ def end(self): @pytest.fixture -def fake_sessions(monkeypatch): - """Patch mod_server/get_current_session; each server() call gets a new session.""" +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(): @@ -65,7 +69,7 @@ def next_session(): monkeypatch.setattr(shiny_mod, "mod_server", fake_mod_server) monkeypatch.setattr(shiny_mod, "get_current_session", next_session) - return sessions + return sessions, calls class TestServerDataSourceRegistersDeferredTable: @@ -75,8 +79,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 +88,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): @@ -105,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"): @@ -122,7 +135,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 +146,253 @@ 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 +class TestServerDataSourceSessionIsolation: + def test_instance_tables_unchanged_by_session_registration( + self, 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"] + assert qc.table_names() == [] + assert captured_mod_server[0]["table_set"].table_names == ["users"] - def test_add_table_still_blocked_after_server_init( + def test_each_session_gets_its_own_source( self, users_df, other_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) + qc.server(data_source=other_users_df) + + 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_session_table_shadows_config_time_table( + self, users_df, other_users_df, captured_mod_server + ): + qc = shiny_mod.QueryChat(users_df, "users") + config_source = qc._data_sources["users"] - with pytest.raises(RuntimeError, match="Cannot add tables while a server session"): - qc.add_table(other_users_df, "other") + qc.server(data_source=other_users_df) + 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] -class TestServerDataSourceCleanupSafety: - def test_second_session_does_not_clean_up_first_sessions_source( + def test_sessions_do_not_see_each_others_tables( 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 = shiny_mod.QueryChat() + qc.add_table(users_df, "orders") - qc.server(data_source=users_df) - first_source = qc._data_sources["users"] + qc.server(data_source=other_users_df, table_name="returns") + qc.server(data_source=pd.DataFrame({"id": [7]}), table_name="orders") - with patch.object(first_source, "cleanup") as mock_cleanup: - qc.server(data_source=other_users_df) - mock_cleanup.assert_not_called() + assert captured_mod_server[0]["table_set"].table_names == ["orders", "returns"] + assert captured_mod_server[1]["table_set"].table_names == ["orders"] - def test_public_add_table_replace_still_cleans_up_old_source( - self, users_df, other_users_df + def test_greeting_tables_snapshot_is_per_session( + self, users_df, captured_mod_server ): - """ - 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 = shiny_mod.QueryChat(None, table_name="users") + qc.server(data_source=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() + assert captured_mod_server[0]["greeting_tables"] == ["users"] + assert qc.greeter.tables == [] - def test_second_session_does_not_clean_up_first_sessions_query_executor( + def test_greeter_build_client_forwards_table_set( 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") + qc.server(data_source=other_users_df) - with patch.object(first_executor, "cleanup") as mock_cleanup: - qc.server(data_source=other_users_df) - mock_cleanup.assert_not_called() + seen = [] - 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") + def factory(tables, prompt, base=None, *, table_set=None): + seen.append(table_set) + return MagicMock() - with patch.object(first_executor, "cleanup") as mock_cleanup: - qc.add_table(other_users_df, "users", replace=True) - mock_cleanup.assert_called_once() + 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] - def test_first_server_call_cleans_up_constructor_registered_source( - self, users_df, other_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 ): - """No session can still be using it, so cleanup-on-replace holds.""" + sessions, calls = session_runs + qc = shiny_mod.QueryChat(None, table_name="users") + qc.server(data_source=users_df) + 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") - constructor_source = qc._data_sources["users"] + qc.server() + config_source = qc._data_sources["users"] - with patch.object(constructor_source, "cleanup") as mock_cleanup: - qc.server(data_source=other_users_df) - mock_cleanup.assert_called_once() + assert sessions[0]._ended_callbacks == [] - def test_second_session_skips_cleanup_even_when_first_cleaned_up( - self, users_df, other_users_df, captured_mod_server + 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 ): + sessions, _calls = session_runs qc = shiny_mod.QueryChat(users_df, "users") - + config_source = qc._data_sources["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() + with patch.object(config_source, "cleanup") as cleanup: + sessions[0].end() + cleanup.assert_not_called() + qc.cleanup() + cleanup.assert_called_once() -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. - """ - - def test_ended_sessions_source_is_cleaned_up_on_replace( - self, users_df, other_users_df, fake_sessions + def test_failed_registration_closes_its_source_and_leaves_instance_untouched( + self, users_df, session_runs, monkeypatch ): - qc = shiny_mod.QueryChat(None, table_name="users") - qc.server(data_source=users_df) - first_source = qc._data_sources["users"] + import duckdb + import polars as pl - fake_sessions[0].end() + sessions, calls = session_runs + qc = shiny_mod.QueryChat(users_df, "users") + created = [] + real_normalize = shiny_mod.normalize_data_source - with patch.object(first_source, "cleanup") as mock_cleanup: - qc.server(data_source=other_users_df) - mock_cleanup.assert_called_once() + def spy(data_source, table_name): + source = real_normalize(data_source, table_name) + created.append(source) + return source - 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"] + monkeypatch.setattr(shiny_mod, "normalize_data_source", spy) - fake_sessions[0].end() # s1 ends; s2 still live + with pytest.raises(ValueError, match="same DataFrame backend"): + qc.server(data_source=pl.DataFrame({"id": [1]}), table_name="other") - 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() + (session_source,) = created + with pytest.raises(duckdb.ConnectionException): + session_source.execute_query("SELECT 1") + assert qc.table_names() == ["users"] + assert sessions[0]._ended_callbacks == [] - def test_add_table_allowed_once_all_sessions_have_ended( - self, users_df, other_users_df, fake_sessions - ): - qc = shiny_mod.QueryChat(None, table_name="users") - qc.server(data_source=users_df) + qc.server() + assert calls[-1]["table_set"].table_names == ["users"] - fake_sessions[0].end() + 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 - qc.add_table(other_users_df, "other") # must not raise - assert qc.table_names() == ["users", "other"] + _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") -class TestServerDataSourceGreetingSnapshot: - def test_server_passes_greeting_tables_snapshot_to_mod_server( - self, users_df, captured_mod_server + def test_client_resolution_failure_still_cleans_up_session_source( + self, users_df, session_runs, monkeypatch ): """ - Greeting generation runs lazily, after a later session may have - mutated the live greeter.tables -- hence the call-time snapshot. + 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") - qc.server(data_source=users_df) + created = [] + real_normalize = shiny_mod.normalize_data_source - assert captured_mod_server[0]["greeting_tables"] == ["users"] + 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) -class TestServerDataSourceMixedWithConfigTimeAddTable: - def test_unnamed_registration_replaces_config_time_table( - self, users_df, other_users_df, captured_mod_server - ): - qc = shiny_mod.QueryChat() - qc.add_table(users_df, "orders") + with pytest.raises(ValueError, match="not a known chatlas provider"): + qc.server(data_source=users_df, client="not-a-real-provider") - qc.server(data_source=other_users_df) + (session_source,) = created + assert sessions[0]._ended_callbacks != [] - # 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] + sessions[0].end() + with pytest.raises(duckdb.ConnectionException): + session_source.execute_query("SELECT 1") - def test_replacing_config_time_table_on_first_server_call_cleans_it_up( - self, users_df, other_users_df, captured_mod_server + def test_mod_server_failure_does_not_mark_sessions_started( + self, users_df, 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"] + """A failed mod_server() call must not lock out later add_table()/remove_table().""" + monkeypatch.setattr(shiny_mod, "get_current_session", lambda: MagicMock()) - with patch.object(config_source, "cleanup") as mock_cleanup: - qc.server(data_source=other_users_df) - mock_cleanup.assert_called_once() + def failing_mod_server(*args, **kwargs): + raise RuntimeError("boom") - 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") + monkeypatch.setattr(shiny_mod, "mod_server", failing_mod_server) - qc.server(data_source=other_users_df, table_name="returns") + qc = shiny_mod.QueryChat(users_df, "users") + with pytest.raises(RuntimeError, match="boom"): + qc.server() - 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] + 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) - def test_later_session_snapshot_includes_earlier_sessions_table( - self, users_df, other_users_df, captured_mod_server + def test_client_override_close_failure_warns_on_session_end( + self, users_df, session_runs, monkeypatch ): - """The registry is shared and cumulative across sessions.""" - qc = shiny_mod.QueryChat() - qc.add_table(users_df, "orders") + """ + 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) - # 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 = 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 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 diff --git a/pkg-py/tests/test_table_set.py b/pkg-py/tests/test_table_set.py new file mode 100644 index 000000000..2a9c01855 --- /dev/null +++ b/pkg-py/tests/test_table_set.py @@ -0,0 +1,88 @@ +"""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_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) + 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") diff --git a/pkg-r/NEWS.md b/pkg-r/NEWS.md index f792aeb51..871c3e373 100644 --- a/pkg-r/NEWS.md +++ b/pkg-r/NEWS.md @@ -61,6 +61,17 @@ * Fixed a module-namespace desync 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. + +* 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. + +* 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 ## New features diff --git a/pkg-r/R/DBISource.R b/pkg-r/R/DBISource.R index 725c02d8c..723434d40 100644 --- a/pkg-r/R/DBISource.R +++ b/pkg-r/R/DBISource.R @@ -18,16 +18,16 @@ #' # 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( "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) @@ -213,15 +213,22 @@ 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) } + ), + 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/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/PinSource.R b/pkg-r/R/PinSource.R index 84bbde013..062eb7d11 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/R/QueryChat.R b/pkg-r/R/QueryChat.R index 8b3b560a3..2bb36df48 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,52 +106,145 @@ 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) }, - auto_fill_data_description = function(sources = private$.data_sources) { + # 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) { - return() + return(private$.data_description) } - if (private$.data_description_mode == "inferred") { - private$.data_description <- NULL - private$.data_description_mode <- "empty" + desc <- sources[[1]]$get_data_description() + if (nzchar(desc %||% "")) { + return(desc) } - if (private$.data_description_mode == "empty") { - desc <- sources[[1]]$get_data_description() - if (nzchar(desc %||% "")) { - private$.data_description <- desc - private$.data_description_mode <- "inferred" - } + NULL + }, + + # 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(NULL) + } + if (length(sources) != 1) { + return(NULL) + } + 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) }, - build_system_prompt = function(data_sources = NULL) { - sources <- data_sources %||% private$.data_sources + 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 + ) { 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)) + } + warn_on_cleanup_failure(old_set$cleanup_executor(), "query executor") + for (source in replaced) { + warn_on_cleanup_failure(source$cleanup(), "data source") + } + invisible(NULL) }, create_session_client = function( + table_set, client_spec = NULL, tools = NA, handoff_available = FALSE, @@ -168,7 +262,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 +272,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, @@ -298,10 +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 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. + #' @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( @@ -381,9 +473,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 +509,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 +526,16 @@ 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." ) } + if ( is_data_source(data_source) && !identical(data_source$table_name, table_name) @@ -452,41 +547,56 @@ 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) + cleanup_normalized <- function() { + if (!inherits(data_source, "DataSource")) { + warn_on_cleanup_failure(normalized$cleanup(), "data source") + } + } + next_sources <- current next_sources[[table_name]] <- normalized + 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, + data_description = candidate_description + ), + error = function(e) { + cleanup_normalized() + stop(e) + } + ) - private$auto_fill_data_description(next_sources) + # 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, + # and doesn't commit the data description either. tryCatch( - { - private$build_system_prompt(data_sources = next_sources) - }, + private$check_late_change("add_table", destructive = exists), error = function(e) { - if (!inherits(data_source, "DataSource")) { - normalized$cleanup() - } + warn_on_cleanup_failure(new_set$cleanup_executor(), "query executor") + cleanup_normalized() stop(e) } ) + private$commit_data_description(pending_description) - 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 +610,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 +631,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,13 +643,15 @@ 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." + ) } if ( @@ -562,47 +674,67 @@ 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]] + cleanup_normalized <- function() { + for (source in normalized) { + warn_on_cleanup_failure(source$cleanup(), "data source") + } } - - next_sources <- private$.data_sources + next_sources <- current for (table_name in tables) { next_sources[[table_name]] <- normalized[[table_name]] } + 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, + data_description = candidate_description + ), + error = function(e) { + cleanup_normalized() + stop(e) + } + ) - private$auto_fill_data_description(next_sources) - private$build_system_prompt(data_sources = next_sources) + # 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) + 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 + private$swap_table_set(new_set, replaced = replaced) - if (!is.null(private$.query_executor)) { - tryCatch(private$.query_executor$cleanup(), error = function(e) NULL) - private$.query_executor <- NULL - } - - 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 +742,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 +810,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 +879,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 +1183,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 +1230,14 @@ QueryChat <- R6::R6Class( ) } + 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 +1250,63 @@ 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")) { + warn_on_cleanup_failure( + session_source$cleanup(), + "session data source" + ) + } + 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 + # 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" + ) + if (owns_session_source) { + warn_on_cleanup_failure( + session_source$cleanup(), + "session data source" + ) + } + }) } resolved_client_spec <- client %||% private$.client_spec @@ -1142,6 +1314,7 @@ QueryChat <- R6::R6Class( create_session_client <- function(...) { private$create_session_client( + table_set = table_set, client_spec = base_client, ... ) @@ -1167,15 +1340,16 @@ QueryChat <- R6::R6Class( result <- 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 ) + private$.sessions_started <- TRUE result }, @@ -1193,15 +1367,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 +1402,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 +1437,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 +1654,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/QueryExecutor.R b/pkg-r/R/QueryExecutor.R index 1347e3d88..6acca1f34 100644 --- a/pkg-r/R/QueryExecutor.R +++ b/pkg-r/R/QueryExecutor.R @@ -255,5 +255,45 @@ 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 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." + ) + ) + } + + # 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) +} + +# 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..2fe8d4f98 --- /dev/null +++ b/pkg-r/R/TableSet.R @@ -0,0 +1,81 @@ +# 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)) { + tryCatch( + private$.executor$cleanup(), + finally = { + private$.executor <- NULL + } + ) + } + 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/R/TblSqlSource.R b/pkg-r/R/TblSqlSource.R index 7ab55e7ac..63846bf5e 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( @@ -55,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 @@ -105,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), @@ -119,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), @@ -128,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 @@ -142,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 @@ -173,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 ) @@ -188,9 +189,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/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/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 -} diff --git a/pkg-r/man/DBISource.Rd b/pkg-r/man/DBISource.Rd index 539af25a8..c7d605083 100644 --- a/pkg-r/man/DBISource.Rd +++ b/pkg-r/man/DBISource.Rd @@ -23,14 +23,21 @@ 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}{ \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{ @@ -204,14 +211,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/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
}} } @@ -242,6 +245,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 +284,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 +323,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 +582,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 +640,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/man/TblSqlSource.Rd b/pkg-r/man/TblSqlSource.Rd index f15a1cccc..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}{ @@ -221,14 +222,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/helper-fixtures.R b/pkg-r/tests/testthat/helper-fixtures.R index 6ee167ee9..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 @@ -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-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-PinSource.R b/pkg-r/tests/testthat/test-PinSource.R index 2230dc31f..f871db943 100644 --- a/pkg-r/tests/testthat/test-PinSource.R +++ b/pkg-r/tests/testthat/test-PinSource.R @@ -398,3 +398,47 @@ 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$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 0b810131c..3458429eb 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)", { @@ -1262,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()", { @@ -1340,29 +1362,244 @@ describe("QueryChat$add_tables()", { ) }) - it("calling after server initialization raises error", { + it("system prompt built exactly once for multiple tables", { conn <- local_multi_table_conn() qc <- QueryChat$new(NULL, "placeholder", greeting = "Test") - qc$.__enclos_env__$private$.server_initialized <- TRUE - expect_error( + warns <- character(0) + withCallingHandlers( qc$add_tables(conn), - "after server initialization" + warning = function(w) { + warns <<- c(warns, conditionMessage(w)) + invokeRestart("muffleWarning") + } ) + multi_table_warns <- warns[grepl("Multiple tables", warns)] + expect_length(multi_table_warns, 1L) }) - it("system prompt built exactly once for multiple tables", { + 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 <- QueryChat$new(NULL, "placeholder", greeting = "Test") + 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( - qc$add_tables(conn), + expect_error(qc$add_tables(conn), "same type"), warning = function(w) { warns <<- c(warns, conditionMessage(w)) invokeRestart("muffleWarning") } ) - multi_table_warns <- warns[grepl("Multiple tables", warns)] - expect_length(multi_table_warns, 1L) + 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", { + 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$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", { + # 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) + }) + + 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 + ) }) }) 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-TableSet.R b/pkg-r/tests/testthat/test-TableSet.R new file mode 100644 index 000000000..4011062c3 --- /dev/null +++ b/pkg-r/tests/testthat/test-TableSet.R @@ -0,0 +1,104 @@ +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")) + 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", { + 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))) +}) diff --git a/pkg-r/tests/testthat/test-TblSqlSource.R b/pkg-r/tests/testthat/test-TblSqlSource.R index d1abcedf2..02b6e76b3 100644 --- a/pkg-r/tests/testthat/test-TblSqlSource.R +++ b/pkg-r/tests/testthat/test-TblSqlSource.R @@ -386,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)) +}) 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..78de4aaa7 --- /dev/null +++ b/pkg-r/tests/testthat/test-server_data_source.R @@ -0,0 +1,281 @@ +# 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$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("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("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" + ) + }) + + 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") + }) +}) 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).