feat(pkg-py): route ggsql viz execution through ggsql.CachingReader

pkg-py/CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## Unreleased
 
+### New features
+
+* **Faster iterative visualization, and support for layer-level data sources.** Visualization queries now execute as a whole through `ggsql.CachingReader`: the underlying SQL runs against your data source once and the result is memoized for the chat session, so re-rendering a chart with a different encoding no longer re-queries the source (a win for remote databases especially). Visualization layers can also draw from other relations in the query — e.g. raw points from the final `SELECT` with a summary line from a CTE — via `DRAW ... MAPPING ... FROM <source>`. Caching can be tuned or disabled with the `GGSQL_CACHE_TTL`, `GGSQL_CACHE_MAX_BYTES`, and `GGSQL_CACHE_DISABLED` environment variables. (#321)
+
 ### Bug fixes
 
 * JSON pins now fall back to `pin_read()` when DuckDB's JSON extension cannot be installed or loaded, including on builds that disable extension loading. Native DuckDB JSON reading is retained when the extension is available.

pkg-py/docs/visualize.qmd

@@ -96,6 +96,22 @@ Each chart has controls in its footer:
 Note that visualization queries are independent of any active dashboard filter set by the `update` tool.
 They always run against the full dataset.
 
+### Caching
+
+To keep iterative chart refinement fast, querychat memoizes the data portion of visualization queries: when the LLM issues identical SQL within 5 minutes (for example, when only the visual encoding changes), the cached result is reused instead of re-reading your data source.
+Each chat session has its own cache, bounded at 512 MiB, which is freed when the session ends.
+
+This assumes the data doesn't change mid-session.
+If it can — for example, update queries are enabled (`QUERYCHAT_ENABLE_UPDATE_QUERIES=true`), or another process writes to the database — charts may lag those changes by up to the TTL below.
+
+Tune or disable the cache with environment variables (set before the app starts):
+
+| Variable | Default | Purpose |
+|---|---|---|
+| `GGSQL_CACHE_TTL` | `300` | Seconds a cached result is treated as fresh |
+| `GGSQL_CACHE_MAX_BYTES` | `512mb` | Per-session memory budget before least-recently-used eviction |
+| `GGSQL_CACHE_DISABLED` | unset | Set to any non-empty value except `0` to always re-read the source |
+
 Learn more about the ggsql grammar at [ggsql.org](https://ggsql.org/).
 
 ## See also

pkg-py/src/querychat/_query_executor.py

@@ -3,10 +3,11 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Any, cast
+from typing import TYPE_CHECKING, Any, TypeAlias, Union, cast
 
 import duckdb
 import narwhals.stable.v1 as nw
+from narwhals.typing import IntoFrame
 
 from ._datasource import (
     ColumnMeta,
@@ -22,15 +23,30 @@
 if TYPE_CHECKING:
     from collections.abc import Mapping
 
+    import ibis
+
     from ._datasource import DataFrameSource, DataSource, PolarsLazySource
     from ._pin_source import PinSource
 
 
+QueryResult: TypeAlias = Union[IntoFrame, "ibis.Table"]
+"""
+Result of ``QueryExecutor.execute_query()``.
+
+Any narwhals-compatible frame (``IntoFrame`` is narwhals' open-ended
+union of frame protocols, which grows as narwhals adds backends, so
+third-party executors are covered without changes here), plus lazy ibis
+tables. Consumers should normalize results with ``as_narwhals()``, which
+accepts every one of these (materializing lazy ibis tables via
+``.execute()``), rather than assume a fixed set of concrete types.
+"""
+
+
 class QueryExecutor(ABC):
     """Thin abstraction that tools use for query execution and validation."""
 
     @abstractmethod
-    def execute_query(self, query: str) -> Any: ...
+    def execute_query(self, query: str) -> QueryResult: ...
 
     @abstractmethod
     def test_query(
@@ -107,7 +123,7 @@ def __init__(self, sources: dict[str, DataFrameSource | PinSource]):
             self._conn.close()
             raise
 
-    def execute_query(self, query: str) -> Any:
+    def execute_query(self, query: str):
         check_query(query)
         result = self._conn.execute(query)
         return self._convert_result(result)
@@ -170,7 +186,7 @@ def __init__(self, sources: dict[str, PolarsLazySource]):
         for name, source in sources.items():
             self._table_columns[name] = list(source.get_data().collect_schema().keys())
 
-    def execute_query(self, query: str) -> Any:
+    def execute_query(self, query: str):
         check_query(query)
         return self._ctx.execute(query)
 
@@ -216,7 +232,7 @@ def __init__(self, data_sources: dict[str, DataSource]):
         self._data_sources = data_sources
         self._primary = next(iter(data_sources.values()))
 
-    def execute_query(self, query: str) -> Any:
+    def execute_query(self, query: str):
         return self._primary.execute_query(query)
 
     def test_query(

pkg-py/src/querychat/_querychat_base.py

@@ -8,7 +8,7 @@
 import re
 import warnings
 from pathlib import Path
-from typing import TYPE_CHECKING, Generic, Literal, Optional
+from typing import TYPE_CHECKING, Generic, Literal, Optional, cast
 
 import chatlas
 import narwhals.stable.v1 as nw
@@ -885,7 +885,13 @@ def _df_for_source(
     ) -> IntoFrameT:
         if sql:
             with contextlib.suppress(Exception):
-                return self._require_query_executor("df").execute_query(sql)
+                # The executor is built from this source, so its result
+                # backend matches the source's IntoFrameT (QueryResult is
+                # the wider union across all source types).
+                return cast(
+                    "IntoFrameT",
+                    self._require_query_executor("df").execute_query(sql),
+                )
         return data_source.get_data()
 
     def df(self, state: AppStateDict | None, *, table: str | None = None) -> IntoFrameT:

pkg-py/src/querychat/_querychat_core.py

@@ -33,10 +33,8 @@
 if TYPE_CHECKING:
     from collections.abc import AsyncIterator, Iterator
 
-    from narwhals.stable.v1.typing import IntoFrame
-
     from ._datasource import DataSource
-    from ._query_executor import QueryExecutor
+    from ._query_executor import QueryExecutor, QueryResult
 
 
 ClientFactory = Callable[
@@ -210,7 +208,7 @@ def get_active_data_source(self) -> DataSource:
             return self.data_sources[self.active_table]
         return next(iter(self.data_sources.values()))
 
-    def get_current_data(self) -> IntoFrame:
+    def get_current_data(self) -> QueryResult:
         """Get current data, falling back to default if query fails."""
         data_source = self.get_active_data_source()
         if self.sql:

pkg-py/src/querychat/_shiny_module.py

@@ -3,7 +3,7 @@
 import warnings
 from dataclasses import dataclass
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Generic, TypedDict, Union
+from typing import TYPE_CHECKING, Any, Generic, TypedDict, Union, cast
 
 import chatlas
 import shinychat
@@ -17,7 +17,7 @@
 from ._querychat_core import warn_multi_table_flat_accessor
 from ._table_accessor import TableAccessor
 from ._viz_altair_widget import AltairWidget
-from ._viz_ggsql import execute_ggsql
+from ._viz_ggsql import ExecutorReader
 from ._viz_utils import has_viz_tool, preload_viz_deps_server, preload_viz_deps_ui
 
 if TYPE_CHECKING:
@@ -266,7 +266,10 @@ def _make_table_state(
         def filtered_df() -> IntoFrameT:
             query = table_sql.get()
             if query:
-                return exec.execute_query(query)
+                # The executor is built from this source, so its result
+                # backend matches the source's IntoFrameT (QueryResult is
+                # the wider union across all source types).
+                return cast("IntoFrameT", exec.execute_query(query))
             return source.get_data()
 
         return TableState(sql=table_sql, title=table_title, df=filtered_df)
@@ -504,17 +507,23 @@ def restore_viz_widgets(
     saved_widgets: list[VizWidgetEntry],
 ) -> list[VizWidgetEntry]:
     """Re-execute ggsql queries, register widgets, and return restored entries."""
-    from ggsql import validate
+    from ggsql import CachingReader, validate
     from shinywidgets import register_widget
 
     restored: list[VizWidgetEntry] = []
+    if not saved_widgets:
+        return restored
+
+    # One reader for the whole restore pass: restored widgets often share a
+    # base read, so the memo avoids re-executing it for each entry.
+    reader = CachingReader(ExecutorReader(executor))
 
     for entry in saved_widgets:
         widget_id = entry["widget_id"]
         ggsql_str = entry["ggsql"]
         try:
-            validated = validate(ggsql_str)
-            spec = execute_ggsql(executor, validated)
+            validate(ggsql_str)  # early syntax check; errors skip the entry
+            spec = reader.execute(ggsql_str)
             altair_widget = AltairWidget.from_ggsql(spec, widget_id=widget_id)
             register_widget(widget_id, altair_widget.widget)
             restored.append(entry)

pkg-py/src/querychat/_viz_ggsql.py

@@ -1,102 +1,59 @@
-"""Helpers for ggsql integration."""
+"""A ggsql reader backed by a querychat executor."""
 
 from __future__ import annotations
 
-import re
+import warnings
 from typing import TYPE_CHECKING
 
-from ._utils import to_polars
+from ._utils import as_narwhals
 
 if TYPE_CHECKING:
-    import ggsql
+    import pyarrow as pa
 
+    from ._datasource import DataSource
     from ._query_executor import QueryExecutor
 
 
-def execute_ggsql(executor: QueryExecutor, validated: ggsql.Validated) -> ggsql.Spec:
+class ExecutorReader:
     """
-    Execute a pre-validated ggsql query against a QueryExecutor, returning a Spec.
-
-    Executes the SQL portion through the executor (preserving database pushdown),
-    then feeds the result into a ggsql DuckDBReader to produce a Spec.
-
-    Parameters
-    ----------
-    executor
-        The querychat QueryExecutor to execute the SQL portion against.
-    validated
-        A pre-validated ggsql query (from ``ggsql.validate()``).
-
-    Returns
-    -------
-    ggsql.Spec
-        The writer-independent plot specification.
-
+    ggsql reader protocol over a querychat executor (source reads only).
+
+    Designed exclusively for composition with ``ggsql.CachingReader``: the
+    caching reader stages every derived table into its own in-memory cache,
+    so this object only ever serves base reads of the user's data, executed
+    through the querychat executor (preserving database pushdown).
+
+    ``register``/``unregister`` are part of ggsql's reader protocol but are
+    never called in this composition — ``CachingReader`` routes all writes
+    to its private cache — so they only warn (loudly, in case a future
+    ggsql changes that routing, which would otherwise silently drop
+    writes).
     """
-    from ggsql import DuckDBReader
 
-    visual = validated.visual()
-    if has_layer_level_source(visual):
-        # Short term, querychat only supports visual layers that can be replayed
-        # from one final SQL result. Long term, the cleaner fix is likely to use
-        # ggsql's native remote-reader execution path (for example via ODBC-backed
-        # Readers) instead of reconstructing multi-relation scope here.
-        raise ValueError(
-            "Layer-specific sources are not currently supported in querychat visual "
-            "queries. Rewrite the query so that all layers come from the final SQL "
-            "result."
+    def __init__(self, executor: QueryExecutor | DataSource) -> None:
+        self._executor = executor
+
+    def execute_sql(self, sql: str) -> pa.Table:
+        """Run a base read through the executor, as a pyarrow Table."""
+        # Narwhals normalizes every execute_query() result type (pandas,
+        # polars, pyarrow, lazy ibis, ...) and to_arrow() converts directly,
+        # so this path does not require polars (only pyarrow, which ggsql
+        # itself guarantees).
+        return as_narwhals(self._executor.execute_query(sql)).to_arrow()
+
+    def register(self, name, table, replace) -> None:
+        # Writes belong to CachingReader's cache, never the primary.
+        warnings.warn(
+            f"ExecutorReader.register({name!r}, table={type(table).__name__}, "
+            f"replace={replace}) called: ggsql routed a write to the primary "
+            "reader, which is a no-op by design. The registered table will "
+            "not be visible to subsequent reads.",
+            stacklevel=2,
         )
 
-    pl_df = to_polars(executor.execute_query(validated.sql()))
-
-    reader = DuckDBReader("duckdb://memory")
-    table = extract_visualise_table(visual)
-
-    if table is not None:
-        # VISUALISE [mappings] FROM <table> — register data under the
-        # referenced table name and execute the visual part directly.
-        name = table[1:-1] if table.startswith('"') and table.endswith('"') else table
-        reader.register(name, pl_df)
-        return reader.execute(visual)
-    else:
-        # SELECT ... VISUALISE — no FROM in VISUALISE clause, so register
-        # under a synthetic name and prepend a SELECT.
-        reader.register("_data", pl_df)
-        return reader.execute(f"SELECT * FROM _data {visual}")
-
-
-def extract_visualise_table(visual: str) -> str | None:
-    """
-    Extract the table name from ``VISUALISE ... FROM <table>`` if present.
-
-    This reimplements a small part of ggsql's parsing because the current
-    Python bindings do not expose the top-level VISUALISE source directly.
-    """
-    draw_pos = re.search(r"\bDRAW\b", visual, re.IGNORECASE)
-    vis_clause = visual[: draw_pos.start()] if draw_pos else visual
-    m = re.search(r'\bFROM\s+("[^"]+?"|\S+)', vis_clause, re.IGNORECASE)
-    return m.group(1) if m else None
-
-
-def has_layer_level_source(visual: str) -> bool:
-    """
-    Return ``True`` when a DRAW clause defines its own ``FROM <source>``.
-
-    Querychat currently replays the VISUALISE portion against a single local
-    relation, so layer-specific sources cannot be preserved reliably.
-    """
-    clauses = re.split(
-        r"(?=\b(?:DRAW|SCALE|PROJECT|FACET|PLACE|LABEL|THEME)\b)",
-        visual,
-        flags=re.IGNORECASE,
-    )
-    for clause in clauses:
-        if not re.match(r"^\s*DRAW\b", clause, re.IGNORECASE):
-            continue
-        if re.search(
-            r'\bMAPPING\b[\s\S]*?\bFROM\s+("[^"]+?"|\S+)',
-            clause,
-            re.IGNORECASE,
-        ):
-            return True
-    return False
+    def unregister(self, name) -> None:
+        warnings.warn(
+            f"ExecutorReader.unregister({name!r}) called: ggsql routed a "
+            "write to the primary reader, which is a no-op by design.",
+            stacklevel=2,
+        )

pkg-py/src/querychat/_viz_tools.py

@@ -22,7 +22,7 @@
 from ._tool_names import TOOL_VISUALIZE
 from ._utils import querychat_tool_starts_open, read_prompt_template, truncate_error
 from ._viz_altair_widget import AltairWidget, fit_chart_to_container
-from ._viz_ggsql import execute_ggsql
+from ._viz_ggsql import ExecutorReader
 
 if TYPE_CHECKING:
     from collections.abc import Callable
@@ -188,13 +188,20 @@ def visualize_impl(
     update_fn: Callable[[VisualizeData], None],
 ) -> Callable[[str, str], ContentToolResult]:
     """Create the visualize implementation function."""
-    from ggsql import VegaLiteWriter, validate
+    from ggsql import CachingReader, VegaLiteWriter, validate
+
+    # One CachingReader per chat client (this closure is created once per
+    # client): the memo only pays off across viz calls within a session, and
+    # the reader's cache is freed when the client is dropped. Created lazily
+    # so clients that never visualize don't pay for an idle DuckDB instance.
+    reader = None
 
     def visualize(
         ggsql: str,
         title: str,
     ) -> ContentToolResult:
         """Execute a ggsql query and render the visualization."""
+        nonlocal reader
         markdown = f"```sql\n{ggsql}\n```"
 
         try:
@@ -210,7 +217,9 @@ def visualize(
                     "\n".join(error["message"] for error in validated.errors())
                 )
 
-            spec = execute_ggsql(executor, validated)
+            if reader is None:
+                reader = CachingReader(ExecutorReader(executor))
+            spec = reader.execute(ggsql)
 
             raw_chart = VegaLiteWriter().render_chart(spec)
             altair_widget = AltairWidget(copy.deepcopy(raw_chart))