feat(pkg-py): Add cleanup param

pkg-py/CHANGELOG.md

@@ -7,6 +7,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [UNRELEASED]
 
+### Improvements
+
+* If a `QueryChat` instance is created inside a Shiny app context, the data source is automatically cleaned up when the app session ends. This can be configured manually via the `cleanup` parameter of the `QueryChat` constructor. (#164)
 
 
 ## [0.3.0] - 2025-12-10

pkg-py/src/querychat/_querychat.py

@@ -36,6 +36,7 @@ def __init__(
         categorical_threshold: int = 20,
         extra_instructions: Optional[str | Path] = None,
         prompt_template: Optional[str | Path] = None,
+        cleanup: Optional[bool] = None,
     ):
         self._data_source = normalize_data_source(data_source, table_name)
 
@@ -57,6 +58,11 @@ def __init__(
             prompt_template=prompt_template,
         )
 
+        # By default, only close the connection automatically if we're in a Shiny app session
+        if cleanup is None:
+            cleanup = get_current_session() is not None
+        self._cleanup = cleanup
+
         # Fork and empty chat now so the per-session forks are fast
         client = as_querychat_client(client)
         self._client = copy.deepcopy(client)
@@ -345,6 +351,10 @@ class QueryChat(QueryChatBase):
           `data_source.get_schema()`
         - `{{data_description}}`: The optional data description provided
         - `{{extra_instructions}}`: Any additional instructions provided
+    cleanup
+        Whether or not to automatically run `data_source.cleanup()` when the
+        Shiny session/app stops. By default, cleanup only occurs if the
+        `QueryChat` instance is created within a Shiny session.
 
     """
 
@@ -421,6 +431,7 @@ def title():
             greeting=self.greeting,
             client=self._client,
             enable_bookmarking=enable_bookmarking,
+            cleanup=self._cleanup,
         )
 
 
@@ -514,6 +525,10 @@ def data_table():
           `data_source.get_schema()`
         - `{{data_description}}`: The optional data description provided
         - `{{extra_instructions}}`: Any additional instructions provided
+    cleanup
+        Whether or not to automatically run `data_source.cleanup()` when the
+        Shiny session/app stops. By default, cleanup only occurs if the
+        `QueryChat` instance is created within a Shiny session.
 
     """
 
@@ -530,6 +545,7 @@ def __init__(
         extra_instructions: Optional[str | Path] = None,
         prompt_template: Optional[str | Path] = None,
         enable_bookmarking: Literal["auto", True, False] = "auto",
+        cleanup: Optional[bool] = None,
     ):
         # Sanity check: Express should always have a (stub/real) session
         session = get_current_session()
@@ -563,12 +579,20 @@ def __init__(
         else:
             enable = enable_bookmarking
 
+        if isinstance(session, ExpressStubSession):
+            cleanup = False
+        elif cleanup is None:
+            cleanup = True
+
+        self._cleanup = cleanup
+
         self._vals = mod_server(
             self.id,
             data_source=self._data_source,
             greeting=self.greeting,
             client=self._client,
             enable_bookmarking=enable,
+            cleanup=cleanup,
         )
 
     def df(self) -> pd.DataFrame:

pkg-py/src/querychat/_querychat_module.py

@@ -93,6 +93,7 @@ def mod_server(
     greeting: str | None,
     client: chatlas.Chat,
     enable_bookmarking: bool,
+    cleanup: bool = False,
 ):
     # Reactive values to store state
     sql = ReactiveStringOrNone(None)
@@ -187,6 +188,13 @@ def _on_restore(x: RestoreState) -> None:
             if "querychat_has_greeted" in vals:
                 has_greeted.set(vals["querychat_has_greeted"])
 
+    if cleanup:
+        # Clean up the data source when the session ends
+
+        @session.on_ended
+        def _cleanup() -> None:
+            data_source.cleanup()
+
     return ServerValues(df=filtered_df, sql=sql, title=title, client=chat)