feat(py): close querychat-created chatlas clients in cleanup()

pkg-py/CHANGELOG.md

@@ -5,6 +5,12 @@ 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]
+
+### New features
+
+* `QueryChat.cleanup()` now also closes the chatlas client, releasing its provider resources (HTTP connection pools, database sessions) — but only when querychat created the client itself (i.e., `client` was `None` or a string spec like `"openai/gpt-4o"`). A user-supplied `chatlas.Chat` instance is never closed by querychat; its lifecycle remains the caller's responsibility. In long-lived applications, call `qc.cleanup()` when the app shuts down.
+
 ## [0.7.0] - 2026-07-10
 
 ### New features

pkg-py/src/querychat/_querychat_base.py

@@ -113,6 +113,12 @@ def __init__(
         self._extra_instructions = extra_instructions
         self._categorical_threshold = categorical_threshold
 
+        # Ownership rule: querychat closes the chatlas client only if it
+        # created it. A `None` client is resolved later from the
+        # QUERYCHAT_CLIENT env var or the "openai" default, and a string spec
+        # goes through ChatAuto -- both are querychat-created. A
+        # user-supplied Chat instance is never closed by querychat.
+        self._owns_client = client is None or isinstance(client, str)
         self._base_client: chatlas.Chat | None = (
             resolve_client(client) if client is not None else None
         )
@@ -684,11 +690,28 @@ def _mark_server_initialized(self) -> None:
         self._server_initialized = True
 
     def cleanup(self) -> None:
-        """Clean up resources associated with all data sources."""
+        """
+        Clean up resources held by this object.
+
+        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"`). A user-supplied `chatlas.Chat` instance is never
+        closed here -- its lifecycle remains the caller's responsibility.
+
+        Note that session, console, and greeter clients are deepcopy clones
+        that share the base client's provider, so closing the base client
+        releases their underlying resources as well.
+
+        Safe to call multiple times. In long-lived applications, call this
+        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()
+        if self._owns_client and self._base_client is not None:
+            self._base_client.close()
 
 
 def normalize_data_source(

pkg-py/tests/test_cleanup.py

@@ -0,0 +1,93 @@
+"""Tests for QueryChatBase.cleanup() resource lifecycle."""
+
+from unittest.mock import patch
+
+import chatlas
+import pandas as pd
+import pytest
+from chatlas import ChatOpenAI
+from querychat._querychat_base import QueryChatBase
+
+
+@pytest.fixture
+def sample_df():
+    return pd.DataFrame(
+        {
+            "id": [1, 2, 3],
+            "name": ["Alice", "Bob", "Charlie"],
+            "age": [25, 30, 35],
+        },
+    )
+
+
+class TestClientOwnership:
+    """querychat closes the chatlas client only if it created it."""
+
+    def test_owns_client_flag(self, monkeypatch, sample_df):
+        monkeypatch.setenv("OPENAI_API_KEY", "sk-dummy-key-for-testing")
+
+        # String spec: querychat-created
+        assert QueryChatBase(sample_df, "users", client="openai")._owns_client
+        # None (deferred default/env): querychat-created
+        assert QueryChatBase(sample_df, "users")._owns_client
+        # User-supplied instance: not owned
+        assert not QueryChatBase(
+            sample_df, "users", client=ChatOpenAI()
+        )._owns_client
+
+    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_cleanup_closes_deferred_default_client(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.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()
+        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
+        qc.cleanup()
+        assert clone.provider._client.is_closed()
+
+
+class TestCleanupDataSources:
+    """Existing executor/source cleanup behavior is preserved."""
+
+    def test_cleanup_closes_data_source(self, monkeypatch, sample_df):
+        monkeypatch.setenv("OPENAI_API_KEY", "sk-dummy-key-for-testing")
+        qc = QueryChatBase(sample_df, "users", client="openai")
+        source = qc._data_sources["users"]
+        with patch.object(source, "cleanup") as mock_cleanup:
+            qc.cleanup()
+            mock_cleanup.assert_called_once()
+
+    def test_cleanup_is_idempotent(self, monkeypatch, sample_df):
+        monkeypatch.setenv("OPENAI_API_KEY", "sk-dummy-key-for-testing")
+        qc = QueryChatBase(sample_df, "users", client="openai")
+        qc.cleanup()
+        qc.cleanup()  # should not raise

pkg-r/R/TblSqlSource.R

@@ -171,10 +171,15 @@ TblSqlSource <- R6::R6Class(
         return(query)
       }
 
+      cte_body <- qualify_cte_table_refs(
+        private$conn,
+        as.character(private$tbl_cte)
+      )
+
       sprintf(
         "WITH %s AS (\n%s\n)\n%s",
         DBI::dbQuoteIdentifier(private$conn, self$table_name),
-        private$tbl_cte,
+        cte_body,
         query
       )
     },
@@ -196,3 +201,36 @@ TblSqlSource <- R6::R6Class(
     }
   )
 )
+
+# Get the current schema name for a DBI connection.
+# Tries `SELECT current_schema()` (DuckDB, PostgreSQL, etc.) and falls
+# back to "main" for databases that don't support it (SQLite).
+# @param conn A DBI connection
+# @return A character string with the schema name
+# @noRd
+get_current_schema <- function(conn) {
+  tryCatch(
+    DBI::dbGetQuery(conn, "SELECT current_schema()")[[1]],
+    error = function(e) "main"
+  )
+}
+
+# Schema-qualify table references in a CTE body to avoid circular
+# references. DuckDB 1.3.0+ and SQLite reject a CTE whose name matches
+# a table it references (e.g., `WITH t AS (SELECT ... FROM t) SELECT * FROM t`).
+# Prefixing the FROM/JOIN table with the schema name (e.g., `FROM main.t`)
+# disambiguates the CTE name from the physical table.
+# @param conn A DBI connection
+# @param cte_body A SQL string (the CTE body from dbplyr::remote_query())
+# @return The CTE body with table references schema-qualified
+# @noRd
+qualify_cte_table_refs <- function(conn, cte_body) {
+  schema <- as.character(DBI::dbQuoteIdentifier(conn, get_current_schema(conn)))
+
+  # Match FROM/JOIN followed by a table name (word or quoted identifier)
+  # that is not already schema-qualified (no dot immediately after).
+  # Subquery FROM clauses (FROM (SELECT ...)) are skipped because "(" is
+  # neither a word character nor a quoted identifier.
+  pattern <- "((?:FROM|JOIN)\\s+)([a-zA-Z_][a-zA-Z0-9_]*|\"[^\"]+\")(?![.])"
+  gsub(pattern, paste0("\\1", schema, ".\\2"), cte_body, perl = TRUE)
+}

pkg-r/tests/testthat/apps/basic/app.R

@@ -24,12 +24,27 @@ dbDisconnect(conn)
 # Setup database source and QueryChat instance
 db_conn <- dbConnect(RSQLite::SQLite(), temp_db)
 
+# ellmer 0.5.0 moved model details out of `Provider` and into a new `Model`
+# class: `Provider()` no longer takes `model` as its second positional
+# argument, and `Chat$new()` now requires a separate `model` argument. `Model`
+# doesn't exist before 0.5.0, so branch on its presence to support both.
+mock_chat_client <- if (
+  exists("Model", where = asNamespace("ellmer"), inherits = FALSE)
+) {
+  MockChat$new(
+    ellmer::Provider("test", "test"),
+    model = ellmer::Model(name = "test")
+  )
+} else {
+  MockChat$new(ellmer::Provider("test", "test", "test"))
+}
+
 # Create QueryChat instance
 qc <- QueryChat$new(
   data_source = db_conn,
   table_name = "iris",
   greeting = "Welcome to the test app!",
-  client = MockChat$new(ellmer::Provider("test", "test", "test"))
+  client = mock_chat_client
 )
 
 ui <- page_sidebar(

pkg-r/tests/testthat/helper-fixtures.R

@@ -173,7 +173,23 @@ mock_ellmer_chat_client <- function(
     private = private
   )
 
-  MockChat$new(ellmer::Provider("test", "test", "test"))
+  new_mock_chat(MockChat)
+}
+
+# ellmer 0.5.0 moved model details out of `Provider` and into a new `Model`
+# class: `Provider()` no longer takes `model` as its second positional
+# argument, and `Chat$new()` now requires a separate `model` argument. `Model`
+# doesn't exist before 0.5.0, so branch on its presence to support both.
+new_mock_chat <- function(MockChat, ...) {
+  if (exists("Model", where = asNamespace("ellmer"), inherits = FALSE)) {
+    MockChat$new(
+      ellmer::Provider("test", "test"),
+      model = ellmer::Model(name = "test"),
+      ...
+    )
+  } else {
+    MockChat$new(ellmer::Provider("test", "test", "test"), ...)
+  }
 }
 
 # shinychat::chat_restore() validates that `client` is an ellmer::Chat() R6

pyproject.toml

@@ -24,7 +24,10 @@ dependencies = [
     "shiny>=1.6.2",
     "shinychat @ git+https://github.com/posit-dev/shinychat.git@main",
     "htmltools",
-    "chatlas>=0.21.2",
+    # TODO: before merging, replace this pin with a released version
+    # (chatlas>=X.Y.Z) once https://github.com/posit-dev/chatlas/pull/427
+    # (Chat.close()/close_async()) is merged and released.
+    "chatlas @ git+https://github.com/posit-dev/chatlas.git@refs/pull/427/head",
     "narwhals>=2.2.0",
     "chevron",
     "sqlalchemy>=2.0.0", # Using 2.0+ for improved type hints and API