From f1bc26f519fb4d4ecfd136dbd079dd336c3779e7 Mon Sep 17 00:00:00 2001 From: Carson Date: Fri, 10 Jul 2026 11:47:54 -0500 Subject: [PATCH 1/4] fix(pkg-py): give each Streamlit e2e test its own server process The Streamlit e2e fixtures (app_04_streamlit, app_09_streamlit_custom) were module-scoped: one Streamlit subprocess served every test in the class, with each test doing a fresh page.goto() against it. Observed in CI and reproduced locally (including on a commit that predates any related application changes, via an isolated git worktree): after the first test's browser session completed successfully, the server stopped re-running its script for any later session -- the static page shell still loaded, but no reactive content (e.g. the chat greeting) ever streamed in, with no error logged on either side. The condition was permanent for that server process; a fresh browser context/page didn't recover it, only a new server did (confirmed by instrumenting the render path: it ran exactly once across the whole affected test class, including pytest-rerunfailures' automatic reruns). Other e2e fixtures in this file (Shiny, Gradio, Dash) use the same module-scoped-server / fresh-page-per-test pattern without hitting this -- many sequential sessions against one server work fine for those apps in the same CI run, including the *other* Streamlit example (09-streamlit-custom-app.py, which doesn't wait on the chat message in its setup and so never surfaced this). So this isn't a blanket "Streamlit can't handle repeat sessions" issue; the actual trigger wasn't pinned down. Hypothetical fix: give each Streamlit test its own server process (function-scoped fixture) instead of sharing one across a class, trading some subprocess-startup overhead for not accumulating whatever state causes this. Validated against the reported failure locally -- it resolves the hang. Not fully validated: with fresh servers per test, a different set of tests (the ones that submit a chat message and wait on a real LLM response) showed new flakiness in local runs, plausibly from more concurrent OpenAI calls in a shorter window; that wasn't investigated further here. --- pkg-py/tests/playwright/conftest.py | 19 +++++++++++++++---- 1 file changed, 15 insertions(+), 4 deletions(-) diff --git a/pkg-py/tests/playwright/conftest.py b/pkg-py/tests/playwright/conftest.py index 961af01f3..3a2db45d3 100644 --- a/pkg-py/tests/playwright/conftest.py +++ b/pkg-py/tests/playwright/conftest.py @@ -359,9 +359,20 @@ def _stop_streamlit_server(process: subprocess.Popen) -> None: process.kill() -@pytest.fixture(scope="module") +@pytest.fixture def app_04_streamlit() -> Generator[str, None, None]: - """Start the 04-streamlit-app.py Streamlit server for testing.""" + """ + Start the 04-streamlit-app.py Streamlit server for testing. + + Function-scoped (a fresh subprocess per test), unlike the Shiny/Gradio/Dash + app fixtures in this file. A single Streamlit server was observed to stop + re-running its script for new browser sessions after the first test's + session completed successfully -- the static page shell still loads, but + no reactive content (e.g. the chat greeting) ever streams in, with no + error on either side, and the condition persists for the rest of the + server's life (a fresh browser context/page doesn't recover it; only a + new server process does). + """ app_path = str(EXAMPLES_DIR / "04-streamlit-app.py") def start_factory(): @@ -381,9 +392,9 @@ def streamlit_cleanup(process, _): _stop_streamlit_server(process) -@pytest.fixture(scope="module") +@pytest.fixture def app_09_streamlit_custom() -> Generator[str, None, None]: - """Start the 09-streamlit-custom-app.py Streamlit server for testing.""" + """Start the 09-streamlit-custom-app.py Streamlit server for testing. Function-scoped; see app_04_streamlit().""" app_path = str(EXAMPLES_DIR / "09-streamlit-custom-app.py") def start_factory(): From a73d23d24567d3db09adb3b5925c408e775679b2 Mon Sep 17 00:00:00 2001 From: Carson Date: Fri, 10 Jul 2026 14:34:21 -0500 Subject: [PATCH 2/4] fix(pkg-py): fix segfault behind Streamlit e2e flakiness The e2e hang/failure symptom ("SQL never updates", no exception anywhere) was a Streamlit subprocess segfault, invisible because its stdout/stderr are redirected to DEVNULL. Confirmed with PYTHONFAULTHANDLER and traced to two compounding issues: 1. The Streamlit example apps reconstructed QueryChat (and its DuckDB connection + polars data registration) on every script rerun. Streamlit's ScriptRunner runs each rerun in a fresh OS thread, so this repeatedly touched polars'/DuckDB's native code from new threads. 2. pyarrow==25.0.0 (and polars' own Arrow conversion) segfaults when its native conversion code is invoked from more than ~2 distinct threads over a process's lifetime -- reproduced in isolation with no Streamlit or querychat code involved, bisected away by pinning pyarrow==21.0.0 / polars==1.20.0. Fixes: - Pin pyarrow/polars to versions that don't exhibit the segfault. - Cache QueryChat construction in both Streamlit examples via st.cache_resource so it's built once instead of on every rerun. - DataFrameSource/DuckDBExecutor/PinSource now register a materialized pyarrow.Table with DuckDB instead of the raw polars object, and open a fresh short-lived connection per query instead of holding one across threads -- defense in depth against the same crash class. --- pkg-py/examples/04-streamlit-app.py | 10 +++- pkg-py/examples/09-streamlit-custom-app.py | 8 ++- pkg-py/src/querychat/_datasource.py | 54 ++++++++++++------ pkg-py/src/querychat/_pin_source.py | 7 ++- pkg-py/src/querychat/_query_executor.py | 66 ++++++++++++++-------- pkg-py/tests/test_multi_table.py | 17 ++++-- pyproject.toml | 4 +- 7 files changed, 114 insertions(+), 52 deletions(-) diff --git a/pkg-py/examples/04-streamlit-app.py b/pkg-py/examples/04-streamlit-app.py index a04b189ba..6c60950b2 100644 --- a/pkg-py/examples/04-streamlit-app.py +++ b/pkg-py/examples/04-streamlit-app.py @@ -7,10 +7,18 @@ from pathlib import Path +import streamlit as st + from querychat.data import titanic from querychat.streamlit import QueryChat greeting = Path(__file__).parent / "greeting.md" -qc = QueryChat(titanic(), "titanic", greeting=greeting) + +@st.cache_resource +def _get_qc() -> QueryChat: + return QueryChat(titanic(), "titanic", greeting=greeting) + + +qc = _get_qc() qc.app() diff --git a/pkg-py/examples/09-streamlit-custom-app.py b/pkg-py/examples/09-streamlit-custom-app.py index d3f763b2a..0ed974b58 100644 --- a/pkg-py/examples/09-streamlit-custom-app.py +++ b/pkg-py/examples/09-streamlit-custom-app.py @@ -22,7 +22,13 @@ initial_sidebar_state="expanded", ) -qc = QueryChat(titanic(), "titanic", greeting=greeting) + +@st.cache_resource +def _get_qc() -> QueryChat: + return QueryChat(titanic(), "titanic", greeting=greeting) + + +qc = _get_qc() qc.sidebar() st.title("Titanic Data Explorer") diff --git a/pkg-py/src/querychat/_datasource.py b/pkg-py/src/querychat/_datasource.py index 556bfac69..475c358de 100644 --- a/pkg-py/src/querychat/_datasource.py +++ b/pkg-py/src/querychat/_datasource.py @@ -348,13 +348,30 @@ def __init__(self, df: nw.DataFrame, table_name: str): native_namespace = nw.get_native_namespace(df) self._df_lib = native_namespace.__name__ - self._conn = duckdb.connect(database=":memory:") - # NOTE: if native representation is polars, pyarrow is required for registration - self._conn.register(table_name, self._df.to_native()) - duckdb_lock_down(self._conn) - - # Store original column names for validation - self._colnames = list(self._df.columns) + # Materialize a pyarrow.Table once rather than registering the native (e.g. + # polars) object with DuckDB on every query: DuckDB's polars registration + # pulls batches through polars' own Rust runtime on whatever thread issues + # the query, which segfaults if that thread differs from the one that + # created the data (e.g. Streamlit reruns the script in a fresh thread + # each time). A materialized Arrow table has no such thread affinity, so + # each query below opens its own short-lived connection around it instead + # of sharing one DuckDB connection across threads. + self._arrow_table = self._df.to_arrow() + self._closed = False + + # Validate the data registers cleanly and cache column names up front. + with self._connect() as conn: + result = conn.execute(f'SELECT * FROM "{table_name}" LIMIT 0') + self._colnames = [desc[0] for desc in result.description] + + def _connect(self) -> duckdb.DuckDBPyConnection: + """Open a fresh, locked-down connection registered with this source's data.""" + if self._closed: + raise duckdb.ConnectionException("Connection already closed!") + conn = duckdb.connect(database=":memory:") + conn.register(self.table_name, self._arrow_table) + duckdb_lock_down(conn) + return conn def get_db_type(self) -> str: """ @@ -389,13 +406,15 @@ def get_schema(self, *, categorical_threshold: int) -> str: return format_schema(self.table_name, metas) def get_column_metas(self) -> list[ColumnMeta]: - result = self._conn.execute(f'SELECT * FROM "{self.table_name}" LIMIT 0') - return [duckdb_column_meta(desc[0], desc[1]) for desc in result.description] + with self._connect() as conn: + result = conn.execute(f'SELECT * FROM "{self.table_name}" LIMIT 0') + return [duckdb_column_meta(desc[0], desc[1]) for desc in result.description] def populate_column_stats( self, columns: list[ColumnMeta], categorical_threshold: int ) -> None: - duckdb_column_stats(self._conn, self.table_name, columns, categorical_threshold) + with self._connect() as conn: + duckdb_column_stats(conn, self.table_name, columns, categorical_threshold) def execute_query(self, query: str) -> IntoDataFrameT: """ @@ -420,8 +439,9 @@ def execute_query(self, query: str) -> IntoDataFrameT: """ check_query(query) - result = self._conn.execute(query) - return self._convert_result(result) + with self._connect() as conn: + result = conn.execute(query) + return self._convert_result(result) def _convert_result(self, result: duckdb.DuckDBPyConnection) -> IntoDataFrameT: """ @@ -472,8 +492,9 @@ def test_query( """ check_query(query) - result = self._conn.execute(f"{query} LIMIT 1") - native_result = self._convert_result(result) + with self._connect() as conn: + result = conn.execute(f"{query} LIMIT 1") + native_result = self._convert_result(result) if require_all_columns: wrapped = nw.from_native(native_result) @@ -506,15 +527,14 @@ def get_data(self) -> IntoDataFrameT: def cleanup(self) -> None: """ - Close the DuckDB connection. + Mark this source closed; further queries raise duckdb.ConnectionException. Returns ------- None """ - if self._conn: - self._conn.close() + self._closed = True class SQLAlchemySource(DataSource[nw.DataFrame]): diff --git a/pkg-py/src/querychat/_pin_source.py b/pkg-py/src/querychat/_pin_source.py index b145d0e9d..4ee499e1b 100644 --- a/pkg-py/src/querychat/_pin_source.py +++ b/pkg-py/src/querychat/_pin_source.py @@ -150,9 +150,12 @@ def __init__( 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]) + # Materialize to a pyarrow.Table rather than registering the + # polars object directly -- see the comment in + # DataFrameSource.__init__. + arrow_tbl = pl.read_ipc(paths[0]).to_arrow() vname = f"__pin_staging_{effective_table_name}" - conn.register(vname, arrow_df) + conn.register(vname, arrow_tbl) conn.execute( f'CREATE TABLE "{effective_table_name}" AS SELECT * FROM "{vname}"' ) diff --git a/pkg-py/src/querychat/_query_executor.py b/pkg-py/src/querychat/_query_executor.py index f1c4e2287..79a76e63e 100644 --- a/pkg-py/src/querychat/_query_executor.py +++ b/pkg-py/src/querychat/_query_executor.py @@ -79,23 +79,38 @@ class DuckDBExecutor(QueryExecutor): 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()) - - # Cache column names per table before lockdown + # Materialize each source's data once rather than registering the native + # (e.g. polars) object on every query -- see the comment in + # DataFrameSource.__init__. Each query below opens its own short-lived + # connection around these tables instead of sharing one DuckDB connection + # across threads. + self._arrow_tables = { + name: source._df.to_arrow() for name, source in sources.items() + } + self._closed = False + + # Cache column names per table 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] - - duckdb_lock_down(self._conn) + with self._connect() as conn: + for name in sources: + result = conn.execute(f'SELECT * FROM "{name}" LIMIT 0') + self._table_columns[name] = [desc[0] for desc in result.description] + + def _connect(self) -> duckdb.DuckDBPyConnection: + """Open a fresh, locked-down connection registered with all tables.""" + if self._closed: + raise duckdb.ConnectionException("Connection already closed!") + conn = duckdb.connect(database=":memory:") + for name, arrow_table in self._arrow_tables.items(): + conn.register(name, arrow_table) + duckdb_lock_down(conn) + return conn def execute_query(self, query: str) -> Any: check_query(query) - result = self._conn.execute(query) - return self._convert_result(result) + with self._connect() as conn: + result = conn.execute(query) + return self._convert_result(result) def _convert_result(self, result: duckdb.DuckDBPyConnection) -> Any: if self._df_lib == "polars": @@ -114,29 +129,32 @@ def test_query( self, query: str, *, table_name: str, require_all_columns: bool = False ) -> None: check_query(query) - result = self._conn.execute(f"{query} LIMIT 1") + with self._connect() as conn: + result = conn.execute(f"{query} LIMIT 1") - if require_all_columns: - result_columns = {desc[0] for desc in result.description} - self._validate_missing_columns( - result_columns, self._table_columns[table_name] - ) + if require_all_columns: + result_columns = {desc[0] for desc in result.description} + self._validate_missing_columns( + result_columns, self._table_columns[table_name] + ) def get_db_type(self) -> str: return "DuckDB" def cleanup(self) -> None: - if self._conn: - self._conn.close() + """Mark this executor closed; further queries raise duckdb.ConnectionException.""" + self._closed = True def get_column_metas(self, table_name: str) -> list[ColumnMeta]: - result = self._conn.execute(f'SELECT * FROM "{table_name}" LIMIT 0') - return [duckdb_column_meta(desc[0], desc[1]) for desc in result.description] + with self._connect() as conn: + result = conn.execute(f'SELECT * FROM "{table_name}" LIMIT 0') + return [duckdb_column_meta(desc[0], desc[1]) for desc in result.description] def populate_column_stats( self, table_name: str, columns: list[ColumnMeta], categorical_threshold: int ) -> None: - duckdb_column_stats(self._conn, table_name, columns, categorical_threshold) + with self._connect() as conn: + duckdb_column_stats(conn, table_name, columns, categorical_threshold) class PolarsSQLExecutor(QueryExecutor): diff --git a/pkg-py/tests/test_multi_table.py b/pkg-py/tests/test_multi_table.py index 2e42aed9b..76eb5a3c0 100644 --- a/pkg-py/tests/test_multi_table.py +++ b/pkg-py/tests/test_multi_table.py @@ -293,17 +293,24 @@ class TestMultiTableCleanup: def test_cleanup_all_sources(self, orders_df, customers_df): """Test that cleanup() cleans up all data sources.""" + import duckdb + qc = QueryChat(orders_df, "orders", greeting="Hello!") qc.add_table(customers_df, "customers") + orders_source = qc._data_sources["orders"] + customers_source = qc._data_sources["customers"] - # Both sources should have connections before cleanup - assert qc._data_sources["orders"]._conn is not None - assert qc._data_sources["customers"]._conn is not None + # Both sources should be usable before cleanup + orders_source.execute_query("SELECT * FROM orders LIMIT 1") + customers_source.execute_query("SELECT * FROM customers LIMIT 1") qc.cleanup() - # Connections should be closed after cleanup - # (DuckDB connections don't have is_closed, but they're closed) + # Both sources should be closed after cleanup + with pytest.raises(duckdb.ConnectionException): + orders_source.execute_query("SELECT * FROM orders LIMIT 1") + with pytest.raises(duckdb.ConnectionException): + customers_source.execute_query("SELECT * FROM customers LIMIT 1") @pytest.fixture diff --git a/pyproject.toml b/pyproject.toml index bec6902f5..3471701dd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -100,8 +100,8 @@ dev = [ "pyright>=1.1.401", "tox-uv>=1.11.4", "pytest>=8.4.0", - "polars>=1.0.0", - "pyarrow>=14.0.0", + "polars==1.20.0", + "pyarrow==21.0.0", "ibis-framework[duckdb]>=9.0.0", "ggsql>=0.3.2", "altair>=6.0", From b06ae8538ad3e870f7ed67fd5dbcff6eafd7ec72 Mon Sep 17 00:00:00 2001 From: Carson Date: Fri, 10 Jul 2026 14:44:25 -0500 Subject: [PATCH 3/4] fix(pkg-py): loosen pyarrow/polars pins to a version range Bisected the pyarrow segfault (see previous commit) to a regression introduced in pyarrow 25.0.0 specifically -- 21.0.0 through 24.0.0 are all fine when their native conversion code is invoked from many OS threads. polars needs no pin at all; the crash was entirely in pyarrow's own code. Replace the exact `pyarrow==21.0.0`/`polars==1.20.0` pins with `pyarrow<25.0.0` (polars left unconstrained), and add the same bound to the public `polars` extra in [project.optional-dependencies] -- the exact pins only lived in the dev/test dependency group, so real users installing `querychat[polars]` were still exposed to the same segfault in their own Streamlit apps. --- pyproject.toml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 3471701dd..cddceb52c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,7 +43,7 @@ classifiers = [ [project.optional-dependencies] # For SQLAlchemySource and sample data, one of polars or pandas is required pandas = ["pandas"] -polars = ["polars", "pyarrow"] # duckdb requires pyarrow for polars DataFrame registration +polars = ["polars", "pyarrow<25.0.0"] # duckdb requires pyarrow for polars DataFrame registration; pyarrow>=25.0.0 segfaults when its native conversion code runs on more than ~2 OS threads over a process's lifetime (e.g. Streamlit reruns each script on a fresh thread) -- see #268 ibis = ["ibis-framework>=9.0.0", "pandas"] # pandas required for ibis .execute() to return DataFrames pins = ["pins>=0.9.1"] # Web framework extras @@ -100,8 +100,8 @@ dev = [ "pyright>=1.1.401", "tox-uv>=1.11.4", "pytest>=8.4.0", - "polars==1.20.0", - "pyarrow==21.0.0", + "polars>=1.0.0", + "pyarrow<25.0.0", # see comment on the `polars` extra above "ibis-framework[duckdb]>=9.0.0", "ggsql>=0.3.2", "altair>=6.0", From bbb4821cc3790e581714308a9c9f11c418e16365 Mon Sep 17 00:00:00 2001 From: Carson Date: Fri, 10 Jul 2026 15:07:53 -0500 Subject: [PATCH 4/4] fix(pkg-py): revert to a minimal, pin-only fix Verified across multiple full e2e-suite runs that the pyarrow<25.0.0 pin alone is sufficient -- every segfault traced back to this single regression. The st.cache_resource caching in the Streamlit examples and the DataFrameSource/DuckDBExecutor/PinSource rewrite to materialize pyarrow.Table and use short-lived connections were not fixing anything the pin doesn't already fix, so revert them to keep this release-blocking fix as small as possible. --- pkg-py/examples/04-streamlit-app.py | 10 +--- pkg-py/examples/09-streamlit-custom-app.py | 8 +-- pkg-py/src/querychat/_datasource.py | 54 ++++++------------ pkg-py/src/querychat/_pin_source.py | 7 +-- pkg-py/src/querychat/_query_executor.py | 66 ++++++++-------------- pkg-py/tests/test_multi_table.py | 17 ++---- 6 files changed, 50 insertions(+), 112 deletions(-) diff --git a/pkg-py/examples/04-streamlit-app.py b/pkg-py/examples/04-streamlit-app.py index 6c60950b2..a04b189ba 100644 --- a/pkg-py/examples/04-streamlit-app.py +++ b/pkg-py/examples/04-streamlit-app.py @@ -7,18 +7,10 @@ from pathlib import Path -import streamlit as st - from querychat.data import titanic from querychat.streamlit import QueryChat greeting = Path(__file__).parent / "greeting.md" - -@st.cache_resource -def _get_qc() -> QueryChat: - return QueryChat(titanic(), "titanic", greeting=greeting) - - -qc = _get_qc() +qc = QueryChat(titanic(), "titanic", greeting=greeting) qc.app() diff --git a/pkg-py/examples/09-streamlit-custom-app.py b/pkg-py/examples/09-streamlit-custom-app.py index 0ed974b58..d3f763b2a 100644 --- a/pkg-py/examples/09-streamlit-custom-app.py +++ b/pkg-py/examples/09-streamlit-custom-app.py @@ -22,13 +22,7 @@ initial_sidebar_state="expanded", ) - -@st.cache_resource -def _get_qc() -> QueryChat: - return QueryChat(titanic(), "titanic", greeting=greeting) - - -qc = _get_qc() +qc = QueryChat(titanic(), "titanic", greeting=greeting) qc.sidebar() st.title("Titanic Data Explorer") diff --git a/pkg-py/src/querychat/_datasource.py b/pkg-py/src/querychat/_datasource.py index 475c358de..556bfac69 100644 --- a/pkg-py/src/querychat/_datasource.py +++ b/pkg-py/src/querychat/_datasource.py @@ -348,30 +348,13 @@ def __init__(self, df: nw.DataFrame, table_name: str): native_namespace = nw.get_native_namespace(df) self._df_lib = native_namespace.__name__ - # Materialize a pyarrow.Table once rather than registering the native (e.g. - # polars) object with DuckDB on every query: DuckDB's polars registration - # pulls batches through polars' own Rust runtime on whatever thread issues - # the query, which segfaults if that thread differs from the one that - # created the data (e.g. Streamlit reruns the script in a fresh thread - # each time). A materialized Arrow table has no such thread affinity, so - # each query below opens its own short-lived connection around it instead - # of sharing one DuckDB connection across threads. - self._arrow_table = self._df.to_arrow() - self._closed = False - - # Validate the data registers cleanly and cache column names up front. - with self._connect() as conn: - result = conn.execute(f'SELECT * FROM "{table_name}" LIMIT 0') - self._colnames = [desc[0] for desc in result.description] - - def _connect(self) -> duckdb.DuckDBPyConnection: - """Open a fresh, locked-down connection registered with this source's data.""" - if self._closed: - raise duckdb.ConnectionException("Connection already closed!") - conn = duckdb.connect(database=":memory:") - conn.register(self.table_name, self._arrow_table) - duckdb_lock_down(conn) - return conn + self._conn = duckdb.connect(database=":memory:") + # NOTE: if native representation is polars, pyarrow is required for registration + self._conn.register(table_name, self._df.to_native()) + duckdb_lock_down(self._conn) + + # Store original column names for validation + self._colnames = list(self._df.columns) def get_db_type(self) -> str: """ @@ -406,15 +389,13 @@ def get_schema(self, *, categorical_threshold: int) -> str: return format_schema(self.table_name, metas) def get_column_metas(self) -> list[ColumnMeta]: - with self._connect() as conn: - result = conn.execute(f'SELECT * FROM "{self.table_name}" LIMIT 0') - return [duckdb_column_meta(desc[0], desc[1]) for desc in result.description] + result = self._conn.execute(f'SELECT * FROM "{self.table_name}" LIMIT 0') + return [duckdb_column_meta(desc[0], desc[1]) for desc in result.description] def populate_column_stats( self, columns: list[ColumnMeta], categorical_threshold: int ) -> None: - with self._connect() as conn: - duckdb_column_stats(conn, self.table_name, columns, categorical_threshold) + duckdb_column_stats(self._conn, self.table_name, columns, categorical_threshold) def execute_query(self, query: str) -> IntoDataFrameT: """ @@ -439,9 +420,8 @@ def execute_query(self, query: str) -> IntoDataFrameT: """ check_query(query) - with self._connect() as conn: - result = conn.execute(query) - return self._convert_result(result) + result = self._conn.execute(query) + return self._convert_result(result) def _convert_result(self, result: duckdb.DuckDBPyConnection) -> IntoDataFrameT: """ @@ -492,9 +472,8 @@ def test_query( """ check_query(query) - with self._connect() as conn: - result = conn.execute(f"{query} LIMIT 1") - native_result = self._convert_result(result) + result = self._conn.execute(f"{query} LIMIT 1") + native_result = self._convert_result(result) if require_all_columns: wrapped = nw.from_native(native_result) @@ -527,14 +506,15 @@ def get_data(self) -> IntoDataFrameT: def cleanup(self) -> None: """ - Mark this source closed; further queries raise duckdb.ConnectionException. + Close the DuckDB connection. Returns ------- None """ - self._closed = True + if self._conn: + self._conn.close() class SQLAlchemySource(DataSource[nw.DataFrame]): diff --git a/pkg-py/src/querychat/_pin_source.py b/pkg-py/src/querychat/_pin_source.py index 4ee499e1b..b145d0e9d 100644 --- a/pkg-py/src/querychat/_pin_source.py +++ b/pkg-py/src/querychat/_pin_source.py @@ -150,12 +150,9 @@ def __init__( f"Pin '{name}' contains {len(paths)} files, but PinSource " "requires a single-file pin (as created by pin_write())." ) - # Materialize to a pyarrow.Table rather than registering the - # polars object directly -- see the comment in - # DataFrameSource.__init__. - arrow_tbl = pl.read_ipc(paths[0]).to_arrow() + arrow_df = pl.read_ipc(paths[0]) vname = f"__pin_staging_{effective_table_name}" - conn.register(vname, arrow_tbl) + conn.register(vname, arrow_df) conn.execute( f'CREATE TABLE "{effective_table_name}" AS SELECT * FROM "{vname}"' ) diff --git a/pkg-py/src/querychat/_query_executor.py b/pkg-py/src/querychat/_query_executor.py index 79a76e63e..f1c4e2287 100644 --- a/pkg-py/src/querychat/_query_executor.py +++ b/pkg-py/src/querychat/_query_executor.py @@ -79,38 +79,23 @@ class DuckDBExecutor(QueryExecutor): def __init__(self, sources: dict[str, DataFrameSource]): self._df_lib = get_shared_dataframe_backend(sources) - # Materialize each source's data once rather than registering the native - # (e.g. polars) object on every query -- see the comment in - # DataFrameSource.__init__. Each query below opens its own short-lived - # connection around these tables instead of sharing one DuckDB connection - # across threads. - self._arrow_tables = { - name: source._df.to_arrow() for name, source in sources.items() - } - self._closed = False - - # Cache column names per table + self._conn = duckdb.connect(database=":memory:") + + for name, source in sources.items(): + self._conn.register(name, source.get_data()) + + # Cache column names per table before lockdown self._table_columns: dict[str, list[str]] = {} - with self._connect() as conn: - for name in sources: - result = conn.execute(f'SELECT * FROM "{name}" LIMIT 0') - self._table_columns[name] = [desc[0] for desc in result.description] - - def _connect(self) -> duckdb.DuckDBPyConnection: - """Open a fresh, locked-down connection registered with all tables.""" - if self._closed: - raise duckdb.ConnectionException("Connection already closed!") - conn = duckdb.connect(database=":memory:") - for name, arrow_table in self._arrow_tables.items(): - conn.register(name, arrow_table) - duckdb_lock_down(conn) - return conn + 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] + + duckdb_lock_down(self._conn) def execute_query(self, query: str) -> Any: check_query(query) - with self._connect() as conn: - result = conn.execute(query) - return self._convert_result(result) + result = self._conn.execute(query) + return self._convert_result(result) def _convert_result(self, result: duckdb.DuckDBPyConnection) -> Any: if self._df_lib == "polars": @@ -129,32 +114,29 @@ def test_query( self, query: str, *, table_name: str, require_all_columns: bool = False ) -> None: check_query(query) - with self._connect() as conn: - result = conn.execute(f"{query} LIMIT 1") + result = self._conn.execute(f"{query} LIMIT 1") - if require_all_columns: - result_columns = {desc[0] for desc in result.description} - self._validate_missing_columns( - result_columns, self._table_columns[table_name] - ) + if require_all_columns: + result_columns = {desc[0] for desc in result.description} + self._validate_missing_columns( + result_columns, self._table_columns[table_name] + ) def get_db_type(self) -> str: return "DuckDB" def cleanup(self) -> None: - """Mark this executor closed; further queries raise duckdb.ConnectionException.""" - self._closed = True + if self._conn: + self._conn.close() def get_column_metas(self, table_name: str) -> list[ColumnMeta]: - with self._connect() as conn: - result = conn.execute(f'SELECT * FROM "{table_name}" LIMIT 0') - return [duckdb_column_meta(desc[0], desc[1]) for desc in result.description] + result = self._conn.execute(f'SELECT * FROM "{table_name}" LIMIT 0') + return [duckdb_column_meta(desc[0], desc[1]) for desc in result.description] def populate_column_stats( self, table_name: str, columns: list[ColumnMeta], categorical_threshold: int ) -> None: - with self._connect() as conn: - duckdb_column_stats(conn, table_name, columns, categorical_threshold) + duckdb_column_stats(self._conn, table_name, columns, categorical_threshold) class PolarsSQLExecutor(QueryExecutor): diff --git a/pkg-py/tests/test_multi_table.py b/pkg-py/tests/test_multi_table.py index 76eb5a3c0..2e42aed9b 100644 --- a/pkg-py/tests/test_multi_table.py +++ b/pkg-py/tests/test_multi_table.py @@ -293,24 +293,17 @@ class TestMultiTableCleanup: def test_cleanup_all_sources(self, orders_df, customers_df): """Test that cleanup() cleans up all data sources.""" - import duckdb - qc = QueryChat(orders_df, "orders", greeting="Hello!") qc.add_table(customers_df, "customers") - orders_source = qc._data_sources["orders"] - customers_source = qc._data_sources["customers"] - # Both sources should be usable before cleanup - orders_source.execute_query("SELECT * FROM orders LIMIT 1") - customers_source.execute_query("SELECT * FROM customers LIMIT 1") + # Both sources should have connections before cleanup + assert qc._data_sources["orders"]._conn is not None + assert qc._data_sources["customers"]._conn is not None qc.cleanup() - # Both sources should be closed after cleanup - with pytest.raises(duckdb.ConnectionException): - orders_source.execute_query("SELECT * FROM orders LIMIT 1") - with pytest.raises(duckdb.ConnectionException): - customers_source.execute_query("SELECT * FROM customers LIMIT 1") + # Connections should be closed after cleanup + # (DuckDB connections don't have is_closed, but they're closed) @pytest.fixture