diff --git a/lumen/sources/duckdb.py b/lumen/sources/duckdb.py index 1db0bc953..320bf123f 100644 --- a/lumen/sources/duckdb.py +++ b/lumen/sources/duckdb.py @@ -10,6 +10,7 @@ import numpy.core.multiarray # noqa: F401 import pandas as pd import param +import sqlglot from ..config import config from ..serializers import Serializer @@ -362,13 +363,35 @@ def create_sql_expr_source( params = {} source_params = dict(self.param.values(), **kwargs) - preserved_tables = {} - for table_name, sql_expr in tables.items(): - if table_name in self._file_based_tables: - preserved_tables[table_name] = self._file_based_tables[table_name] - else: - preserved_tables[table_name] = sql_expr - source_params['tables'] = preserved_tables + + # Only preserve existing tables if reusing the connection + # If uri or initializers changed, start fresh with only new tables + all_tables = tables + if 'uri' not in kwargs and 'initializers' not in kwargs: + # Reuse connection - start with ALL existing tables (upsert behavior) + # Only applies when self.tables is a dict (list-based tables don't have SQL expressions) + if isinstance(self.tables, dict): + all_tables = dict(self.tables) + # Update with new tables (overwrites if exists, adds if new) + all_tables.update(tables) + else: + # New connection - only use the new tables, but include file-based dependencies + all_tables = dict(tables) + # Analyze SQL expressions to find table dependencies + for sql_expr in tables.values(): + if not isinstance(sql_expr, str): + continue + try: + parsed = sqlglot.parse_one(sql_expr, dialect='duckdb') + except Exception: + continue # If parsing fails, continue without dependencies + # Find all table references in the SQL + # Add file-based tables that are referenced but not already included + for table_obj in parsed.find_all(sqlglot.exp.Table): + table = table_obj.name + if table in self._file_based_tables and table not in all_tables: + all_tables[table] = self._file_based_tables[table] + source_params['tables'] = all_tables if params: source_params['table_params'] = params @@ -382,9 +405,13 @@ def create_sql_expr_source( return source for table, sql_expr in tables.copy().items(): + # Skip file paths - they're already handled by __init__ + if self._is_file_path(sql_expr): + continue + equivalent_sql_exprs = ( - self.sql_expr.format(table=f'"{table_name}"'), - self.sql_expr.format(table=table_name), + self.sql_expr.format(table=f'"{table}"'), + self.sql_expr.format(table=table), ) if table in self.tables: # do not need to re-materialize existing @@ -416,9 +443,11 @@ def create_sql_expr_source( finally: cursor.close() - # keep references of the original file-based tables so views can be recreated - source.tables.update(**{table: self._file_based_tables[table] for table in self._file_based_tables if table not in tables}) - source._file_based_tables.update(self._file_based_tables) + # Preserve file-based metadata for tables that weren't overwritten + source._file_based_tables = { + k: v for k, v in self._file_based_tables.items() + if k not in tables + } return source def execute(self, sql_query: str, params: list | dict | None = None, *args, **kwargs): diff --git a/lumen/sources/rest_duckdb.py b/lumen/sources/rest_duckdb.py new file mode 100644 index 000000000..4542f6ad9 --- /dev/null +++ b/lumen/sources/rest_duckdb.py @@ -0,0 +1,190 @@ +""" +RESTDuckDBSource - DuckDB source with URL parameterization support. + +Enables dynamic URL query parameter updates for REST API endpoints, +making it compatible with LLM-based agents like SQLAgent. +""" +from __future__ import annotations + +from typing import Any, ClassVar +from urllib.parse import urlencode, urlparse, urlunparse + +import param +import sqlglot + +from duckdb import InvalidInputException + +from lumen.sources.base import cached + +from .duckdb import DuckDBSource + + +class RESTDuckDBSource(DuckDBSource): + """ + DuckDBSource subclass that supports parameterized REST API URLs. + + This source allows defining URL templates with dynamic query parameters + that can be updated at runtime, enabling LLM agents to modify API calls + on the fly. + """ + + cache_httpfs = param.Boolean( + default=True, + doc=""" + Whether to cache HTTP responses using DuckDB's cache_httpfs extension. + When True, repeated requests to the same URL return cached results. + """, + ) + + data_format = param.String( + default="json", + doc=""" + Default data format for REST tables if not specified in url_params. + Used to determine the appropriate DuckDB read function. + """, + ) + + source_type: ClassVar[str] = 'rest_duckdb' + + tables = param.Dict( + default={}, + doc=""" + REST table configurations are specified as dictionaries with: + - 'url': Base URL of the REST endpoint + - 'url_params': Dict of query parameters (including 'format' if the API supports it) + - 'required_params': Optional list of parameter names that must be provided + - 'read_fn': Optional override for the DuckDB read function ('json', 'csv', 'parquet'). + If not specified, auto-detects from url_params['format'] or URL extension. + - 'read_options': Optional dict of DuckDB read_* function options + Alternatively, a table can be a SQL expression string as in the base DuckDBSource. + """ + ) + + # Map format to DuckDB read function (using auto variants where available) + _format_to_reader: ClassVar[dict[str, str]] = { + 'json': 'read_json_auto', + 'csv': 'read_csv_auto', + 'parquet': 'read_parquet', + 'ndjson': 'read_ndjson_auto', + } + + _cached_rest_tables = param.Dict(default={}, doc="Internal dict for REST tables to their last used table_params.") + + def _is_rest_table(self, table: str) -> bool: + if table not in self.tables: + raise ValueError(f"Table '{table}' not found in source tables.") + return isinstance(self.tables[table], dict) and 'url' in self.tables[table] + + def _ensure_sql_expr_materialized(self, sql_expr: str, url_params: dict | None = None) -> None: + if not isinstance(sql_expr, str): + # Not a SQL expression, skip + return + + table_objs = sqlglot.parse_one(sql_expr).find_all(sqlglot.exp.Table) + tables = {table_obj.name for table_obj in table_objs if self._is_rest_table(table_obj.name)} + for table in tables: + if not self._is_rest_table(table): + return + table_params = self.tables[table] + if self._cached_rest_tables.get(table) != table_params: + url_params = {**table_params.get("url_params", {}), **(url_params or {})} + df = self.get(table, url_params=url_params) + self._connection.from_df(df).to_view(table) + self._cached_rest_tables[table] = table_params + + def get_sql_expr(self, table: str) -> str: + if self._is_rest_table(table): + table_params = self.tables[table] + # Use table-specific read_fn, or fall back to format-based lookup + read_fn = table_params.get('read_fn') + if read_fn: + # Allow 'json' or 'read_json_auto' style + read_fn = self._format_to_reader.get(read_fn, read_fn) + else: + data_format = table_params.get('url_params', {}).get('format', self.data_format) + read_fn = self._format_to_reader.get(data_format, self._format_to_reader['json']) + + # Handle read_options + read_options = table_params.get('read_options', {}) + if read_options: + options_str = ', '.join(f"{k}={v!r}" for k, v in read_options.items()) + return f"SELECT * FROM {read_fn}(?, {options_str})" + return f"SELECT * FROM {read_fn}(?)" + return super().get_sql_expr(table) + + @cached + def get(self, table: str, url_params: dict[str, Any] | None = None, **query): + if not self._is_rest_table(table): + return super().get(table, **query) + + table_params = self.tables[table].copy() + url_params = {**table_params.get("url_params", {}), **(url_params or {})} + required_params = table_params.get("required_params", []) + missing_params = [p for p in required_params if p not in url_params or url_params[p] is None] + if missing_params: + raise ValueError( + f"Missing required parameters for table '{table}': {missing_params}" + ) + + last_exc = None + url = self.render_table_url(table, url_params=url_params) + data_format = url_params.get("format", self.data_format) + for try_data_format in (data_format, 'csv'): + with self.param.update(table_params={table: [url]}, data_format=try_data_format): + try: + return super().get(table, **query) + except InvalidInputException as exc: + last_exc = exc + continue + + if last_exc is not None: + raise last_exc + + def render_table_url(self, table: str, url_params: dict[str, Any] | None = None) -> str: + """ + Get the current full URL for a REST table. + + Parameters + ---------- + table : str + Name of the REST table + url_params : dict[str, Any] | None + Optional URL parameters to override or add to the table's url params + + Returns + ------- + str + Full URL with current query parameters + """ + if not self._is_rest_table(table): + raise ValueError(f"Table '{table}' is not a REST table.") + + table_params = self.tables[table] + url = table_params["url"] + if url_params is None: + url_params = table_params["url_params"] + parsed = urlparse(url) + return urlunparse(( + parsed.scheme, + parsed.netloc, + parsed.path, + parsed.params, + urlencode(url_params), + parsed.fragment, + )) + + def execute(self, sql_query: str, params: list | dict | None = None, url_params: dict[str, Any] | None = None, *args, **kwargs): + # First ensure all REST tables in the query are materialized + self._ensure_sql_expr_materialized(sql_query, url_params=url_params) + return super().execute(sql_query, *args, params=params, **kwargs) + + def to_spec(self) -> dict[str, Any]: + spec = super().to_spec() + spec.pop("_cached_rest_tables", None) + return spec + + def create_sql_expr_source(self, tables: dict, materialize: bool = True, params: dict | None = None, **kwargs) -> RESTDuckDBSource: + for sql_expr in tables.values(): + self._ensure_sql_expr_materialized(sql_expr) + source = super().create_sql_expr_source(tables, materialize, params, **kwargs) + return source diff --git a/lumen/tests/sources/test_duckdb.py b/lumen/tests/sources/test_duckdb.py index edd2d441c..3062f0048 100644 --- a/lumen/tests/sources/test_duckdb.py +++ b/lumen/tests/sources/test_duckdb.py @@ -719,21 +719,42 @@ def test_detour_roundtrip(sample_csv_files): preserves the original SQL file-based tables so that it can be re-serialized without error. """ - source = DuckDBSource(tables=sample_csv_files) - df = source.get("customers") - new_source = source.create_sql_expr_source( - tables={"limited_customers": 'SELECT * FROM customers LIMIT 1'} - ) - limited_df = new_source.get("limited_customers") - assert len(limited_df) == 1 - assert limited_df.iloc[[0]].equals(df.iloc[[0]]) + files = sample_csv_files + original_cwd = os.getcwd() + + try: + os.chdir(files['dir']) + + # Create source with file-based tables + source = DuckDBSource( + uri=':memory:', + tables={ + 'customers': 'customers.csv', + 'orders': 'orders.csv' + } + ) + df = source.get("customers") + + # Create a derived source with a new SQL expression + new_source = source.create_sql_expr_source( + tables={"limited_customers": 'SELECT * FROM customers LIMIT 1'} + ) + limited_df = new_source.get("limited_customers") + assert len(limited_df) == 1 + assert limited_df.iloc[[0]].equals(df.iloc[[0]]) - read_source = source.from_spec(new_source.to_spec()) - read_df = read_source.get("limited_customers") - assert len(read_df) == 1 - assert read_df.iloc[[0]].equals(df.iloc[[0]]) - assert read_source.tables["limited_customers"] == 'SELECT * FROM customers LIMIT 1' - assert "customers" in read_source.tables + # Serialize and deserialize + spec = new_source.to_spec() + + read_source = DuckDBSource.from_spec(spec) + read_df = read_source.get("limited_customers") + assert len(read_df) == 1 + assert read_df.iloc[[0]].equals(df.iloc[[0]]) + assert read_source.tables["limited_customers"] == 'SELECT * FROM customers LIMIT 1' + assert "customers" in read_source.tables + assert "orders" in read_source.tables + finally: + os.chdir(original_cwd) def test_table_params_basic(sample_csv_files): @@ -960,3 +981,240 @@ def test_table_params_serialization(sample_csv_files): assert restored_result.iloc[0]['id'] == 2 finally: os.chdir(original_cwd) + + +def test_create_sql_expr_source_preserves_all_existing_tables(sample_csv_files): + """Test that create_sql_expr_source preserves ALL existing tables (upsert behavior).""" + files = sample_csv_files + original_cwd = os.getcwd() + + try: + os.chdir(files['dir']) + + # Create initial source with multiple tables + source = DuckDBSource( + uri=':memory:', + tables={ + 'customers': 'customers.csv', + 'orders': 'orders.csv', + 'existing_view': 'SELECT * FROM customers WHERE id > 1' + } + ) + + # Verify initial state + assert set(source.get_tables()) == {'customers', 'orders', 'existing_view'} + + # Create new source with additional tables - should preserve ALL existing ones + new_tables = { + 'new_table': 'SELECT * FROM orders WHERE total > 200' + } + + new_source = source.create_sql_expr_source(new_tables) + + # ALL tables should be present: original + new + expected_tables = {'customers', 'orders', 'existing_view', 'new_table'} + assert set(new_source.get_tables()) == expected_tables + + # Verify all tables are accessible and work correctly + assert len(new_source.get('customers')) == 3 + assert len(new_source.get('orders')) == 3 + assert len(new_source.get('existing_view')) == 2 # id > 1 + assert len(new_source.get('new_table')) == 2 # total > 200 + + finally: + os.chdir(original_cwd) + + +def test_create_sql_expr_source_upserts_existing_tables(sample_csv_files): + """Test that create_sql_expr_source overwrites tables with same name (upsert behavior).""" + files = sample_csv_files + original_cwd = os.getcwd() + + try: + os.chdir(files['dir']) + + # Create initial source + source = DuckDBSource( + uri=':memory:', + tables={ + 'customers': 'customers.csv', + 'orders': 'orders.csv', + 'filtered_customers': 'SELECT * FROM customers WHERE id = 1' # Original: just Alice + } + ) + + # Verify initial state + initial_result = source.get('filtered_customers') + assert len(initial_result) == 1 + assert initial_result.iloc[0]['name'] == 'Alice' + + # Create new source that OVERWRITES filtered_customers but keeps others + new_tables = { + 'filtered_customers': 'SELECT * FROM customers WHERE id > 1', # New: Bob and Charlie + 'new_table': 'SELECT * FROM orders WHERE total > 200' + } + + new_source = source.create_sql_expr_source(new_tables) + + # Should have all tables + expected_tables = {'customers', 'orders', 'filtered_customers', 'new_table'} + assert set(new_source.get_tables()) == expected_tables + + # filtered_customers should have NEW definition (id > 1, not id = 1) + updated_result = new_source.get('filtered_customers') + assert len(updated_result) == 2 + assert set(updated_result['name']) == {'Bob', 'Charlie'} + + # Original tables should still work + assert len(new_source.get('customers')) == 3 + assert len(new_source.get('orders')) == 3 + assert len(new_source.get('new_table')) == 2 + + finally: + os.chdir(original_cwd) + + +def test_create_sql_expr_source_new_connection_only_new_tables(sample_csv_files): + """Test that create_sql_expr_source with new connection only includes new tables.""" + files = sample_csv_files + original_cwd = os.getcwd() + + try: + os.chdir(files['dir']) + + # Create initial source with multiple tables + source = DuckDBSource( + uri=':memory:', + tables={ + 'customers': 'customers.csv', + 'orders': 'orders.csv', + 'existing_view': 'SELECT * FROM customers WHERE id > 1' + } + ) + + # Verify initial state + assert set(source.get_tables()) == {'customers', 'orders', 'existing_view'} + + # Create new source with a DIFFERENT URI - should NOT preserve old tables + new_tables = { + 'products': files['customers'] # Reusing customers.csv as "products" + } + + new_source = source.create_sql_expr_source(new_tables, uri=':memory:') + + # Should ONLY have the new table, not the old ones + assert set(new_source.get_tables()) == {'products'} + + # Old tables should NOT be accessible + assert 'customers' not in new_source.get_tables() + assert 'orders' not in new_source.get_tables() + assert 'existing_view' not in new_source.get_tables() + + finally: + os.chdir(original_cwd) + + +def test_create_sql_expr_source_new_connection_includes_file_dependencies(sample_csv_files): + """Test that new connection includes file-based tables referenced in SQL.""" + files = sample_csv_files + original_cwd = os.getcwd() + + try: + os.chdir(files['dir']) + + # Create initial source with file-based tables + source = DuckDBSource( + uri=':memory:', + tables={ + 'customers': 'customers.csv', + 'orders': 'orders.csv', + } + ) + + # Create new source with new connection that REFERENCES file-based tables + new_tables = { + 'summary': 'SELECT c.name, COUNT(o.id) as order_count FROM customers c LEFT JOIN orders o ON c.id = o.customer_id GROUP BY c.name' + } + + new_source = source.create_sql_expr_source(new_tables, uri=':memory:') + + # Should have the new table AND the file-based dependencies + expected_tables = {'summary', 'customers', 'orders'} + assert set(new_source.get_tables()) == expected_tables + + # The summary query should actually work (dependencies are present) + result = new_source.get('summary') + assert len(result) == 3 # 3 customers + assert 'order_count' in result.columns + + # File-based tables should be accessible + assert len(new_source.get('customers')) == 3 + assert len(new_source.get('orders')) == 3 + + finally: + os.chdir(original_cwd) + + +def test_create_sql_expr_source_with_list_tables(): + """Test that create_sql_expr_source works when self.tables is a list.""" + # Create an in-memory source with actual data + df = pd.DataFrame({ + 'A': [0, 1, 2, 3, 4], + 'B': [0, 0, 1, 1, 1], + 'C': ['foo1', 'foo2', 'foo3', 'foo4', 'foo5'] + }) + + # Use from_df which creates dict-based tables, then manually convert to list + source = DuckDBSource.from_df({'test_table': df}) + # Simulate a list-based source (though unusual in practice) + source.tables = ['test_table'] # Override with list + + # Verify it's a list + assert isinstance(source.tables, list) + + # Create new source with SQL expressions + new_tables = { + 'filtered': 'SELECT * FROM test_table WHERE A > 2' + } + + new_source = source.create_sql_expr_source(new_tables) + + # Should only have the new table (list-based tables don't get preserved) + assert 'filtered' in new_source.get_tables() + assert isinstance(new_source.tables, dict) + assert 'filtered' in new_source.tables + + # The new table should work + result = new_source.get('filtered') + assert len(result) == 2 # A values 3 and 4 + assert all(result['A'] > 2) + + +def test_create_sql_expr_source_reuse_connection_with_list_tables(): + """Test that reusing connection with list tables just uses new tables.""" + # Create an in-memory source with actual data + df = pd.DataFrame({ + 'A': [0, 1, 2, 3, 4], + 'B': [0, 0, 1, 1, 1], + 'C': ['foo1', 'foo2', 'foo3', 'foo4', 'foo5'] + }) + + source = DuckDBSource.from_df({'test_table': df}) + # Simulate a list-based source + source.tables = ['test_table'] # Override with list + + # Create new source reusing connection + new_tables = { + 'filtered': 'SELECT * FROM test_table WHERE A > 2' + } + + # No uri or initializers provided - reusing connection + new_source = source.create_sql_expr_source(new_tables) + + # Should only have the new tables (since original was a list, not dict) + assert set(new_source.get_tables()) == {'filtered'} + + # But the connection is reused, so we can still query the original table + # via the connection even if it's not in new_source.tables + result = new_source.execute('SELECT * FROM test_table') + assert len(result) == 5 # Original table still exists in the connection diff --git a/lumen/tests/sources/test_rest_duckdb.py b/lumen/tests/sources/test_rest_duckdb.py new file mode 100644 index 000000000..54b3ab0ba --- /dev/null +++ b/lumen/tests/sources/test_rest_duckdb.py @@ -0,0 +1,511 @@ +"""Tests for RESTDuckDBSource.""" +import pandas as pd +import pytest + +try: + from lumen.sources.rest_duckdb import RESTDuckDBSource + pytestmark = pytest.mark.xdist_group("duckdb") +except ImportError: + pytestmark = pytest.mark.skip(reason="DuckDB is not installed") + + +# Table configurations as constants +DAILY_TABLE_CONFIG = { + 'url': 'https://mesonet.agron.iastate.edu/cgi-bin/request/daily.py', + 'url_params': { + 'stations': 'ABR', + 'sts': '2025-12-08', + 'ets': '2025-12-09', + 'network': 'SD_ASOS', + 'format': 'csv' + }, +} + +RAOB_TABLE_CONFIG = { + 'url': 'https://mesonet.agron.iastate.edu/cgi-bin/request/raob.py', + 'url_params': { + 'station': 'KABR', + 'sts': '2025-12-08T15:49', + 'ets': '2025-12-09T15:49', + 'format': 'csv' + }, +} + +PENGUINS_CSV_URL = 'https://raw.githubusercontent.com/mwaskom/seaborn-data/master/penguins.csv' + + +@pytest.fixture(scope="session") +def single_table_source(): + """Fixture providing a RESTDuckDBSource with one REST table.""" + config = { + 'uri': ':memory:', + 'tables': { + 'daily': DAILY_TABLE_CONFIG, + } + } + source = RESTDuckDBSource(**config) + # Pre-materialize to avoid repeated API calls + daily_df = source.get('daily') + source._connection.from_df(daily_df).to_view('daily') + source._cached_rest_tables["daily"] = DAILY_TABLE_CONFIG + return source + + +@pytest.fixture(scope="session") +def multi_table_source(): + """Fixture providing a RESTDuckDBSource with two REST tables.""" + config = { + 'uri': ':memory:', + 'tables': { + 'daily': DAILY_TABLE_CONFIG, + 'raob': RAOB_TABLE_CONFIG, + } + } + source = RESTDuckDBSource(**config) + # Pre-materialize both tables + daily_df = source.get('daily') + raob_df = source.get('raob') + source._connection.from_df(daily_df).to_view('daily') + source._connection.from_df(raob_df).to_view('raob') + source._cached_rest_tables["daily"] = DAILY_TABLE_CONFIG + source._cached_rest_tables["raob"] = RAOB_TABLE_CONFIG + return source + + +@pytest.fixture(scope="session") +def mixed_table_source(): + """Fixture providing a RESTDuckDBSource with REST table and CSV file.""" + config = { + 'uri': ':memory:', + 'tables': { + 'daily': DAILY_TABLE_CONFIG, + 'penguins': PENGUINS_CSV_URL, + } + } + source = RESTDuckDBSource(**config) + # Pre-materialize REST table + daily_df = source.get('daily') + source._connection.from_df(daily_df).to_view('daily') + source._cached_rest_tables["daily"] = DAILY_TABLE_CONFIG + # CSV table doesn't need pre-materialization + return source + + +class TestRESTDuckDBSourceBasics: + """Test basic functionality and initialization.""" + + def test_source_type(self): + """Test that source_type is correctly set.""" + assert RESTDuckDBSource.source_type == 'rest_duckdb' + + def test_resolve_module_type(self): + """Test that the source can be resolved by module path.""" + assert RESTDuckDBSource._get_type('lumen.sources.rest_duckdb.RESTDuckDBSource') is RESTDuckDBSource + + def test_initialization(self, single_table_source): + """Test that RESTDuckDBSource initializes correctly.""" + assert single_table_source.uri == ':memory:' + assert 'daily' in single_table_source.tables + assert isinstance(single_table_source.tables['daily'], dict) + assert 'url' in single_table_source.tables['daily'] + + +class TestRESTTableOperations: + """Test REST-specific table operations.""" + + def test_render_table_url(self, single_table_source): + """Test that render_table_url constructs correct URLs.""" + url = single_table_source.render_table_url('daily') + assert 'https://mesonet.agron.iastate.edu/cgi-bin/request/daily.py' in url + assert 'stations=ABR' in url + assert 'sts=2025-12-08' in url + assert 'format=csv' in url + + def test_get_table(self, single_table_source): + """Test that get() retrieves data correctly.""" + df = single_table_source.get('daily') + + assert isinstance(df, pd.DataFrame) + assert not df.empty + assert 'station' in df.columns + assert 'day' in df.columns + assert 'max_temp_f' in df.columns + assert len(df) == 2 + assert all(df['station'] == 'ABR') + + def test_get_multiple_tables(self, multi_table_source): + """Test that both tables can be retrieved.""" + daily_df = multi_table_source.get('daily') + raob_df = multi_table_source.get('raob') + + assert isinstance(daily_df, pd.DataFrame) + assert isinstance(raob_df, pd.DataFrame) + assert not daily_df.empty + + def test_invalid_table_name(self, single_table_source): + """Test that accessing non-existent table raises appropriate error.""" + with pytest.raises(Exception): + single_table_source.get('nonexistent_table') + + +class TestSQLExecution: + """Test SQL query execution.""" + + def test_execute_simple_query(self, single_table_source): + """Test basic SQL execution.""" + result = single_table_source.execute("SELECT * FROM daily LIMIT 5") + + assert isinstance(result, pd.DataFrame) + assert len(result) <= 5 + assert 'station' in result.columns + + def test_execute_with_filter(self, single_table_source): + """Test SQL with WHERE clause.""" + result = single_table_source.execute("SELECT * FROM daily WHERE max_temp_f > 20") + + assert isinstance(result, pd.DataFrame) + if not result.empty: + assert all(result['max_temp_f'] > 20) + + def test_execute_aggregate(self, single_table_source): + """Test SQL aggregation functions.""" + result = single_table_source.execute("SELECT COUNT(*) as count FROM daily") + + assert isinstance(result, pd.DataFrame) + assert 'count' in result.columns + assert result['count'].iloc[0] > 0 + + def test_execute_materializes_rest_tables(self, single_table_source): + """Test that execute() automatically materializes REST tables.""" + result = single_table_source.execute("SELECT COUNT(*) as cnt FROM daily") + + assert isinstance(result, pd.DataFrame) + assert 'daily' in single_table_source._cached_rest_tables + + def test_invalid_sql_query(self, single_table_source): + """Test that invalid SQL raises appropriate error.""" + with pytest.raises(Exception): + single_table_source.execute("SELECT * FROM nonexistent_table") + + +class TestSQLExpressionSource: + """Test create_sql_expr_source functionality.""" + + def test_create_simple_expression(self, single_table_source): + """Test creating a source with a simple SQL expression.""" + new_source = single_table_source.create_sql_expr_source({ + 'daily_1': "SELECT * FROM daily LIMIT 1" + }) + + assert 'daily_1' in new_source.tables + df = new_source.get('daily_1') + assert isinstance(df, pd.DataFrame) + assert len(df) == 1 + + def test_create_multiple_expressions(self, single_table_source): + """Test creating multiple SQL expressions at once.""" + new_source = single_table_source.create_sql_expr_source({ + 'daily_1': "SELECT * FROM daily LIMIT 1", + 'daily_high_temp': "SELECT * FROM daily WHERE max_temp_f > 30" + }) + + assert 'daily_1' in new_source.tables + assert 'daily_high_temp' in new_source.tables + + df1 = new_source.get('daily_1') + df_high = new_source.get('daily_high_temp') + + assert len(df1) == 1 + assert isinstance(df_high, pd.DataFrame) + + def test_preserves_original_tables(self, single_table_source): + """Test that creating SQL expr source preserves original tables.""" + new_source = single_table_source.create_sql_expr_source({ + 'daily_1': "SELECT * FROM daily LIMIT 1" + }) + + # Original table should still be accessible + daily_df = new_source.get('daily') + assert isinstance(daily_df, pd.DataFrame) + assert len(daily_df) > 1 + + def test_rest_table_dependency_materialization(self, single_table_source): + """Test that REST tables in SQL expressions are materialized.""" + new_source = single_table_source.create_sql_expr_source({ + 'daily_filtered': "SELECT * FROM daily WHERE max_temp_f > 20" + }) + + # REST table should be accessible and materialized + assert 'daily' in new_source.tables + daily_df = new_source.get('daily') + filtered_df = new_source.get('daily_filtered') + + assert len(filtered_df) <= len(daily_df) + if not filtered_df.empty: + assert all(filtered_df['max_temp_f'] > 20) + + def test_upsert_behavior(self, single_table_source): + """Test that new tables with same name override existing ones.""" + source1 = single_table_source.create_sql_expr_source({ + 'summary': "SELECT COUNT(*) as total_days FROM daily" + }) + result1 = source1.get('summary') + assert 'total_days' in result1.columns + + source2 = source1.create_sql_expr_source({ + 'summary': "SELECT AVG(max_temp_f) as avg_temp FROM daily" + }) + result2 = source2.get('summary') + assert 'avg_temp' in result2.columns + assert 'total_days' not in result2.columns + + def test_multiple_rest_dependencies(self, multi_table_source): + """Test SQL expression that references multiple REST tables.""" + new_source = multi_table_source.create_sql_expr_source({ + 'combined': """ + SELECT station, day, max_temp_f FROM daily + UNION ALL + SELECT station, validUTC as day, tmpc as max_temp_f FROM raob + LIMIT 10 + """ + }) + + # Both REST tables should be materialized + assert 'daily' in new_source.tables + assert 'raob' in new_source.tables + assert 'combined' in new_source.tables + + result = new_source.get('combined') + assert isinstance(result, pd.DataFrame) + assert len(result) <= 10 + + def test_preserves_rest_configs(self, single_table_source): + """Test that REST table configs are preserved in derived sources.""" + new_source = single_table_source.create_sql_expr_source({ + 'daily_subset': "SELECT * FROM daily LIMIT 5" + }) + + # REST config should be preserved + assert isinstance(new_source.tables['daily'], dict) + assert 'url' in new_source.tables['daily'] + url = new_source.render_table_url('daily') + assert 'https://mesonet.agron.iastate.edu' in url + + +class TestSerialization: + """Test source serialization.""" + + def test_to_spec_basic(self, single_table_source): + """Test that to_spec() returns correct specification.""" + spec = single_table_source.to_spec() + + assert isinstance(spec, dict) + assert spec['uri'] == ':memory:' + assert 'tables' in spec + assert spec['type'] == 'rest_duckdb' + assert '_cached_rest_tables' not in spec + + def test_to_spec_with_sql_expressions(self, single_table_source): + """Test to_spec() on derived source with SQL expressions.""" + new_source = single_table_source.create_sql_expr_source({ + 'daily_1': "SELECT * FROM daily LIMIT 1" + }) + + spec = new_source.to_spec() + assert 'daily_1' in spec['tables'] + assert spec['tables']['daily_1'] == "SELECT * FROM daily LIMIT 1" + + +class TestDataValidation: + """Test data type and content validation.""" + + def test_column_presence(self, single_table_source): + """Test that expected columns are present.""" + df = single_table_source.get('daily') + + expected_columns = ['station', 'day', 'max_temp_f', 'min_temp_f', + 'max_dewpoint_f', 'min_dewpoint_f', 'precip_in'] + for col in expected_columns: + assert col in df.columns + + def test_data_types(self, single_table_source): + """Test that data types are correctly inferred.""" + df = single_table_source.get('daily') + + assert pd.api.types.is_numeric_dtype(df['max_temp_f']) + assert pd.api.types.is_numeric_dtype(df['min_temp_f']) + assert pd.api.types.is_numeric_dtype(df['precip_in']) + + +class TestRequiredParams: + """Test required_params validation.""" + + def test_required_params_missing_raises_error(self): + """Test that missing required params raises ValueError.""" + config = { + 'uri': ':memory:', + 'tables': { + 'daily': { + 'url': 'https://mesonet.agron.iastate.edu/cgi-bin/request/daily.py', + 'url_params': { + 'stations': 'ABR', + 'network': 'SD_ASOS', + 'format': 'csv' + }, + 'required_params': ['stations', 'sts', 'ets'], + }, + } + } + source = RESTDuckDBSource(**config) + + with pytest.raises(ValueError, match="Missing required parameters.*sts.*ets"): + source.get('daily') + + def test_required_params_provided_via_url_params_arg(self): + """Test that required params can be provided via url_params argument.""" + config = { + 'uri': ':memory:', + 'tables': { + 'daily': { + 'url': 'https://mesonet.agron.iastate.edu/cgi-bin/request/daily.py', + 'url_params': { + 'network': 'SD_ASOS', + 'format': 'csv' + }, + 'required_params': ['stations', 'sts', 'ets'], + }, + } + } + source = RESTDuckDBSource(**config) + df = source.get('daily', url_params={ + 'stations': 'ABR', + 'sts': '2025-12-08', + 'ets': '2025-12-09', + }) + + assert isinstance(df, pd.DataFrame) + assert not df.empty + + +class TestReadFnAndReadOptions: + """Test read_fn and read_options configuration.""" + + def test_read_fn_and_read_options_in_sql_expr(self): + """Test that read_fn and read_options are included in the SQL expression.""" + config = { + 'uri': ':memory:', + 'tables': { + 'data': { + 'url': 'https://example.com/data.csv', + 'url_params': {}, + 'read_fn': 'csv', + 'read_options': { + 'header': True, + 'delim': ',', + }, + }, + } + } + source = RESTDuckDBSource(**config) + sql_expr = source.get_sql_expr('data') + + assert 'read_csv_auto' in sql_expr + assert 'header=True' in sql_expr + assert "delim=','" in sql_expr + + def test_read_fn_falls_back_to_url_params_format(self): + """Test that read_fn falls back to url_params['format'].""" + config = { + 'uri': ':memory:', + 'tables': { + 'data': { + 'url': 'https://example.com/data', + 'url_params': {'format': 'csv'}, + }, + } + } + source = RESTDuckDBSource(**config) + sql_expr = source.get_sql_expr('data') + + assert 'read_csv_auto' in sql_expr + + +class TestMixedTableTypes: + """Test mixing REST tables with regular CSV tables.""" + + def test_mixed_source_has_both_table_types(self, mixed_table_source): + """Test that mixed source contains both REST and CSV tables.""" + assert 'daily' in mixed_table_source.tables + assert 'penguins' in mixed_table_source.tables + + # daily is REST table (dict config) + assert isinstance(mixed_table_source.tables['daily'], dict) + assert 'url' in mixed_table_source.tables['daily'] + + # penguins is CSV table (string URL) + assert isinstance(mixed_table_source.tables['penguins'], str) + + def test_get_csv_table(self, mixed_table_source): + """Test retrieving CSV table from mixed source.""" + df = mixed_table_source.get('penguins') + + assert isinstance(df, pd.DataFrame) + assert not df.empty + assert 'species' in df.columns + assert 'island' in df.columns + assert 'bill_length_mm' in df.columns + + def test_get_rest_table_from_mixed(self, mixed_table_source): + """Test retrieving REST table from mixed source.""" + df = mixed_table_source.get('daily') + + assert isinstance(df, pd.DataFrame) + assert not df.empty + assert 'station' in df.columns + + def test_sql_join_rest_and_csv(self, mixed_table_source): + """Test SQL query joining REST and CSV tables.""" + result = mixed_table_source.execute(""" + SELECT d.station, p.species, COUNT(*) as count + FROM daily d + CROSS JOIN penguins p + WHERE p.species = 'Adelie' + GROUP BY d.station, p.species + LIMIT 5 + """) + + assert isinstance(result, pd.DataFrame) + assert not result.empty + assert 'station' in result.columns + assert 'species' in result.columns + assert all(result['species'] == 'Adelie') + + def test_create_sql_expr_with_mixed_tables(self, mixed_table_source): + """Test creating SQL expressions that reference both table types.""" + new_source = mixed_table_source.create_sql_expr_source({ + 'rest_summary': "SELECT station, AVG(max_temp_f) as avg_temp FROM daily GROUP BY station", + 'csv_summary': "SELECT species, COUNT(*) as count FROM penguins GROUP BY species", + 'combined': """ + SELECT 'weather' as source_type, station as name FROM daily + UNION ALL + SELECT 'penguin' as source_type, species as name FROM penguins + LIMIT 10 + """ + }) + + # All tables should exist + assert 'rest_summary' in new_source.tables + assert 'csv_summary' in new_source.tables + assert 'combined' in new_source.tables + + # Verify they work + rest_df = new_source.get('rest_summary') + csv_df = new_source.get('csv_summary') + combined_df = new_source.get('combined') + + assert isinstance(rest_df, pd.DataFrame) + assert isinstance(csv_df, pd.DataFrame) + assert isinstance(combined_df, pd.DataFrame) + assert 'avg_temp' in rest_df.columns + assert 'species' in csv_df.columns + assert 'source_type' in combined_df.columns