Session-scoped data source lifecycle for R and Python

[cpsievert]

Summary

QueryChat$server(data_source = ) / .server(data_source=...) used to mutate the shared instance to register a per-session table, which meant concurrent sessions could step on each other's tables, and a live session's data source could get closed out from under it when another session registered a replacement. This lands the same fix in both R and Python: each session now gets its own read-only TableSet (tables + system prompt + query executor) built on top of the instance's tables, so a session never observes another session's or the instance's later changes, and cleanup only ever closes what querychat itself created.

Also: cleanup() now closes each resource independently and warns instead of raising if one fails, so a single broken connection can't leave the rest open. And connections/engines you pass in yourself (DBI connections, SQLAlchemy engines, tbl_sql backends) are never closed by querychat — only resources querychat opened itself (e.g. DataFrameSource's in-memory DuckDB) are.

Behavior changes

• $server(data_source = ) / .server(data_source=...) registers a table 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; the session's own data source is cleaned up when the session ends.
• Adding a new table with add_table()/add_tables() after a session has started now warns instead of raising. Replacing or removing an existing table after a session has started still raises — that's still illegal, since a running session may be using it.
• cleanup() follows one rule: querychat closes only what it created. DBISource/TblSqlSource/SQLAlchemySource no longer disconnect/dispose your connection or engine.

Fixes from Copilot review (four rounds)

• Python: add_table()/add_tables() no longer move a replaced table to the end of table_names() (a real ordering regression vs. both the prior Python behavior and current R behavior).
• R & Python: $server()/.server()/app_server()/Express's _ensure_server_started() only mark the instance as having a live session after the module server successfully starts — previously a failed registration could permanently lock the instance out of later add_table()/remove_table() calls.
• Python: every teardown step in _swap_table_set() and a session-scoped .server(client=...) override's close callback now go through warn_on_failure(), so one broken resource can't abort a replace/remove call after the new state is already committed, or raise out of session.on_ended.
• R & Python: rollback cleanup in add_table()'s and .server(data_source=)'s failure paths now goes through the warning helper too — previously, if cleaning up a just-created source itself failed, that cleanup error replaced the original build/compatibility error the caller was supposed to see.
• Python: TableSet.data_sources/system_prompt are now read-only properties, matching R's enforced immutability contract (previously they were plain assignable attributes).
• R: stale DBISource/TblSqlSource roxygen examples that said cleanup() disconnects the connection (and told users not to call it to keep the connection open) now show the correct pattern — call cleanup() (a no-op) and disconnect yourself.

Not applied (didn't hold up under verification, replied inline where a comment thread existed): Copilot's naming-convention suggestion to drop the leading underscore from two new private methods contradicts this repo's own convention (private methods on classes keep the underscore, even in already-private modules); its ownership-boundary claim about cleanup() closing caller-constructed DataSource instances conflates two different cleanup semantics — the ownership boundary is the underlying connection/engine, not who instantiated the wrapper, and DataFrameSource/PinSource always own their internal DuckDB connection regardless of construction site.

Filed as a separate follow-up (not blocking): #311 — the same "don't warn/lock before confirming success" fix only landed for R's add_table(); R's add_tables() and Python's add_table()/add_tables() still check the late-change guard before the build step that can fail.

Test plan

• make py-check (format, types, tests) — 959 passed
• make r-check-tests — 2167 passed, 0 failures
• make r-check-format
• make r-check-package (R CMD check) — 0 errors, 0 warnings, 1 pre-existing unrelated NOTE (modifyList import, handoff feature, predates this branch)
• CI green (R-CMD-check matrix, Python 3.10-3.14, e2e, docs)
• Regression tests added for every fix above

Known follow-ups (not blocking)

• Late-change warning fires before a failed add_table()/add_tables() call is known to succeed #311 — late-change warning ordering gap in R's add_tables() and Python's add_table()/add_tables().
• Python's TableSet.cleanup_executor() doesn't yet reset its cached executor after cleanup the way R's does (R: closing then rebuilding gives a fresh executor; Python currently would hand back the closed one). No live bug today (traced every caller; nothing rebuilds a cleaned-up set), but worth matching R's stricter contract.
• Python's build_query_executor() uses a local import inside the function where a top-level import would do (no circular import risk).

🤖 Generated with Claude Code

[Copilot]

🟡 Changes recommended

Unresolved moderate issues remain in Python ordering, teardown/ownership, and session state, along with R lifecycle issues and stale cleanup examples.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

This PR adds session-scoped data-source snapshots for R and Python QueryChat, improving concurrent-session isolation and resource cleanup.

Changes:

• Adds immutable per-session TableSet instances and executors.
• Preserves caller-owned connections and engines.
• Expands lifecycle, cleanup, isolation, and documentation coverage.

Open comments remain on Python table ordering, teardown and ownership handling, session-start state, corresponding R lifecycle issues, and cleanup documentation examples.

File summaries

File	Reviewed change
pkg-r/vignettes/build.Rmd	Documents session-scoped sources.
pkg-r/tests/testthat/test-TblSqlSource.R	Tests caller-owned connection cleanup.
pkg-r/tests/testthat/test-TableSet.R	Tests immutable table sets.
pkg-r/tests/testthat/test-server_data_source.R	Tests session isolation.
pkg-r/tests/testthat/test-QueryChat.R	Tests lifecycle and late table changes.
pkg-r/tests/testthat/test-querychat_module.R	Tests module table-set integration.
pkg-r/tests/testthat/test-PinSource.R	Tests source ownership behavior.
pkg-r/tests/testthat/test-DBISource.R	Tests connection ownership.
pkg-r/tests/testthat/helper-fixtures.R	Provides table-set test fixtures.
pkg-r/R/TblSqlSource.R	Updates non-owning cleanup behavior.
pkg-r/R/TableSet.R	Implements immutable table sets.
pkg-r/R/QueryExecutor.R	Updates executor lifecycle support.
pkg-r/R/QueryChatGreeter.R	Supports session-specific greetings.
pkg-r/R/QueryChat.R	Implements session lifecycle and table management.
pkg-r/R/querychat_module.R	Integrates session table sets.
pkg-r/R/PinSource.R	Updates source cleanup semantics.
pkg-r/R/DBISource.R	Preserves caller-owned connections.
pkg-r/R/DataSource.R	Documents resource ownership.
pkg-r/NEWS.md	Records lifecycle changes.
pkg-r/man/TblSqlSource.Rd	Documents TblSqlSource.
pkg-r/man/QueryChat.Rd	Documents QueryChat behavior.
pkg-r/man/PinSource.Rd	Documents PinSource behavior.
pkg-r/man/DBISource.Rd	Documents DBISource ownership.
pkg-r/man/DataSource.Rd	Documents data-source lifecycle.
pkg-py/tests/test_table_set.py	Tests Python table sets.
pkg-py/tests/test_state.py	Tests lifecycle state handling.
pkg-py/tests/test_shiny_module.py	Tests module table-set integration.
pkg-py/tests/test_server_data_source.py	Tests session-specific sources.
pkg-py/tests/test_querychat.py	Tests QueryChat lifecycle behavior.
pkg-py/tests/test_pin_source.py	Tests pin-source behavior.
pkg-py/tests/test_multi_table.py	Tests multi-table lifecycle.
pkg-py/tests/test_multi_table_frameworks.py	Tests multi-table framework support.
pkg-py/tests/test_deferred_shiny.py	Tests deferred Shiny behavior.
pkg-py/tests/test_datasource.py	Tests data-source ownership.
pkg-py/tests/test_cleanup.py	Tests cleanup behavior.
pkg-py/tests/test_base.py	Tests late table changes.
pkg-py/src/querychat/_table_set.py	Implements immutable table sets.
pkg-py/src/querychat/_shiny.py	Registers session-scoped sources.
pkg-py/src/querychat/_shiny_module.py	Consumes session table sets.
pkg-py/src/querychat/_querychat_greeter.py	Supports greeting snapshots.
pkg-py/src/querychat/_querychat_base.py	Implements table lifecycle and cleanup.
pkg-py/src/querychat/_query_executor.py	Updates executor handling.
pkg-py/src/querychat/_datasource.py	Defines ownership semantics.
pkg-py/docs/build.qmd	Documents session behavior.
pkg-py/CHANGELOG.md	Records lifecycle changes.

Review details

Suppressed comments (9)

pkg-py/src/querychat/_querychat_base.py:616

• This merge also removes every replaced key before re-adding it, so add_tables(..., replace=True) moves replaced tables to the end instead of preserving their documented registration order. That can change the primary table selected by the UI and flat accessors. Build a copy and update it so existing keys retain their positions and only new tables append.

        merged = {
            **{k: v for k, v in self._data_sources.items() if k not in normalized},
            **normalized,
        }

pkg-py/src/querychat/_querychat_base.py:616

• This merge has the same ordering regression for add_tables(): filtering out every name being added and then merging normalized moves replacements after all untouched tables. That violates the table_names() ordering contract and can change the first/primary source when an earlier table is replaced. Update a copy of the existing mapping instead, so replacements retain their positions and new names append.

        merged = {
            **{k: v for k, v in self._data_sources.items() if k not in normalized},
            **normalized,
        }

pkg-py/src/querychat/_querychat_base.py:189

• This new helper is defined in the already-private _querychat_base.py module, but its leading underscore makes the internal API inconsistent with the repository's Python naming guideline for private modules. Rename it without the prefix and update its call sites.

    def _check_late_change(self, method_name: str, *, destructive: bool) -> None:

pkg-py/src/querychat/_querychat_base.py:206

• This new helper is defined in the already-private _querychat_base.py module, but its leading underscore makes the internal API inconsistent with the repository's Python naming guideline for private modules. Rename it without the prefix and update its call sites.

    def _swap_table_set(self, new_set: TableSet, *, replaced: list[DataSource]) -> None:

pkg-py/src/querychat/_querychat_base.py:693

• This unconditionally calls cleanup() on every registered DataSource, including a DataFrameSource/PinSource that the caller constructed and passed to QueryChat. That closes a caller-owned DuckDB/backend resource, contradicting the DataSource.cleanup() contract and this method's docstring that caller-provided resources are never closed. Track source ownership when normalizing and only clean up sources created by querychat; apply the same ownership check in the replacement path.

                warn_on_failure(source.cleanup, "data source")

pkg-py/src/querychat/_querychat_base.py:222

• Executor cleanup failures are silently discarded here. The new cleanup behavior is to warn while continuing with the remaining teardown steps (and the R replacement path does so); use warn_on_failure() instead of suppressing the exception without any diagnostic.

        with contextlib.suppress(Exception):
            old_set.cleanup_executor()

pkg-py/src/querychat/_querychat_base.py:224

• This replacement path also closes every replaced DataSource, so replacing a caller-constructed DataFrameSource/PinSource shuts down a resource that querychat did not create. Apply the same ownership check as the main cleanup() path before invoking this callback.

            source.cleanup()

pkg-r/R/QueryChat.R:1313

• sessions_started is set before mod_server() has successfully initialized the session. If module setup fails (for example while building the executor or client), this call raises but the instance is permanently treated as having an active session, so later add_table() calls warn and replacements/removals are rejected even though no session started. Set the flag only after mod_server() returns successfully, as the state transition should represent a completed registration.

      private$.sessions_started <- TRUE

pkg-r/R/QueryChat.R:241

• The replacement path invokes cleanup() on every replaced DataSource, including a DataFrameSource/PinSource that the caller constructed and supplied to QueryChat. That closes a caller-owned DuckDB/backend resource, contradicting the DataSource$cleanup() contract and the documented ownership rule. Apply the same source-ownership check used for session-scoped sources before cleaning this resource.

        warn_on_cleanup_failure(source$cleanup(), "data source")

• Files reviewed: 45/45 changed files
• Comments generated: 6
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

🔵 Needs a closer look

Four moderate findings remain across Python cleanup/startup handling and stale R ownership examples.

Review details

Suppressed comments (4)

Previously missed (1) — in code that hasn't changed since the last review.

pkg-py/src/querychat/_shiny.py:1069

• _sessions_started is still set before mod_server() runs in the Express path. If module startup raises, this instance is left in the late-configuration state, so later add_table()/remove_table() calls are warned or rejected even though no session started (and _server_attempted also prevents a retry). Move this assignment until after mod_server() returns, as in QueryChat.server() and app_server().

pkg-py/src/querychat/_querychat_base.py:223

• Unlike the executor teardown immediately above, this calls each replaced source's cleanup() directly. If one cleanup raises, _swap_table_set() aborts after self._table_set has already been swapped and later replaced sources are not cleaned. Use the same best-effort warning wrapper for each source so replacement remains consistent and cleanup continues.

            source.cleanup()

pkg-r/R/DBISource.R:217

• The class examples near the top of this file (and the generated man/DBISource.Rd) still say that DBISource$cleanup() disconnects the connection and advise skipping it to keep the connection open. With this change cleanup is a no-op, so those examples contradict the new contract and can mislead users about connection ownership. Update the examples to call DBI::dbDisconnect() explicitly and regenerate the Rd file.

    #' No-op: the DBI connection is owned by the caller. Disconnect it
    #' yourself with `DBI::dbDisconnect()` when your application shuts down.

pkg-r/R/TblSqlSource.R:193

• The class example near the top of this file (and man/TblSqlSource.Rd) still says that $cleanup() closes the DB connection. The changed method is now explicitly a no-op for caller-owned tbl_sql connections, so the example contradicts the new ownership behavior. Update the example to disconnect the caller-owned connection explicitly and regenerate the Rd file.

    #' @description
    #' No-op: the connection behind the `tbl_sql` is owned by the caller.
    #'
    #' @return `NULL` (invisibly)

• Files reviewed: 45/45 changed files
• Comments generated: 0 new
• Review effort level: Lite

[Copilot]

🔵 Needs a closer look

Five unresolved moderate review issues affect rollback cleanup and Python TableSet immutability.

Review details

Suppressed comments (5)

pkg-py/src/querychat/_querychat_base.py:493

• If cleanup of a newly normalized source fails while building the replacement set, this direct call masks the original build/compatibility error, and the caller sees only the cleanup failure. Use the existing warning wrapper so the original error is preserved while cleanup remains best-effort.

                normalized.cleanup()

pkg-py/src/querychat/_shiny.py:739

• When a session table-set build fails, this rollback cleanup is called directly. A cleanup failure would replace the original registration/build error and contradict the best-effort cleanup behavior used by the callback below; use the warning wrapper and then re-raise the original error.

                    session_source.cleanup()

pkg-py/src/querychat/_table_set.py:38

• TableSet is documented as immutable/read-only, but these are ordinary assignable attributes: ts.system_prompt = ... and even ts.data_sources = ... both succeed. That can make a running session's prompt or table mapping diverge from its cached executor (the R TableSet rejects these assignments); keep the backing fields private and expose read-only properties, or otherwise reject reassignment.

        self.data_sources: Mapping[str, DataSource] = MappingProxyType(
            dict(data_sources)
        )
        self.system_prompt = system_prompt

pkg-r/R/QueryChat.R:551

• If cleanup of a newly normalized source fails while building the replacement set, this direct call masks the original build/compatibility error, and the caller sees only the cleanup failure. Route this rollback through the existing warning wrapper so the original error is preserved while cleanup remains best-effort.

          normalized$cleanup()

pkg-r/R/QueryChat.R:1250

• When a session table-set build fails, this rollback cleanup is called directly. A cleanup failure would replace the original registration/build error and contradict the best-effort cleanup behavior used by the callback below; warn and then re-raise the original error instead.

              session_source$cleanup()

• Files reviewed: 46/46 changed files
• Comments generated: 0 new
• Review effort level: Lite

[Copilot]

🔵 Needs a closer look

The broad cross-language lifecycle and cleanup changes warrant final human review.

Review details

• Files reviewed: 46/46 changed files
• Comments generated: 0 new
• Review effort level: Lite