diff --git a/pkg-py/CHANGELOG.md b/pkg-py/CHANGELOG.md index 057f4c1e6..403df5432 100644 --- a/pkg-py/CHANGELOG.md +++ b/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 `. 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. diff --git a/pkg-py/docs/visualize.qmd b/pkg-py/docs/visualize.qmd index e8d720b68..1a708fb70 100644 --- a/pkg-py/docs/visualize.qmd +++ b/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 diff --git a/pkg-py/src/querychat/_query_executor.py b/pkg-py/src/querychat/_query_executor.py index fca4857df..fb662aac1 100644 --- a/pkg-py/src/querychat/_query_executor.py +++ b/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( diff --git a/pkg-py/src/querychat/_querychat_base.py b/pkg-py/src/querychat/_querychat_base.py index f516c1825..5abd1ad52 100644 --- a/pkg-py/src/querychat/_querychat_base.py +++ b/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: diff --git a/pkg-py/src/querychat/_querychat_core.py b/pkg-py/src/querychat/_querychat_core.py index 0c33df3c4..1c6ba90bf 100644 --- a/pkg-py/src/querychat/_querychat_core.py +++ b/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: diff --git a/pkg-py/src/querychat/_shiny_module.py b/pkg-py/src/querychat/_shiny_module.py index 3d7656d0a..11d744271 100644 --- a/pkg-py/src/querychat/_shiny_module.py +++ b/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) diff --git a/pkg-py/src/querychat/_viz_ggsql.py b/pkg-py/src/querychat/_viz_ggsql.py index 9e166c151..ff7e9bc82 100644 --- a/pkg-py/src/querychat/_viz_ggsql.py +++ b/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 — 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
`` 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 ``. - - 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, + ) diff --git a/pkg-py/src/querychat/_viz_tools.py b/pkg-py/src/querychat/_viz_tools.py index fa6306d8c..e38ba688f 100644 --- a/pkg-py/src/querychat/_viz_tools.py +++ b/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)) diff --git a/pkg-py/src/querychat/prompts/ggsql-syntax.md b/pkg-py/src/querychat/prompts/ggsql-syntax.md index f9ad904ab..79abd7bb8 100644 --- a/pkg-py/src/querychat/prompts/ggsql-syntax.md +++ b/pkg-py/src/querychat/prompts/ggsql-syntax.md @@ -513,34 +513,15 @@ PROJECT TO polar SETTING inner => 0.5 VISUALISE FROM titanic DRAW bar MAPPING class AS x ``` -3. **In querychat, all layers must come from the final SQL result**: Do not use layer-specific `FROM source` inside `DRAW ... MAPPING ...` clauses. If you need raw data and a summary in one chart, put both into one final relation and distinguish layers with a column such as `layer_type`: +3. **Layers can draw from other relations in the query**: Attach `FROM ` to a layer's `MAPPING` to combine different granularities in one chart — e.g. raw points from the final SELECT with a summary line from a CTE. The `FROM` belongs to `MAPPING` (a bare `DRAW line FROM t` is a parse error). Explicit layer mappings resolve against the layer's source; aesthetics the layer doesn't map are inherited from VISUALISE: ```sql - WITH raw AS ( - SELECT - date, - amount, - region, - 'raw' AS layer_type - FROM sales - ), - summary AS ( - SELECT - date, - AVG(amount) AS amount, - region, - 'summary' AS layer_type - FROM sales - GROUP BY date, region - ), - combined AS ( - SELECT * FROM raw - UNION ALL - SELECT * FROM summary + WITH summary AS ( + SELECT date, AVG(amount) AS amount FROM sales GROUP BY date ) - SELECT * FROM combined + SELECT date, amount FROM sales VISUALISE date AS x, amount AS y - DRAW point MAPPING region AS color FILTER layer_type = 'raw' - DRAW line MAPPING region AS color FILTER layer_type = 'summary' + DRAW point SETTING opacity => 0.3 + DRAW line MAPPING date AS x, amount AS y FROM summary ``` 4. **String values use single quotes**: In SETTING, LABEL, and RENAMING clauses, always use single quotes for string values. Double quotes cause parse errors. 5. **Statistical layers**: When using `histogram`, `bar` (without y), `density`, `smooth`, `violin`, or `boxplot`, the layer computes statistics. Use REMAPPING to access `density`, `intensity`, `proportion`, etc. diff --git a/pkg-py/tests/test_ggsql.py b/pkg-py/tests/test_ggsql.py index e7072808b..1f59718ca 100644 --- a/pkg-py/tests/test_ggsql.py +++ b/pkg-py/tests/test_ggsql.py @@ -2,52 +2,28 @@ import ggsql import narwhals.stable.v1 as nw +import pandas as pd import polars as pl +import pyarrow as pa import pytest -from querychat._datasource import DataFrameSource -from querychat._viz_altair_widget import AltairWidget -from querychat._viz_ggsql import ( - execute_ggsql, - extract_visualise_table, - has_layer_level_source, +from querychat._datasource import ( + DataFrameSource, + IbisSource, + PolarsLazySource, + SQLAlchemySource, ) +from querychat._query_executor import ( + DataSourceExecutor, + DuckDBExecutor, + PolarsSQLExecutor, +) +from querychat._viz_altair_widget import AltairWidget +from querychat._viz_ggsql import ExecutorReader -class TestExtractVisualiseTable: - """Tests for extract_visualise_table() parsing.""" - - def test_bare_identifier(self): - assert ( - extract_visualise_table("VISUALISE x, y FROM mytable DRAW point") - == "mytable" - ) - - def test_quoted_identifier(self): - assert ( - extract_visualise_table('VISUALISE x FROM "my table" DRAW point') - == '"my table"' - ) - - def test_no_from_returns_none(self): - assert extract_visualise_table("VISUALISE x, y DRAW point") is None - - def test_ignores_draw_level_from(self): - visual = "VISUALISE x, y DRAW bar MAPPING z AS fill FROM summary" - assert extract_visualise_table(visual) is None - - -class TestHasLayerLevelSource: - def test_detects_draw_level_from(self): - visual = "VISUALISE x, y DRAW bar MAPPING z AS fill FROM summary" - assert has_layer_level_source(visual) - - def test_ignores_visualise_from(self): - visual = "VISUALISE x, y FROM sales DRAW point MAPPING z AS color" - assert not has_layer_level_source(visual) - - def test_ignores_scale_from(self): - visual = "VISUALISE x, y DRAW point MAPPING z AS color SCALE x FROM [0, 10]" - assert not has_layer_level_source(visual) +def make_reader(source) -> ggsql.CachingReader: + """Create a CachingReader over an ExecutorReader, as the viz tool does.""" + return ggsql.CachingReader(ExecutorReader(source)) class TestGgsqlValidate: @@ -146,7 +122,7 @@ def test_full_pipeline(self): nw_df = nw.from_native(pl.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6]})) ds = DataFrameSource(nw_df, "test_data") query = "SELECT * FROM test_data VISUALISE x, y DRAW point" - spec = execute_ggsql(ds, ggsql.validate(query)) + spec = make_reader(ds).execute(query) altair_widget = AltairWidget.from_ggsql(spec) result = altair_widget.widget.chart.to_dict() assert "$schema" in result @@ -158,7 +134,7 @@ def test_with_filtered_query(self): ) ds = DataFrameSource(nw_df, "test_data") query = "SELECT * FROM test_data WHERE x > 2 VISUALISE x, y DRAW point" - spec = execute_ggsql(ds, ggsql.validate(query)) + spec = make_reader(ds).execute(query) assert spec.metadata()["rows"] == 3 @pytest.mark.ggsql @@ -166,7 +142,7 @@ def test_spec_has_visual(self): nw_df = nw.from_native(pl.DataFrame({"x": [1, 2], "y": [3, 4]})) ds = DataFrameSource(nw_df, "test_data") query = "SELECT * FROM test_data VISUALISE x, y DRAW point" - spec = execute_ggsql(ds, ggsql.validate(query)) + spec = make_reader(ds).execute(query) assert "VISUALISE" in spec.visual() @pytest.mark.ggsql @@ -174,7 +150,7 @@ def test_visualise_from_path(self): nw_df = nw.from_native(pl.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6]})) ds = DataFrameSource(nw_df, "test_data") query = "VISUALISE x, y FROM test_data DRAW point" - spec = execute_ggsql(ds, ggsql.validate(query)) + spec = make_reader(ds).execute(query) assert spec.metadata()["rows"] == 3 assert "VISUALISE" in spec.visual() @@ -185,7 +161,7 @@ def test_with_pandas_dataframe(self): nw_df = nw.from_native(pd.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6]})) ds = DataFrameSource(nw_df, "test_data") query = "SELECT * FROM test_data VISUALISE x, y DRAW point" - spec = execute_ggsql(ds, ggsql.validate(query)) + spec = make_reader(ds).execute(query) altair_widget = AltairWidget.from_ggsql(spec) result = altair_widget.widget.chart.to_dict() assert "$schema" in result @@ -206,49 +182,11 @@ def test_supports_uppercase_column_references_without_renaming(self): "VISUALISE ROOM_TYPE AS x, listings AS y " "DRAW bar" ) - spec = execute_ggsql(ds, ggsql.validate(query)) - assert "VISUALISE" in spec.visual() - - def test_passes_original_column_names_through_to_ggsql(self, monkeypatch): - nw_df = nw.from_native( - pl.DataFrame( - { - "ROOM_TYPE": ["Entire home", "Private room"], - "listings": [10, 20], - } - ) - ) - ds = DataFrameSource(nw_df, "upper_table") - validated = type( - "Validated", - (), - { - "sql": lambda self: "SELECT ROOM_TYPE, listings FROM upper_table", - "visual": lambda self: ( - "VISUALISE ROOM_TYPE AS x, listings AS y DRAW bar" - ), - }, - )() - captured_columns = None - - class FakeReader: - def register(self, name, df): - nonlocal captured_columns - captured_columns = list(df.columns) - - def execute(self, query): - return type("Spec", (), {"visual": lambda self: query})() - - import ggsql - - monkeypatch.setattr(ggsql, "DuckDBReader", lambda *_args: FakeReader()) - - spec = execute_ggsql(ds, validated) + spec = make_reader(ds).execute(query) assert "VISUALISE" in spec.visual() - assert captured_columns == ["ROOM_TYPE", "listings"] @pytest.mark.ggsql - def test_rejects_layer_level_from_sources_with_clear_error(self): + def test_supports_layer_level_from_sources(self): nw_df = nw.from_native( pl.DataFrame( { @@ -272,11 +210,12 @@ def test_rejects_layer_level_from_sources_with_clear_error(self): DRAW bar MAPPING region AS x, total AS y FROM summary """ - with pytest.raises( - ValueError, - match="Layer-specific sources are not currently supported", - ): - execute_ggsql(ds, ggsql.validate(query)) + # Previously raised "Layer-specific sources are not currently + # supported"; CachingReader stages each layer source through its + # cache, so the query now runs end to end. + spec = make_reader(ds).execute(query) + assert spec is not None + assert "VISUALISE" in spec.visual() @pytest.mark.ggsql def test_supports_single_relation_raw_plus_summary_overlay(self): @@ -320,6 +259,249 @@ def test_supports_single_relation_raw_plus_summary_overlay(self): DRAW line MAPPING category AS color FILTER layer_type = 'summary' """ - spec = execute_ggsql(ds, ggsql.validate(query)) + spec = make_reader(ds).execute(query) assert spec.metadata()["rows"] == 4 assert "VISUALISE" in spec.visual() + + +class TestExecutorReaderResultTypes: + """ + ExecutorReader.execute_sql normalizes every execute_query() result type + (the QueryResult union) to a pyarrow Table, without requiring polars. + """ + + @staticmethod + def check(executor): + tbl = ExecutorReader(executor).execute_sql("SELECT x, y FROM t") + assert isinstance(tbl, pa.Table) + assert tbl.schema.names == ["x", "y"] + assert tbl.num_rows == 2 + + @pytest.fixture( + params=[ + pd.DataFrame({"x": [1, 2], "y": [3, 4]}), + pl.DataFrame({"x": [1, 2], "y": [3, 4]}), + pa.table({"x": [1, 2], "y": [3, 4]}), + ], + ids=["pandas", "polars", "pyarrow"], + ) + def frame(self, request): + return request.param + + def test_dataframe_source_backends(self, frame): + self.check(DataFrameSource(nw.from_native(frame), "t")) + + def test_duckdb_executor_backends(self, frame): + source = DataFrameSource(nw.from_native(frame), "t") + self.check(DuckDBExecutor({"t": source})) + + def test_polars_lazy_source(self): + lf = nw.from_native(pl.LazyFrame({"x": [1, 2], "y": [3, 4]})) + self.check(PolarsLazySource(lf, "t")) + + def test_polars_sql_executor(self): + lf = nw.from_native(pl.LazyFrame({"x": [1, 2], "y": [3, 4]})) + self.check(PolarsSQLExecutor({"t": PolarsLazySource(lf, "t")})) + + def test_sqlalchemy_source(self): + sqlalchemy = pytest.importorskip("sqlalchemy") + eng = sqlalchemy.create_engine("sqlite://") + pd.DataFrame({"x": [1, 2], "y": [3, 4]}).to_sql("t", eng, index=False) + self.check(SQLAlchemySource(eng, "t")) + + def test_ibis_source(self): + ibis = pytest.importorskip("ibis") + con = ibis.duckdb.connect() + con.create_table("t", pd.DataFrame({"x": [1, 2], "y": [3, 4]})) + self.check(IbisSource(con.table("t"), "t")) + + def test_datasource_executor_delegation(self): + sqlalchemy = pytest.importorskip("sqlalchemy") + eng = sqlalchemy.create_engine("sqlite://") + pd.DataFrame({"x": [1, 2], "y": [3, 4]}).to_sql("t", eng, index=False) + self.check(DataSourceExecutor({"t": SQLAlchemySource(eng, "t")})) + + def test_register_unregister_warn(self): + # register/unregister are no-ops by design (CachingReader owns all + # writes); they warn loudly in case a future ggsql routes writes + # to the primary reader, which would otherwise silently drop them. + ds = DataFrameSource(nw.from_native(pl.DataFrame({"x": [1]})), "t") + reader = ExecutorReader(ds) + with pytest.warns(UserWarning, match="no-op by design"): + reader.register("t", None, replace=False) + with pytest.warns(UserWarning, match="no-op by design"): + reader.unregister("t") + + def test_empty_result(self): + ds = DataFrameSource(nw.from_native(pl.DataFrame({"x": [1, 2]})), "t") + tbl = ExecutorReader(ds).execute_sql("SELECT * FROM t WHERE x > 100") + assert isinstance(tbl, pa.Table) + assert tbl.num_rows == 0 + + def test_null_column(self): + ds = DataFrameSource(nw.from_native(pl.DataFrame({"x": [1, 2]})), "t") + tbl = ExecutorReader(ds).execute_sql("SELECT NULL AS n FROM t") + assert tbl.schema.names == ["n"] + assert tbl.num_rows == 2 + + def test_duplicate_column_names_raise(self): + # Known limitation (pre-existing): narwhals requires unique column + # names, so a JOIN with overlapping names through a DBAPI-backed + # source fails conversion regardless of the arrow/polars cast. + # (DuckDB-family sources auto-mangle duplicates to x, x_1, so this + # only bites SQLAlchemy/Ibis-style sources.) + import sqlalchemy + + eng = sqlalchemy.create_engine("sqlite://") + pd.DataFrame({"id": [1, 2], "v": [10, 20]}).to_sql("a", eng, index=False) + pd.DataFrame({"id": [1, 2], "w": [30, 40]}).to_sql("b", eng, index=False) + src = SQLAlchemySource(eng, "a") + with pytest.raises(Exception, match=r"[Uu]nique column names"): + ExecutorReader(src).execute_sql("SELECT * FROM a JOIN b ON a.id = b.id") + + +class TestCachingBehavior: + """CachingReader shares memoized source reads across viz executions.""" + + @pytest.mark.ggsql + def test_identical_queries_share_source_reads(self, monkeypatch): + nw_df = nw.from_native(pl.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6]})) + ds = DataFrameSource(nw_df, "test_data") + + calls = 0 + original = ds.execute_query + + def counting(query): + nonlocal calls + calls += 1 + return original(query) + + monkeypatch.setattr(ds, "execute_query", counting) + + query = "SELECT * FROM test_data VISUALISE x, y DRAW point" + reader = make_reader(ds) + reader.execute(query) + reader.execute(query) + assert calls == 1 + + @pytest.mark.ggsql + def test_independent_readers_do_not_share_memo(self, monkeypatch): + """Readers own their caches; nothing is shared implicitly.""" + nw_df = nw.from_native(pl.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6]})) + ds = DataFrameSource(nw_df, "test_data") + + calls = 0 + original = ds.execute_query + + def counting(query): + nonlocal calls + calls += 1 + return original(query) + + monkeypatch.setattr(ds, "execute_query", counting) + + query = "SELECT * FROM test_data VISUALISE x, y DRAW point" + make_reader(ds).execute(query) + make_reader(ds).execute(query) + assert calls == 2 + + +class TestRestoreVizWidgets: + """ + Bookmark restore executes through a CachingReader (PR #321 follow-up). + + Regression guard: an earlier draft passed the executor itself where a + reader was expected, so every widget silently failed to restore (the + AttributeError was swallowed by restore_viz_widgets' except Exception). + """ + + @pytest.fixture + def register_spy(self, monkeypatch): + """Stub out the two places restore_viz_widgets needs a Shiny session.""" + import ipywidgets + import shinywidgets + + registered = [] + monkeypatch.setattr( + shinywidgets, + "register_widget", + lambda widget_id, widget: registered.append(widget_id), + ) + # shinywidgets hooks ipywidgets construction (Widget.on_widget_constructed) + # and requires an active session; unhook so widgets build session-free. + monkeypatch.setattr(ipywidgets.Widget, "_widget_construction_callback", None) + return registered + + @pytest.mark.ggsql + def test_restores_and_registers_widgets(self, register_spy): + from querychat._shiny_module import restore_viz_widgets + + nw_df = nw.from_native(pl.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6]})) + ds = DataFrameSource(nw_df, "test_data") + entries = [ + { + "widget_id": "w1", + "ggsql": "SELECT * FROM test_data VISUALISE x, y DRAW point", + }, + { + "widget_id": "w2", + "ggsql": "SELECT * FROM test_data WHERE x > 1 VISUALISE x, y DRAW point", + }, + ] + + restored = restore_viz_widgets(ds, entries) + + assert [e["widget_id"] for e in restored] == ["w1", "w2"] + assert register_spy == ["w1", "w2"] + + @pytest.mark.ggsql + def test_shared_base_reads_are_memoized(self, monkeypatch, register_spy): + """One reader per restore pass: identical queries share one source read.""" + from querychat._shiny_module import restore_viz_widgets + + nw_df = nw.from_native(pl.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6]})) + ds = DataFrameSource(nw_df, "test_data") + + calls = 0 + original = ds.execute_query + + def counting(query): + nonlocal calls + calls += 1 + return original(query) + + monkeypatch.setattr(ds, "execute_query", counting) + + query = "SELECT * FROM test_data VISUALISE x, y DRAW point" + entries = [ + {"widget_id": "w1", "ggsql": query}, + {"widget_id": "w2", "ggsql": query}, + ] + + restored = restore_viz_widgets(ds, entries) + + assert len(restored) == 2 + assert calls == 1 + + @pytest.mark.ggsql + def test_invalid_entries_are_skipped_with_warning(self, register_spy): + from querychat._shiny_module import restore_viz_widgets + + nw_df = nw.from_native(pl.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6]})) + ds = DataFrameSource(nw_df, "test_data") + entries = [ + { + "widget_id": "bad", + "ggsql": "SELECT * FROM nope VISUALISE x, y DRAW point", + }, + { + "widget_id": "good", + "ggsql": "SELECT * FROM test_data VISUALISE x, y DRAW point", + }, + ] + + with pytest.warns(UserWarning, match="Failed to restore"): + restored = restore_viz_widgets(ds, entries) + + assert [e["widget_id"] for e in restored] == ["good"] + assert register_spy == ["good"] diff --git a/pkg-py/tests/test_viz_footer.py b/pkg-py/tests/test_viz_footer.py index 90c228427..1900dfffe 100644 --- a/pkg-py/tests/test_viz_footer.py +++ b/pkg-py/tests/test_viz_footer.py @@ -49,7 +49,7 @@ def _patch_deps(monkeypatch): mock_altair_widget.widget_id = "querychat_viz_test1234" mock_altair_widget.is_compound = False - monkeypatch.setattr("querychat._viz_ggsql.execute_ggsql", lambda _ds, _q: mock_spec) + monkeypatch.setattr("ggsql.CachingReader.execute", lambda _self, _q: mock_spec) monkeypatch.setattr( "querychat._viz_altair_widget.AltairWidget.from_ggsql", staticmethod(lambda _spec: mock_altair_widget),