diff --git a/pkg-py/CHANGELOG.md b/pkg-py/CHANGELOG.md index 36455431..c0d50cdd 100644 --- a/pkg-py/CHANGELOG.md +++ b/pkg-py/CHANGELOG.md @@ -5,6 +5,10 @@ All notable changes to this project will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [UNRELEASED] + +* Fixed an issue with the query tool when used with SQLAlchemy data sources. (@npelikan #79) + ## [0.2.0] - 2025-09-02 * `querychat.init()` now accepts a `client` argument, replacing the previous `create_chat_callback` argument. (#60) diff --git a/pkg-py/src/querychat/querychat.py b/pkg-py/src/querychat/querychat.py index 3f6ea616..f449ede8 100644 --- a/pkg-py/src/querychat/querychat.py +++ b/pkg-py/src/querychat/querychat.py @@ -224,11 +224,15 @@ def df_to_html(df: IntoFrame, maxrows: int = 5) -> str: HTML string representation of the table """ - if isinstance(df, (nw.LazyFrame, nw.DataFrame)): - df_short = df.lazy().head(maxrows).collect() - nrow_full = df.lazy().select(nw.len()).collect().item() + ndf = nw.from_native(df) + + if isinstance(ndf, (nw.LazyFrame, nw.DataFrame)): + df_short = ndf.lazy().head(maxrows).collect() + nrow_full = ndf.lazy().select(nw.len()).collect().item() else: - raise TypeError("df must be a Narwhals DataFrame or LazyFrame") + raise TypeError( + "Must be able to convert `df` into a Narwhals DataFrame or LazyFrame", + ) # Generate HTML table table_html = df_short.to_pandas().to_html( diff --git a/pkg-py/tests/test_df_to_html.py b/pkg-py/tests/test_df_to_html.py new file mode 100644 index 00000000..63d306c1 --- /dev/null +++ b/pkg-py/tests/test_df_to_html.py @@ -0,0 +1,116 @@ +import sqlite3 +import tempfile +from pathlib import Path + +import pandas as pd +import pytest +from sqlalchemy import create_engine +from src.querychat.datasource import DataFrameSource, SQLAlchemySource +from src.querychat.querychat import df_to_html + + +@pytest.fixture +def sample_dataframe(): + """Create a sample pandas DataFrame for testing.""" + return pd.DataFrame( + { + "id": [1, 2, 3, 4, 5], + "name": ["Alice", "Bob", "Charlie", "Diana", "Eve"], + "age": [25, 30, 35, 28, 32], + "salary": [50000, 60000, 70000, 55000, 65000], + }, + ) + + +@pytest.fixture +def sample_sqlite(): + """Create a temporary SQLite database with test data.""" + temp_db = tempfile.NamedTemporaryFile(delete=False, suffix=".db") # noqa: SIM115 + temp_db.close() + + conn = sqlite3.connect(temp_db.name) + cursor = conn.cursor() + + cursor.execute(""" + CREATE TABLE employees ( + id INTEGER PRIMARY KEY, + name TEXT, + age INTEGER, + salary REAL + ) + """) + + test_data = [ + (1, "Alice", 25, 50000), + (2, "Bob", 30, 60000), + (3, "Charlie", 35, 70000), + (4, "Diana", 28, 55000), + (5, "Eve", 32, 65000), + ] + + cursor.executemany( + "INSERT INTO employees (id, name, age, salary) VALUES (?, ?, ?, ?)", + test_data, + ) + + conn.commit() + conn.close() + + engine = create_engine(f"sqlite:///{temp_db.name}") + yield engine + + # Cleanup + Path(temp_db.name).unlink() + + +def test_df_to_html_with_dataframe_source_result(sample_dataframe): + """Test that df_to_html() works with results from DataFrameSource.execute_query().""" + source = DataFrameSource(sample_dataframe, "employees") + + # Execute query to get pandas DataFrame + result_df = source.execute_query("SELECT * FROM employees WHERE age > 25") + + # This should succeed after the fix + html_output = df_to_html(result_df) + + # Verify the HTML contains expected content + assert isinstance(html_output, str) + assert "