diff --git a/pkg-py/CHANGELOG.md b/pkg-py/CHANGELOG.md index fbe256070..e9726b143 100644 --- a/pkg-py/CHANGELOG.md +++ b/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 diff --git a/pkg-py/src/querychat/_querychat_base.py b/pkg-py/src/querychat/_querychat_base.py index cf4c548a5..a8a3b9b71 100644 --- a/pkg-py/src/querychat/_querychat_base.py +++ b/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( diff --git a/pkg-py/tests/test_cleanup.py b/pkg-py/tests/test_cleanup.py new file mode 100644 index 000000000..b50730e9b --- /dev/null +++ b/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 diff --git a/pkg-r/R/TblSqlSource.R b/pkg-r/R/TblSqlSource.R index 7ab55e7ac..847460ed3 100644 --- a/pkg-r/R/TblSqlSource.R +++ b/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) +} diff --git a/pkg-r/tests/testthat/apps/basic/app.R b/pkg-r/tests/testthat/apps/basic/app.R index efa2a9c7f..6a31f7463 100644 --- a/pkg-r/tests/testthat/apps/basic/app.R +++ b/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( diff --git a/pkg-r/tests/testthat/helper-fixtures.R b/pkg-r/tests/testthat/helper-fixtures.R index 5180cf59d..7ce594ce5 100644 --- a/pkg-r/tests/testthat/helper-fixtures.R +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index cf228908c..6c59dfba5 100644 --- a/pyproject.toml +++ b/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