diff --git a/pkg-py/CHANGELOG.md b/pkg-py/CHANGELOG.md index 085ff509..64bb29b4 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 + +* Multiple pins (and pins mixed with data frames) are now supported in one chat: pin and data-frame tables are materialized into a shared DuckDB connection, so the LLM can join and filter across them. Each pin still keeps its own private connection for standalone use. (#312) + ### Changes * `.server(data_source=)` no longer modifies the `QueryChat` instance. The table is registered for that session only: the instance's tables, greeting tables, and system prompt are unchanged, a same-named instance table is shadowed for that session, and the session's data source is cleaned up when the session ends. This removes the concurrent-session edge cases that `0.8.0` patched around (#300, #302, #303, #304, #308). @@ -21,8 +25,6 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 * `cleanup()` no longer closes a spec-resolved `.server(client=...)` override while its session is still running; it is closed when the session ends. -* Registering a second pins table with `add_table()` now raises a clear error at registration time instead of failing at query time: each pin queries through its own DuckDB connection, so only the first pin's table would be queryable. To combine a pin with other tables, register them in a shared DuckDB connection and pass that instead. - ## [0.8.0] - 2026-09-12 ### New features diff --git a/pkg-py/src/querychat/_datasource.py b/pkg-py/src/querychat/_datasource.py index 546bd425..5dee42c6 100644 --- a/pkg-py/src/querychat/_datasource.py +++ b/pkg-py/src/querychat/_datasource.py @@ -117,6 +117,11 @@ def duckdb_column_meta(name: str, duckdb_type: Any) -> ColumnMeta: return ColumnMeta(name=name, sql_type=sql_type, kind=kind) +def quote_identifier(name: str) -> str: + """Return ``name`` as a double-quoted SQL identifier.""" + return '"' + name.replace('"', '""') + '"' + + def duckdb_column_stats( conn: duckdb.DuckDBPyConnection, table_name: str, @@ -137,7 +142,9 @@ def duckdb_column_stats( return try: - stats_query = f'SELECT {", ".join(select_parts)} FROM "{table_name}"' + stats_query = ( + f"SELECT {', '.join(select_parts)} FROM {quote_identifier(table_name)}" + ) result = conn.execute(stats_query).fetchone() if not result: return @@ -162,7 +169,7 @@ def duckdb_column_stats( try: for col in categorical_cols: cat_result = conn.execute( - f'SELECT DISTINCT "{col.name}" FROM "{table_name}" ' + f'SELECT DISTINCT "{col.name}" FROM {quote_identifier(table_name)} ' f'WHERE "{col.name}" IS NOT NULL ORDER BY "{col.name}"' ).fetchall() col.categories = [str(row[0]) for row in cat_result] @@ -501,6 +508,18 @@ def get_data(self) -> IntoDataFrameT: """ return self._df.to_native() + def register_into( + self, conn: duckdb.DuckDBPyConnection, table_name: str | None = None + ) -> None: + """ + Register this DataFrame in a shared DuckDB connection. + + Internal hook for joining a shared DuckDB executor. The caller owns + ``conn`` and locks it down once all tables are registered. + """ + # NOTE: if native representation is polars, pyarrow is required for registration + conn.register(table_name or self.table_name, self.get_data()) + def cleanup(self) -> None: """ Close the DuckDB connection. diff --git a/pkg-py/src/querychat/_pin_source.py b/pkg-py/src/querychat/_pin_source.py index b145d0e9..8a6a29c8 100644 --- a/pkg-py/src/querychat/_pin_source.py +++ b/pkg-py/src/querychat/_pin_source.py @@ -2,6 +2,7 @@ import re from typing import TYPE_CHECKING, Any, TypeGuard +from uuid import uuid4 import duckdb import narwhals.stable.v1 as nw @@ -14,6 +15,7 @@ duckdb_column_stats, duckdb_lock_down, format_schema, + quote_identifier, ) from ._utils import check_query @@ -65,6 +67,21 @@ def _convert_result(result: duckdb.DuckDBPyConnection) -> nw.DataFrame: return nw.from_native(result.df()) +def stage_frame_as_table( + conn: duckdb.DuckDBPyConnection, frame: Any, table_name: str +) -> None: + """Materialize a DataFrame as a real table via a unique staging view.""" + vname = f"__pin_staging_{table_name}_{uuid4().hex[:8]}" + conn.register(vname, frame) + try: + conn.execute( + f"CREATE TABLE {quote_identifier(table_name)} AS " + f"SELECT * FROM {quote_identifier(vname)}" + ) + finally: + conn.unregister(vname) + + class PinSource(DataSource[nw.DataFrame]): """ DataSource backed by a pin from a pins board. @@ -80,6 +97,13 @@ class PinSource(DataSource[nw.DataFrame]): :class:`~querychat.QueryChat` uses them as the default ``data_description``, which you can override. + Multiple pins + ~~~~~~~~~~~~~ + + Multiple pins (and pins mixed with data frames) can be combined in one + chat: every table is materialized into a shared DuckDB connection, so + the LLM can join and filter across them. + Lazy queries with pins ~~~~~~~~~~~~~~~~~~~~~~ @@ -118,61 +142,24 @@ def __init__( effective_table_name = _sanitize_table_name(table_name or name) self.table_name = effective_table_name + # Retained so the pin can be re-materialized into a shared DuckDB + # connection when it joins a multi-table executor (register_into()). + self._board = board + self._pin_name = name + self._pin_meta_obj = board.pin_meta(name, version=version) - pin_type = self._pin_meta_obj.type + # Snapshot the resolved version so register_into() reads the same pin + # content even if the pin is updated after construction. meta.version + # is a Version/VersionRaw (unwrap .version) or a plain string on some + # boards (e.g. Connect GUIDs). + resolved = getattr(self._pin_meta_obj.version, "version", None) + if not isinstance(resolved, str): + resolved = self._pin_meta_obj.version + self._version: str | None = resolved if isinstance(resolved, str) else version conn = duckdb.connect() try: - if pin_type in DUCKDB_FILE_TYPES: - paths = board.pin_download(name, version=version) - if len(paths) != 1: - raise ValueError( - f"Pin '{name}' contains {len(paths)} files, but PinSource " - "requires a single-file pin (as created by pin_write())." - ) - reader_fn = DUCKDB_READER_FN[pin_type] - if pin_type == "json": - conn.execute("INSTALL json") - conn.execute("LOAD json") - conn.execute( - f'CREATE TABLE "{effective_table_name}" AS ' - f"SELECT * FROM {reader_fn}(?)", - [paths[0]], - ) - elif pin_type == "arrow" and _has_polars(): - # Arrow/IPC files can't be read natively by DuckDB, but - # polars can read them directly — avoiding pin_read() overhead. - import polars as pl - - paths = board.pin_download(name, version=version) - if len(paths) != 1: - raise ValueError( - f"Pin '{name}' contains {len(paths)} files, but PinSource " - "requires a single-file pin (as created by pin_write())." - ) - arrow_df = pl.read_ipc(paths[0]) - vname = f"__pin_staging_{effective_table_name}" - conn.register(vname, arrow_df) - conn.execute( - f'CREATE TABLE "{effective_table_name}" AS SELECT * FROM "{vname}"' - ) - conn.unregister(vname) - else: - import pandas as pd - - data = board.pin_read(name, version=version) - if not isinstance(data, pd.DataFrame): - raise TypeError( - f"Pin '{name}' contains {type(data).__name__}, not a DataFrame. " - "PinSource requires the pin to contain a pandas DataFrame." - ) - vname = f"__pin_staging_{effective_table_name}" - conn.register(vname, data) - conn.execute( - f'CREATE TABLE "{effective_table_name}" AS SELECT * FROM "{vname}"' - ) - conn.unregister(vname) - + self._materialize_into(conn, effective_table_name) duckdb_lock_down(conn) except Exception: conn.close() @@ -184,6 +171,80 @@ def __init__( result = self._conn.execute(f'SELECT * FROM "{effective_table_name}" LIMIT 0') self._colnames = [desc[0] for desc in result.description] + def _materialize_into(self, conn: duckdb.DuckDBPyConnection, table_name: str): + """ + Materialize the pin as a real table in ``conn``. + + Does not lock the connection down; the caller owns ``conn`` and + decides when (or whether) to call :func:`duckdb_lock_down`. + """ + board, name, version = self._board, self._pin_name, self._version + pin_type = self._pin_meta_obj.type + + if pin_type in DUCKDB_FILE_TYPES: + paths = board.pin_download(name, version=version) + if len(paths) != 1: + raise ValueError( + f"Pin '{name}' contains {len(paths)} files, but PinSource " + "requires a single-file pin (as created by pin_write())." + ) + reader_fn = DUCKDB_READER_FN[pin_type] + if pin_type == "json": + conn.execute("INSTALL json") + conn.execute("LOAD json") + conn.execute( + f"CREATE TABLE {quote_identifier(table_name)} AS " + f"SELECT * FROM {reader_fn}(?)", + [paths[0]], + ) + elif pin_type == "arrow" and _has_polars(): + # Arrow/IPC files can't be read natively by DuckDB, but + # polars can read them directly — avoiding pin_read() overhead. + import polars as pl + + paths = board.pin_download(name, version=version) + if len(paths) != 1: + raise ValueError( + f"Pin '{name}' contains {len(paths)} files, but PinSource " + "requires a single-file pin (as created by pin_write())." + ) + arrow_df = pl.read_ipc(paths[0]) + stage_frame_as_table(conn, arrow_df, table_name) + else: + import pandas as pd + + data = board.pin_read(name, version=version) + if not isinstance(data, pd.DataFrame): + raise TypeError( + f"Pin '{name}' contains {type(data).__name__}, not a DataFrame. " + "PinSource requires the pin to contain a pandas DataFrame." + ) + stage_frame_as_table(conn, data, table_name) + + def register_into( + self, conn: duckdb.DuckDBPyConnection, table_name: str | None = None + ) -> None: + """ + Materialize this pin into a shared DuckDB connection. + + Internal hook for joining a shared DuckDB executor (multiple pins, or + pins mixed with data frames). The caller owns ``conn`` and locks it + down once all tables are materialized. + """ + target = table_name or self.table_name + from pins.errors import PinsError + + try: + self._materialize_into(conn, target) + except PinsError as e: + # The snapshotted pin version may have been pruned (e.g. a + # non-versioned board rewritten after construction); fall back to + # this source's own copy so the shared table matches the private + # connection. Other materialization failures still raise. + if "missing version" not in str(e): + raise + stage_frame_as_table(conn, self.get_data().to_native(), target) + def get_db_type(self) -> str: return "DuckDB" diff --git a/pkg-py/src/querychat/_query_executor.py b/pkg-py/src/querychat/_query_executor.py index d889647e..fca4857d 100644 --- a/pkg-py/src/querychat/_query_executor.py +++ b/pkg-py/src/querychat/_query_executor.py @@ -15,6 +15,7 @@ duckdb_column_stats, duckdb_lock_down, format_schema, + quote_identifier, ) from ._utils import check_query @@ -22,6 +23,7 @@ from collections.abc import Mapping from ._datasource import DataFrameSource, DataSource, PolarsLazySource + from ._pin_source import PinSource class QueryExecutor(ABC): @@ -77,22 +79,33 @@ def _validate_missing_columns( class DuckDBExecutor(QueryExecutor): - """Shared DuckDB connection for multi-table DataFrameSource queries.""" - - def __init__(self, sources: dict[str, DataFrameSource]): - self._df_lib = get_shared_dataframe_backend(sources) - self._conn = duckdb.connect(database=":memory:") - - for name, source in sources.items(): - self._conn.register(name, source.get_data()) + """ + Shared DuckDB connection for multi-table DataFrameSource/PinSource queries. - # Cache column names per table before lockdown - self._table_columns: dict[str, list[str]] = {} - for name in sources: - result = self._conn.execute(f'SELECT * FROM "{name}" LIMIT 0') - self._table_columns[name] = [desc[0] for desc in result.description] + Every source materializes its table into one connection (data frames via + ``register()``, pins via their file-based materialization), then the + connection is locked down once. + """ - duckdb_lock_down(self._conn) + def __init__(self, sources: dict[str, DataFrameSource | PinSource]): + self._df_lib = get_shared_duckdb_result_backend(sources) + self._conn = duckdb.connect(database=":memory:") + try: + for name, source in sources.items(): + source.register_into(self._conn, name) + + # Cache column names per table before lockdown + self._table_columns: dict[str, list[str]] = {} + for name in sources: + result = self._conn.execute( + f"SELECT * FROM {quote_identifier(name)} LIMIT 0" + ) + self._table_columns[name] = [desc[0] for desc in result.description] + + duckdb_lock_down(self._conn) + except Exception: + self._conn.close() + raise def execute_query(self, query: str) -> Any: check_query(query) @@ -132,7 +145,9 @@ def cleanup(self) -> None: self._conn.close() def get_column_metas(self, table_name: str) -> list[ColumnMeta]: - result = self._conn.execute(f'SELECT * FROM "{table_name}" LIMIT 0') + result = self._conn.execute( + f"SELECT * FROM {quote_identifier(table_name)} LIMIT 0" + ) return [duckdb_column_meta(desc[0], desc[1]) for desc in result.description] def populate_column_stats( @@ -228,22 +243,38 @@ def populate_column_stats( ) -def get_shared_dataframe_backend(sources: dict[str, DataFrameSource]) -> str: - """Return the shared backend name, rejecting mixed DataFrameSource backends.""" - source_items = iter(sources.items()) - _, first_source = next(source_items) - shared_lib = get_dataframe_backend_name(first_source) +def get_shared_duckdb_result_backend( + sources: dict[str, DataFrameSource | PinSource], +) -> str: + """ + Pick the result DataFrame backend for a shared DuckDB executor. + + DataFrameSources determine the backend (and must agree with each other); + pins have no native backend of their own, so an all-pin group follows + PinSource's own convention: polars when available, pandas otherwise. + """ + from ._datasource import DataFrameSource - for name, source in source_items: + shared_lib: str | None = None + for name, source in sources.items(): + if not isinstance(source, DataFrameSource): + continue source_lib = get_dataframe_backend_name(source) - if source_lib != shared_lib: + if shared_lib is None: + shared_lib = source_lib + elif source_lib != shared_lib: raise ValueError( f"Cannot add table '{name}': all DataFrameSources must use " f"the same DataFrame backend. " f"Existing tables use {shared_lib}, new table uses {source_lib}." ) - return shared_lib + if shared_lib is not None: + return shared_lib + + from ._pin_source import _has_polars + + return "polars" if _has_polars() else "pandas" def validate_source_group_compatibility(data_sources: dict[str, DataSource]) -> None: @@ -272,6 +303,35 @@ def check_source_compatibility( first_source = next(iter(existing.values())) + duckdb_family = (DataFrameSource, PinSource) + new_is_duckdb = isinstance(new_source, duckdb_family) + first_is_duckdb = isinstance(first_source, duckdb_family) + + # DataFrameSources and PinSources may mix freely: both materialize their + # tables into a shared DuckDBExecutor connection. + if new_is_duckdb and first_is_duckdb: + if isinstance(new_source, DataFrameSource): + new_lib = get_dataframe_backend_name(new_source) + for source in existing.values(): + if not isinstance(source, DataFrameSource): + continue + existing_lib = get_dataframe_backend_name(source) + if new_lib != existing_lib: + raise ValueError( + f"Cannot add table '{new_name}': all DataFrameSources " + f"must use the same DataFrame backend. " + f"Existing tables use {existing_lib}, new table uses {new_lib}." + ) + return + + if new_is_duckdb != first_is_duckdb: + raise ValueError( + f"Cannot add {type(new_source).__name__} table '{new_name}': " + f"{type(first_source).__name__} tables can only be combined with " + "other tables of the same type. Pins and data frames may be " + "combined with each other, but not with database-backed sources." + ) + if type(new_source) is not type(first_source): raise ValueError( f"Cannot add {type(new_source).__name__} table '{new_name}': " @@ -279,33 +339,6 @@ def check_source_compatibility( f"Existing tables use {type(first_source).__name__}." ) - # Reached only when the existing sources are also PinSources: a second pin - # would validate here but fail at query time, since each pin queries - # through its own private connection and DataSourceExecutor delegates all - # queries to the first one. - if isinstance(new_source, PinSource): - # ValueError like the neighboring checks: this is a group constraint - # violation, not a wrong-argument-type error (contra TRY004). - raise ValueError( # noqa: TRY004 - f"Cannot add pin '{new_name}': only one pin table is supported per " - "chat. Each pin queries through its own DuckDB connection, so only " - "the first pin's table would be queryable. To combine a pin with " - "other tables, register them in a shared DuckDB connection and " - "pass that instead." - ) - - if isinstance(new_source, DataFrameSource) and isinstance( - first_source, DataFrameSource - ): - new_lib = get_dataframe_backend_name(new_source) - existing_lib = get_dataframe_backend_name(first_source) - if new_lib != existing_lib: - raise ValueError( - f"Cannot add table '{new_name}': all DataFrameSources must use " - f"the same DataFrame backend. " - f"Existing tables use {existing_lib}, new table uses {new_lib}." - ) - if ( isinstance(new_source, SQLAlchemySource) and isinstance(first_source, SQLAlchemySource) @@ -337,8 +370,10 @@ def get_dataframe_backend_name(source: DataFrameSource) -> str: def build_query_executor(sources: Mapping[str, DataSource]) -> QueryExecutor: """Pick the executor for a compatible group of sources.""" from ._datasource import DataFrameSource, PolarsLazySource + from ._pin_source import PinSource - # After validation, every source has the same type as the first one. + # After validation, every source has the same type as the first one, + # or the whole group is in the DuckDB family (DataFrameSource/PinSource). validate_source_group_compatibility(dict(sources)) if len(sources) == 1: @@ -346,8 +381,10 @@ def build_query_executor(sources: Mapping[str, DataSource]) -> QueryExecutor: first_source = next(iter(sources.values())) - if isinstance(first_source, DataFrameSource): - return DuckDBExecutor(cast("dict[str, DataFrameSource]", dict(sources))) + if isinstance(first_source, (DataFrameSource, PinSource)): + return DuckDBExecutor( + cast("dict[str, DataFrameSource | PinSource]", dict(sources)) + ) if isinstance(first_source, PolarsLazySource): return PolarsSQLExecutor(cast("dict[str, PolarsLazySource]", dict(sources))) diff --git a/pkg-py/tests/test_pin_source.py b/pkg-py/tests/test_pin_source.py index 56c61262..c22b851d 100644 --- a/pkg-py/tests/test_pin_source.py +++ b/pkg-py/tests/test_pin_source.py @@ -302,9 +302,9 @@ def test_clears_auto_description_on_source_change(self, board, sample_df): class TestMultiplePins: - """A second pin must fail at registration, not at query time.""" + """Multiple pins (and pin + data frame mixes) share a DuckDB executor.""" - def test_second_pin_rejected_by_compatibility_check(self, board, sample_df): + def test_second_pin_passes_compatibility_check(self, board, sample_df): from querychat._query_executor import check_source_compatibility board.pin_write(sample_df, "pin_a", type="parquet") @@ -312,20 +312,164 @@ def test_second_pin_rejected_by_compatibility_check(self, board, sample_df): first = PinSource(board, "pin_a") second = PinSource(board, "pin_b") try: - with pytest.raises(ValueError, match="only one pin"): - check_source_compatibility({"pin_a": first}, second, "pin_b") + check_source_compatibility({"pin_a": first}, second, "pin_b") finally: first.cleanup() second.cleanup() - def test_add_table_rejects_second_pin(self, board, sample_df): + def test_two_pins_query_through_shared_executor(self, board, sample_df): + from querychat import QueryChat + from querychat._query_executor import DuckDBExecutor + + board.pin_write(sample_df, "pin_a", type="parquet") + cities = pd.DataFrame( + {"name": ["Alice", "Diana"], "city": ["Springfield", "Shelbyville"]} + ) + board.pin_write(cities, "pin_b", type="csv") + + qc = QueryChat(board, "pin_a") + try: + qc.add_table(board, "pin_b") + executor = qc._table_set.executor + assert isinstance(executor, DuckDBExecutor) + + result = nw.from_native( + executor.execute_query( + "SELECT a.name, b.city FROM pin_a a " + "JOIN pin_b b ON a.name = b.name ORDER BY a.name" + ) + ) + assert result.rows(named=True) == [ + {"name": "Alice", "city": "Springfield"}, + {"name": "Diana", "city": "Shelbyville"}, + ] + + # Per-table schema and validation work for both pins + metas = executor.get_column_metas("pin_b") + assert {m.name for m in metas} == {"name", "city"} + executor.test_query( + "SELECT * FROM pin_b", table_name="pin_b", require_all_columns=True + ) + finally: + qc.cleanup() + + def test_pin_and_dataframe_mix(self, board, sample_df): + from querychat import QueryChat + from querychat._query_executor import DuckDBExecutor + + board.pin_write(sample_df, "pin_a", type="parquet") + qc = QueryChat(board, "pin_a") + try: + qc.add_table( + pd.DataFrame({"name": ["Bob"], "dept": ["Engineering"]}), + "employees", + ) + executor = qc._table_set.executor + assert isinstance(executor, DuckDBExecutor) + + result = nw.from_native( + executor.execute_query( + "SELECT e.dept FROM pin_a p JOIN employees e ON p.name = e.name" + ) + ) + assert result.rows(named=True) == [{"dept": "Engineering"}] + finally: + qc.cleanup() + + def test_mixed_dataframe_backends_still_rejected_with_pin(self, board, sample_df): + pl = pytest.importorskip("polars") + + from querychat import QueryChat + + board.pin_write(sample_df, "pin_a", type="parquet") + qc = QueryChat(board, "pin_a") + try: + qc.add_table(sample_df, "pandas_table") + with pytest.raises(ValueError, match="same DataFrame backend"): + qc.add_table(pl.DataFrame({"x": [1]}), "polars_table") + finally: + qc.cleanup() + + def test_pin_with_non_sql_safe_name(self, board, sample_df): + """Registry keys that need quoting work in the shared executor.""" + from querychat import QueryChat + + board.pin_write(sample_df, "sales-2026", type="parquet") + board.pin_write(sample_df, "pin_b", type="parquet") + qc = QueryChat(board, "sales-2026") + try: + qc.add_table(board, "pin_b") + executor = qc._table_set.executor + result = nw.from_native( + executor.execute_query('SELECT COUNT(*) AS n FROM "sales-2026"') + ) + assert result.rows(named=True) == [{"n": 4}] + finally: + qc.cleanup() + + def test_shared_executor_uses_version_resolved_at_construction( + self, board, sample_df + ): + """Updating a pin after construction doesn't change the shared table.""" + from querychat import QueryChat + + board.pin_write(sample_df, "pin_a", type="parquet") + board.pin_write(sample_df, "pin_b", type="parquet") + qc = QueryChat(board, "pin_a") + try: + qc.add_table(board, "pin_b") + # New version of pin_a published after its PinSource was built + board.pin_write(sample_df.head(1), "pin_a", type="parquet") + executor = qc._table_set.executor + result = nw.from_native( + executor.execute_query("SELECT COUNT(*) AS n FROM pin_a") + ) + assert result.rows(named=True) == [{"n": 4}] + finally: + qc.cleanup() + + def test_shared_executor_falls_back_when_snapshot_version_is_pruned( + self, tmp_path, sample_df + ): + """ + Non-versioned boards drop old versions on rewrite; the shared table + must still match the source's own copy. + """ + from querychat import QueryChat + + unversioned = pins.board_folder(str(tmp_path / "unversioned"), versioned=False) + unversioned.pin_write(sample_df, "pin_a", type="parquet") + unversioned.pin_write(sample_df, "pin_b", type="parquet") + qc = QueryChat(unversioned, "pin_a") + try: + qc.add_table(unversioned, "pin_b") + unversioned.pin_write(sample_df.head(1), "pin_a", type="parquet") + executor = qc._table_set.executor + result = nw.from_native( + executor.execute_query("SELECT COUNT(*) AS n FROM pin_a") + ) + assert result.rows(named=True) == [{"n": 4}] + finally: + qc.cleanup() + + def test_executor_cleanup_leaves_pin_connections(self, board, sample_df): + """The shared connection is executor-owned; pins keep their own.""" from querychat import QueryChat board.pin_write(sample_df, "pin_a", type="parquet") board.pin_write(sample_df, "pin_b", type="parquet") qc = QueryChat(board, "pin_a") try: - with pytest.raises(ValueError, match="only one pin"): - qc.add_table(board, "pin_b") + qc.add_table(board, "pin_b") + _ = qc._table_set.executor # build the shared executor + qc._table_set.cleanup_executor() + + # Both pins still answer source-level queries through their + # private connections. + for name in ("pin_a", "pin_b"): + result = qc._data_sources[name].execute_query( + f"SELECT COUNT(*) AS n FROM {name}" + ) + assert result.rows(named=True) == [{"n": 4}] finally: qc.cleanup() diff --git a/pkg-r/NEWS.md b/pkg-r/NEWS.md index 985cab1a..4fb61701 100644 --- a/pkg-r/NEWS.md +++ b/pkg-r/NEWS.md @@ -29,7 +29,7 @@ * **Editable SQL panel**: the SQL panel in `querychat_app()` is now a code editor — tweak the generated SQL and apply it with Ctrl/Cmd+Enter. (#265) -* `PinSource`: chat with datasets pinned to a [pins](https://pins.rstudio.com/) board (parquet, CSV, JSON, RDS); the pin's title, description, and tags serve as the default data description. (#246) +* `PinSource`: chat with datasets pinned to a [pins](https://pins.rstudio.com/) board (parquet, CSV, JSON, RDS); the pin's title, description, and tags serve as the default data description. Multiple pins (and pins mixed with data frames) work in one chat via a shared DuckDB connection. (#246, #312) * Deferred construction is more flexible: `table_name` is now optional in `QueryChat$new(NULL)`, and `$server()` gains a `table_name` parameter so the table can be named per session. (#305) diff --git a/pkg-r/R/DataFrameSource.R b/pkg-r/R/DataFrameSource.R index 92eb13aa..91498a29 100644 --- a/pkg-r/R/DataFrameSource.R +++ b/pkg-r/R/DataFrameSource.R @@ -72,6 +72,28 @@ DataFrameSource <- R6::R6Class( private$.conn <- new_dataframe_connection(df, table_name, engine) }, + #' @description + #' Register this data frame in a shared DuckDB connection. + #' + #' Internal hook for joining a shared `DuckDBExecutor`. The caller owns + #' `con` and locks it down once all tables are registered. + #' + #' @param con A DuckDB DBI connection, owned by the caller. + #' @param table_name Name for the table in `con`. Defaults to the source's + #' own table name. + #' + #' @return `NULL` (invisibly) + register_into = function(con, table_name = self$table_name) { + check_installed("duckdb") + duckdb::duckdb_register( + con, + table_name, + self$get_data(), + experimental = FALSE + ) + invisible(NULL) + }, + #' @description #' Disconnect from the database and shut down the DuckDB instance if used. #' diff --git a/pkg-r/R/PinSource.R b/pkg-r/R/PinSource.R index 062eb7d1..e541548f 100644 --- a/pkg-r/R/PinSource.R +++ b/pkg-r/R/PinSource.R @@ -13,6 +13,11 @@ #' When loaded into DuckDB, the connection's external file access is locked #' down so that LLM-generated SQL cannot reach the filesystem. #' +#' Multiple pins (and pins mixed with data frames) can be combined in one +#' chat: every table is materialized into a shared DuckDB connection, so the +#' LLM can join and filter across them. Pins using `engine = "sqlite"` can't +#' join multi-table chats. +#' #' If the pin has a title, description, or tags, [QueryChat] uses them as #' the default `data_description`, which you can override. #' @@ -96,6 +101,15 @@ PinSource <- R6::R6Class( table_name <- sanitize_table_name(table_name) private$.pin_meta <- pins::pin_meta(board, name, version = version) + # Retained so the pin can be re-materialized into a shared DuckDB + # connection when it joins a multi-table executor (register_into()). + private$.board <- board + private$.name <- name + # Snapshot the resolved version so register_into() reads the same pin + # content even if the pin is updated after construction. + private$.version <- private$.pin_meta$local$version %||% version + private$.engine <- engine + pin_type <- private$.pin_meta$type duckdb_file_types <- c("parquet", "csv", "json") use_duckdb_file_read <- engine == "duckdb" && @@ -107,30 +121,7 @@ PinSource <- R6::R6Class( con_owned <- FALSE on.exit(if (!con_owned) DBI::dbDisconnect(con), add = TRUE) - paths <- pins::pin_download(board, name, version = version) - if (length(paths) != 1) { - cli::cli_abort( - "Pin {.val {name}} contains {length(paths)} files, but PinSource requires a single-file pin (as created by {.fn pins::pin_write})." - ) - } - reader_fn <- switch( - pin_type, - parquet = "read_parquet", - csv = "read_csv_auto", - json = "read_json_auto" - ) - if (pin_type == "json") { - DBI::dbExecute(con, "INSTALL json") - DBI::dbExecute(con, "LOAD json") - } - quoted_path <- DBI::dbQuoteLiteral(con, paths[[1]]) - sql <- sprintf( - "CREATE TABLE %s AS SELECT * FROM %s(%s)", - DBI::dbQuoteIdentifier(con, table_name), - reader_fn, - quoted_path - ) - DBI::dbExecute(con, sql) + private$materialize_duckdb_file(con, table_name) duckdb_lock_down(con) } else { if (engine == "sqlite" && pin_type %in% duckdb_file_types) { @@ -141,7 +132,7 @@ PinSource <- R6::R6Class( ) ) } - data <- pins::pin_read(board, name, version = version) + data <- pins::pin_read(board, name, version = private$.version) if (!is.data.frame(data)) { cli::cli_abort( "Pin {.val {name}} contains {.obj_type_friendly {data}}, not a data frame." @@ -156,6 +147,63 @@ PinSource <- R6::R6Class( con_owned <- TRUE }, + #' @description + #' Materialize this pin into a shared DuckDB connection. + #' + #' Internal hook for joining a shared `DuckDBExecutor`. The caller owns + #' `con` and locks it down once all tables are materialized. + #' + #' @param con A DuckDB DBI connection, owned by the caller. + #' @param table_name Name for the table in `con`. Defaults to the pin's + #' own table name. + #' + #' @return `NULL` (invisibly) + register_into = function(con, table_name = self$table_name) { + if (private$.engine != "duckdb") { + cli::cli_abort( + c( + "Pin {.val {private$.name}} uses {.code engine = \"sqlite\"} and cannot join a shared DuckDB executor.", + "i" = "Use {.code engine = \"duckdb\"} to combine pins with other tables." + ) + ) + } + check_installed("duckdb") + + pin_type <- private$.pin_meta$type + tryCatch( + { + if (pin_type %in% c("parquet", "csv", "json")) { + private$materialize_duckdb_file(con, table_name) + } else { + data <- pins::pin_read( + private$.board, + private$.name, + version = private$.version + ) + if (!is.data.frame(data)) { + cli::cli_abort( + "Pin {.val {private$.name}} contains {.obj_type_friendly {data}}, not a data frame." + ) + } + duckdb::duckdb_register(con, table_name, data, experimental = FALSE) + } + }, + # The snapshotted pin version may have been pruned (e.g. a + # non-versioned board rewritten after construction); fall back to + # this source's own copy so the shared table matches the private + # connection. Other materialization failures still raise. + pins_pin_version_missing = function(e) { + duckdb::duckdb_register( + con, + table_name, + self$get_data(), + experimental = FALSE + ) + } + ) + invisible(NULL) + }, + #' @description #' Get a human-readable description of the pin for use in the system prompt. #' @@ -195,7 +243,53 @@ PinSource <- R6::R6Class( invisible(NULL) } ), + active = list( + #' @field engine The database engine backing this pin (`"duckdb"` or + #' `"sqlite"`, read-only). + engine = function() { + private$.engine + } + ), private = list( - .pin_meta = NULL + .pin_meta = NULL, + .board = NULL, + .name = NULL, + .version = NULL, + .engine = NULL, + # Materialize a parquet/CSV/JSON pin as a real table in `con` via DuckDB's + # native file readers. Does not lock the connection down; the caller owns + # `con` and decides when (or whether) to call duckdb_lock_down(). + materialize_duckdb_file = function(con, table_name) { + pin_type <- private$.pin_meta$type + paths <- pins::pin_download( + private$.board, + private$.name, + version = private$.version + ) + if (length(paths) != 1) { + cli::cli_abort( + "Pin {.val {private$.name}} contains {length(paths)} files, but PinSource requires a single-file pin (as created by {.fn pins::pin_write})." + ) + } + reader_fn <- switch( + pin_type, + parquet = "read_parquet", + csv = "read_csv_auto", + json = "read_json_auto" + ) + if (pin_type == "json") { + DBI::dbExecute(con, "INSTALL json") + DBI::dbExecute(con, "LOAD json") + } + quoted_path <- DBI::dbQuoteLiteral(con, paths[[1]]) + sql <- sprintf( + "CREATE TABLE %s AS SELECT * FROM %s(%s)", + DBI::dbQuoteIdentifier(con, table_name), + reader_fn, + quoted_path + ) + DBI::dbExecute(con, sql) + invisible(NULL) + } ) ) diff --git a/pkg-r/R/QueryChatSystemPrompt.R b/pkg-r/R/QueryChatSystemPrompt.R index 5c1072ae..712f521d 100644 --- a/pkg-r/R/QueryChatSystemPrompt.R +++ b/pkg-r/R/QueryChatSystemPrompt.R @@ -130,7 +130,7 @@ QueryChatSystemPrompt <- R6::R6Class( # data_sources may be empty for a greeting with no included tables. has_sources <- length(self$data_sources) > 0 first_source <- if (has_sources) self$data_sources[[1]] else NULL - db_type <- if (has_sources) first_source$get_db_type() else "SQL" + db_type <- if (has_sources) group_db_type(self$data_sources) else "SQL" # Data dicts can carry global (table-less) descriptions, so they may # render even when no tables are selected (e.g. a generic greeting). has_dicts <- length(self$data_dicts) > 0 diff --git a/pkg-r/R/QueryExecutor.R b/pkg-r/R/QueryExecutor.R index 6acca1f3..375257ee 100644 --- a/pkg-r/R/QueryExecutor.R +++ b/pkg-r/R/QueryExecutor.R @@ -59,22 +59,25 @@ DuckDBExecutor <- R6::R6Class( table_columns = list() ), public = list( - initialize = function(dataframes) { + # `data_sources` is a named list of DataFrameSource and/or PinSource + # objects; each materializes its table into one shared connection, then + # the connection is locked down once. + initialize = function(data_sources) { check_installed("duckdb") private$conn <- DBI::dbConnect(duckdb::duckdb(), dbdir = ":memory:") + conn_ok <- FALSE + on.exit( + if (!conn_ok) DBI::dbDisconnect(private$conn, shutdown = TRUE), + add = TRUE + ) - for (name in names(dataframes)) { - duckdb::duckdb_register( - private$conn, - name, - dataframes[[name]], - experimental = FALSE - ) + for (name in names(data_sources)) { + data_sources[[name]]$register_into(private$conn, name) } # Cache column names per table before lockdown - for (name in names(dataframes)) { + for (name in names(data_sources)) { cols <- colnames( DBI::dbGetQuery( private$conn, @@ -88,6 +91,7 @@ DuckDBExecutor <- R6::R6Class( } duckdb_lock_down(private$conn) + conn_ok <- TRUE }, execute_query = function(query) { @@ -230,14 +234,31 @@ build_query_executor <- function(data_sources) { first_source <- data_sources[[1]] - if (inherits(first_source, "DataFrameSource")) { - dataframes <- lapply(data_sources, function(ds) ds$get_data()) - return(DuckDBExecutor$new(dataframes)) + if ( + inherits(first_source, "DataFrameSource") || + inherits(first_source, "PinSource") + ) { + return(DuckDBExecutor$new(data_sources)) } DataSourceExecutor$new(data_sources) } +# DataFrameSources and PinSources can share a DuckDBExecutor connection. +is_duckdb_family_source <- function(x) { + inherits(x, "DataFrameSource") || inherits(x, "PinSource") +} + +# The db_type a group's executor will report. Multi-table +# DataFrameSource/PinSource groups are served by a shared DuckDB connection +# regardless of each source's own engine. +group_db_type <- function(data_sources) { + if (length(data_sources) > 1 && is_duckdb_family_source(data_sources[[1]])) { + return("DuckDB") + } + data_sources[[1]]$get_db_type() +} + # Validates that a new source is compatible with existing sources. check_source_compatibility <- function(existing_sources, new_source, new_name) { if (length(existing_sources) == 0) { @@ -246,25 +267,42 @@ check_source_compatibility <- function(existing_sources, new_source, new_name) { first_source <- existing_sources[[1]] - if (!identical(class(new_source), class(first_source))) { - cli::cli_abort( - c( - "Cannot add {.cls {class(new_source)[1]}} table {.val {new_name}}: all tables must be the same type.", - "i" = "Existing tables use {.cls {class(first_source)[1]}}." + if ( + is_duckdb_family_source(new_source) && is_duckdb_family_source(first_source) + ) { + # Pins materialized into SQLite can't live in a DuckDB executor, so + # multi-table groups containing sqlite-engine pins are rejected. + for (existing_name in names(existing_sources)) { + src <- existing_sources[[existing_name]] + if (inherits(src, "PinSource") && identical(src$engine, "sqlite")) { + cli::cli_abort( + c( + "Cannot add table {.val {new_name}}: pin {.val {existing_name}} uses {.code engine = \"sqlite\"}, which can't join the shared DuckDB connection used for multi-table chats.", + "i" = "Recreate pin {.val {existing_name}} with {.code engine = \"duckdb\"} to combine it with other tables." + ) + ) + } + } + if ( + inherits(new_source, "PinSource") && + identical(new_source$engine, "sqlite") + ) { + cli::cli_abort( + c( + "Cannot add pin {.val {new_name}} with {.code engine = \"sqlite\"}: multi-table chats are served by a shared DuckDB connection, which SQLite pins can't join.", + "i" = "Use {.code engine = \"duckdb\"} to combine pins with other tables." + ) ) - ) + } + return(invisible(NULL)) } - # Reached only when the existing sources are also PinSources: a second pin - # would validate here but fail at query time, since each pin queries - # through its own private connection and DataSourceExecutor delegates all - # queries to the first one. - if (inherits(new_source, "PinSource")) { + if (!identical(class(new_source), class(first_source))) { cli::cli_abort( c( - "Cannot add pin {.val {new_name}}: only one pin table is supported per chat.", - "i" = "Each pin queries through its own connection, so only the first pin's table would be queryable.", - "i" = "To combine a pin with other tables, register them in a shared DuckDB connection and pass that instead." + "Cannot add {.cls {class(new_source)[1]}} table {.val {new_name}}: all tables must be the same type.", + "i" = "Existing tables use {.cls {class(first_source)[1]}}.", + "i" = "Pins and data frames may be combined with each other, but not with database-backed sources." ) ) } diff --git a/pkg-r/R/handoff_chat.R b/pkg-r/R/handoff_chat.R index 6fb16d64..4fa1a641 100644 --- a/pkg-r/R/handoff_chat.R +++ b/pkg-r/R/handoff_chat.R @@ -103,7 +103,7 @@ apply_handoff_max_tokens_override <- function(chat, max_tokens) { } new_model <- ellmer::Model( name = model@name, - params = modifyList(model@params, list(max_tokens = max_tokens)), + params = utils::modifyList(model@params, list(max_tokens = max_tokens)), extra_args = model@extra_args ) chat$initialize( diff --git a/pkg-r/man/DataFrameSource.Rd b/pkg-r/man/DataFrameSource.Rd index b7402f4f..5c41c0a5 100644 --- a/pkg-r/man/DataFrameSource.Rd +++ b/pkg-r/man/DataFrameSource.Rd @@ -45,6 +45,7 @@ df_sqlite$cleanup() \subsection{Public methods}{ \itemize{ \item \href{#method-DataFrameSource-initialize}{\code{DataFrameSource$new()}} + \item \href{#method-DataFrameSource-register_into}{\code{DataFrameSource$register_into()}} \item \href{#method-DataFrameSource-cleanup}{\code{DataFrameSource$cleanup()}} \item \href{#method-DataFrameSource-clone}{\code{DataFrameSource$clone()}} } @@ -94,6 +95,33 @@ engine from duckdb or RSQLite (in that order).} } } +\if{html}{\out{
}} +\if{html}{\out{}} +\if{latex}{\out{\hypertarget{method-DataFrameSource-register_into}{}}} +\subsection{\code{DataFrameSource$register_into()}}{ + Register this data frame in a shared DuckDB connection. + +Internal hook for joining a shared \code{DuckDBExecutor}. The caller owns +\code{con} and locks it down once all tables are registered. + \subsection{Usage}{ + \if{html}{\out{
}} + \preformatted{DataFrameSource$register_into(con, table_name = self$table_name)} + \if{html}{\out{
}} + } + \subsection{Arguments}{ + \if{html}{\out{
}} + \describe{ + \item{\code{con}}{A DuckDB DBI connection, owned by the caller.} + \item{\code{table_name}}{Name for the table in \code{con}. Defaults to the source's +own table name.} + } + \if{html}{\out{
}} + } + \subsection{Returns}{ + \code{NULL} (invisibly) + } +} + \if{html}{\out{
}} \if{html}{\out{}} \if{latex}{\out{\hypertarget{method-DataFrameSource-cleanup}{}}} diff --git a/pkg-r/man/PinSource.Rd b/pkg-r/man/PinSource.Rd index 59289dc5..3e69b5f7 100644 --- a/pkg-r/man/PinSource.Rd +++ b/pkg-r/man/PinSource.Rd @@ -16,6 +16,11 @@ just like \link{DataFrameSource}. When loaded into DuckDB, the connection's external file access is locked down so that LLM-generated SQL cannot reach the filesystem. +Multiple pins (and pins mixed with data frames) can be combined in one +chat: every table is materialized into a shared DuckDB connection, so the +LLM can join and filter across them. Pins using \code{engine = "sqlite"} can't +join multi-table chats. + If the pin has a title, description, or tags, \link{QueryChat} uses them as the default \code{data_description}, which you can override. } @@ -62,10 +67,19 @@ if (rlang::is_installed(c("pins", "duckdb"))) { \section{Super classes}{ \code{\link[querychat:DataSource]{DataSource}} -> \code{\link[querychat:DBISource]{DBISource}} -> \code{PinSource} } +\section{Active bindings}{ + \if{html}{\out{
}} + \describe{ + \item{\code{engine}}{The database engine backing this pin (\code{"duckdb"} or +\code{"sqlite"}, read-only).} + } + \if{html}{\out{
}} +} \section{Methods}{ \subsection{Public methods}{ \itemize{ \item \href{#method-PinSource-initialize}{\code{PinSource$new()}} + \item \href{#method-PinSource-register_into}{\code{PinSource$register_into()}} \item \href{#method-PinSource-get_data_description}{\code{PinSource$get_data_description()}} \item \href{#method-PinSource-cleanup}{\code{PinSource$cleanup()}} \item \href{#method-PinSource-clone}{\code{PinSource$clone()}} @@ -124,6 +138,33 @@ deserialized via \code{pin_read()} instead.} } } +\if{html}{\out{
}} +\if{html}{\out{}} +\if{latex}{\out{\hypertarget{method-PinSource-register_into}{}}} +\subsection{\code{PinSource$register_into()}}{ + Materialize this pin into a shared DuckDB connection. + +Internal hook for joining a shared \code{DuckDBExecutor}. The caller owns +\code{con} and locks it down once all tables are materialized. + \subsection{Usage}{ + \if{html}{\out{
}} + \preformatted{PinSource$register_into(con, table_name = self$table_name)} + \if{html}{\out{
}} + } + \subsection{Arguments}{ + \if{html}{\out{
}} + \describe{ + \item{\code{con}}{A DuckDB DBI connection, owned by the caller.} + \item{\code{table_name}}{Name for the table in \code{con}. Defaults to the pin's +own table name.} + } + \if{html}{\out{
}} + } + \subsection{Returns}{ + \code{NULL} (invisibly) + } +} + \if{html}{\out{
}} \if{html}{\out{}} \if{latex}{\out{\hypertarget{method-PinSource-get_data_description}{}}} diff --git a/pkg-r/tests/testthat/test-PinSource.R b/pkg-r/tests/testthat/test-PinSource.R index f871db94..69c8b37b 100644 --- a/pkg-r/tests/testthat/test-PinSource.R +++ b/pkg-r/tests/testthat/test-PinSource.R @@ -412,22 +412,159 @@ test_that("PinSource$cleanup() disconnects the connection it opened", { expect_false(DBI::dbIsValid(conn)) }) -describe("PinSource multi-table registration", { +describe("PinSource multi-table support", { skip_if_not_installed("pins") skip_if_not_installed("duckdb") skip_if_not_installed("nanoparquet") - it("rejects a second pin via check_source_compatibility", { + it("accepts a second pin via check_source_compatibility", { ps1 <- local_pin_source(name = "pin_a", type = "parquet") ps2 <- local_pin_source(name = "pin_b", type = "parquet") + expect_no_error( + check_source_compatibility(list(pin_a = ps1), ps2, "pin_b") + ) + }) + + it("queries across two pins through a shared DuckDB executor", { + board <- pins::board_temp() + suppressMessages( + pins::pin_write( + board, + data.frame(id = 1:3, x = c(10, 20, 30)), + "pin_a", + type = "parquet" + ) + ) + suppressMessages( + pins::pin_write( + board, + data.frame(id = 1:3, y = c("a", "b", "c")), + "pin_b", + type = "csv" + ) + ) + + qc <- QueryChat$new(board, "pin_a") + withr::defer(qc$cleanup()) + qc$add_table(board, "pin_b") + + executor <- qc$.__enclos_env__$private$.table_set$executor() + expect_s3_class(executor, "DuckDBExecutor") + + joined <- executor$execute_query( + "SELECT a.x, b.y FROM pin_a a JOIN pin_b b ON a.id = b.id ORDER BY a.id" + ) + expect_equal(joined$x, c(10, 20, 30)) + expect_equal(joined$y, c("a", "b", "c")) + + # Per-table schema and validation work for both pins + schema <- executor$get_schema("pin_b", categorical_threshold = 20) + expect_match(schema, "Table: pin_b") + expect_no_error( + executor$test_query( + "SELECT * FROM pin_b", + "pin_b", + require_all_columns = TRUE + ) + ) + }) + + it("mixes pins and data frames in one executor", { + ps1 <- local_pin_source(name = "pin_a", type = "parquet") + df_source <- local_data_frame_source(new_test_df(), "test") + + executor <- build_query_executor(list(pin_a = ps1, test = df_source)) + withr::defer(executor$cleanup()) + + expect_s3_class(executor, "DuckDBExecutor") + result <- executor$execute_query( + "SELECT COUNT(*) AS n FROM pin_a CROSS JOIN test" + ) + expect_equal(result$n, 10 * nrow(new_test_df())) + }) + + it("materializes the snapshotted version on versioned boards", { + board <- pins::board_temp(versioned = TRUE) + suppressMessages( + pins::pin_write(board, data.frame(x = 1:4), "pin_a", type = "parquet") + ) + ps1 <- PinSource$new(board, "pin_a") + withr::defer(ps1$cleanup()) + ps2 <- local_pin_source(name = "pin_b", type = "parquet") + + # New version of pin_a published after its PinSource was built + suppressMessages( + pins::pin_write(board, data.frame(x = 1), "pin_a", type = "parquet") + ) + + executor <- build_query_executor(list(pin_a = ps1, pin_b = ps2)) + withr::defer(executor$cleanup()) + + expect_equal(nrow(executor$execute_query("SELECT * FROM pin_a")), 4) + }) + + it("falls back to the source's copy when the snapshot version is pruned", { + # Non-versioned boards drop the old version on rewrite, so the version + # snapshotted at construction no longer exists at registration time. + board <- pins::board_temp() + suppressMessages( + pins::pin_write(board, data.frame(x = 1:4), "pin_a", type = "parquet") + ) + ps1 <- PinSource$new(board, "pin_a") + withr::defer(ps1$cleanup()) + ps2 <- local_pin_source(name = "pin_b", type = "parquet") + + suppressMessages( + pins::pin_write(board, data.frame(x = 1), "pin_a", type = "parquet") + ) + + executor <- build_query_executor(list(pin_a = ps1, pin_b = ps2)) + withr::defer(executor$cleanup()) + + expect_equal(nrow(executor$execute_query("SELECT * FROM pin_a")), 4) + }) + + it("leaves pin connections open when the executor is cleaned up", { + ps1 <- local_pin_source(name = "pin_a", type = "parquet") + ps2 <- local_pin_source(name = "pin_b", type = "parquet") + executor <- build_query_executor(list(pin_a = ps1, pin_b = ps2)) + + # The shared connection is executor-owned; pins keep their own. + executor$cleanup() + + expect_true(DBI::dbIsValid(ps1$conn)) + expect_true(DBI::dbIsValid(ps2$conn)) + expect_equal(nrow(ps2$execute_query("SELECT * FROM pin_b")), 10) + }) + + it("rejects adding a sqlite-engine pin to a multi-table group", { + skip_if_not_installed("RSQLite") + ps1 <- local_pin_source(name = "pin_a", type = "parquet") + ps_sqlite <- suppressWarnings( + local_pin_source(name = "pin_b", engine = "sqlite") + ) + + expect_error( + check_source_compatibility(list(pin_a = ps1), ps_sqlite, "pin_b"), + "engine = \"sqlite\"" + ) + }) + + it("rejects adding any table when an existing pin uses sqlite", { + skip_if_not_installed("RSQLite") + ps_sqlite <- suppressWarnings( + local_pin_source(name = "pin_a", engine = "sqlite") + ) + df_source <- local_data_frame_source(new_test_df(), "test") + expect_error( - check_source_compatibility(list(pin_a = ps1), ps2, "pin_b"), - "only one pin" + check_source_compatibility(list(pin_a = ps_sqlite), df_source, "test"), + "engine = \"sqlite\"" ) }) - it("QueryChat$add_table() rejects a second pin", { + it("QueryChat$add_table() accepts a second pin", { board <- pins::board_temp() suppressMessages( pins::pin_write(board, mtcars[1:5, ], "pin_a", type = "parquet") @@ -439,6 +576,6 @@ describe("PinSource multi-table registration", { qc <- QueryChat$new(board, "pin_a") withr::defer(qc$cleanup()) - expect_error(qc$add_table(board, "pin_b"), "only one pin") + expect_no_error(qc$add_table(board, "pin_b")) }) }) diff --git a/pkg-r/tests/testthat/test-QueryExecutor.R b/pkg-r/tests/testthat/test-QueryExecutor.R index 8e39e072..16420f75 100644 --- a/pkg-r/tests/testthat/test-QueryExecutor.R +++ b/pkg-r/tests/testthat/test-QueryExecutor.R @@ -35,14 +35,16 @@ describe("DuckDBExecutor", { skip_if_not_installed("duckdb") it("registers multiple data frames for cross-table JOINs", { - users <- new_users_df() scores <- data.frame( id = 1:5, score = c(90, 85, 92, 78, 88), stringsAsFactors = FALSE ) - dataframes <- list(users = users, scores = scores) - executor <- DuckDBExecutor$new(dataframes) + sources <- list( + users = local_data_frame_source(new_users_df(), "users"), + scores = local_data_frame_source(scores, "scores") + ) + executor <- DuckDBExecutor$new(sources) withr::defer(executor$cleanup()) result <- executor$execute_query( @@ -55,9 +57,8 @@ describe("DuckDBExecutor", { }) it("enforces require_all_columns per table in test_query()", { - users <- new_users_df() - dataframes <- list(users = users) - executor <- DuckDBExecutor$new(dataframes) + sources <- list(users = local_data_frame_source(new_users_df(), "users")) + executor <- DuckDBExecutor$new(sources) withr::defer(executor$cleanup()) # Query that drops a column should fail with require_all_columns = TRUE @@ -81,9 +82,8 @@ describe("DuckDBExecutor", { }) it("locks down the connection (DDL like CREATE TABLE should fail)", { - users <- new_users_df() - dataframes <- list(users = users) - executor <- DuckDBExecutor$new(dataframes) + sources <- list(users = local_data_frame_source(new_users_df(), "users")) + executor <- DuckDBExecutor$new(sources) withr::defer(executor$cleanup()) expect_error( @@ -92,14 +92,16 @@ describe("DuckDBExecutor", { }) it("returns correct get_db_type()", { - executor <- DuckDBExecutor$new(list(users = new_users_df())) + sources <- list(users = local_data_frame_source(new_users_df(), "users")) + executor <- DuckDBExecutor$new(sources) withr::defer(executor$cleanup()) expect_equal(executor$get_db_type(), "DuckDB") }) it("gets schema for a named table", { - executor <- DuckDBExecutor$new(list(users = new_users_df())) + sources <- list(users = local_data_frame_source(new_users_df(), "users")) + executor <- DuckDBExecutor$new(sources) withr::defer(executor$cleanup()) schema <- executor$get_schema("users", categorical_threshold = 20) @@ -154,6 +156,26 @@ describe("build_query_executor()", { }) }) +describe("group_db_type()", { + skip_if_not_installed("duckdb") + + it("is DuckDB for multi-source family groups, even with sqlite members", { + skip_if_not_installed("RSQLite") + sqlite_src <- local_data_frame_source( + new_test_df(), + "test", + engine = "sqlite" + ) + duckdb_src <- local_data_frame_source(new_users_df(), "users") + + expect_equal( + group_db_type(list(test = sqlite_src, users = duckdb_src)), + "DuckDB" + ) + expect_equal(group_db_type(list(test = sqlite_src)), "SQLite") + }) +}) + describe("check_source_compatibility()", { skip_if_not_installed("duckdb")