From a099d98e352eb6d8c379e7a9f8dd26a2f3d3ab84 Mon Sep 17 00:00:00 2001 From: Gordon Murray Date: Wed, 15 Apr 2026 19:32:32 +0100 Subject: [PATCH 1/2] refactor(backend): per-request DuckDB connections Replaces the module-level DuckDBManager with a per-request context manager. The previous design shared one connection across all requests and closed/recreated it inside _apply_s3_config, which was not concurrency-safe: one request could close the connection another was actively using. - _duckdb_connection(config) yields a fresh in-memory conn, loads extensions, applies S3+Iceberg config, and closes on exit - _apply_s3_config, _attach_iceberg_catalog, _validate_iceberg_table, _convert_to_iceberg_query become module-level functions taking conn - Route handlers delegate to run_connection_test / run_query impls - Lifespan pre-installs extensions once at startup so first request does not pay the download cost (DuckDB caches binaries on disk) - Drop PRAGMA enable_profiling / profiling_output: the output file was never read and writing a shared filename from concurrent connections would reintroduce the race we are fixing Closes #9 --- backend/main.py | 499 ++++++++++++++++++++++-------------------------- 1 file changed, 226 insertions(+), 273 deletions(-) diff --git a/backend/main.py b/backend/main.py index f6e5170..0b084da 100644 --- a/backend/main.py +++ b/backend/main.py @@ -8,7 +8,7 @@ import os import re import time -from contextlib import asynccontextmanager +from contextlib import asynccontextmanager, contextmanager from typing import Any, List, Literal, Optional import sqlglot @@ -239,242 +239,205 @@ class QueryResponse(BaseModel): truncated: bool = False -class DuckDBManager: - """Manages DuckDB connections and query execution.""" +# --- DuckDB session management --------------------------------------------- +# A fresh in-memory connection is opened per request, configured with the +# caller's S3/Iceberg settings, and closed on exit. The previous design kept +# a single module-level connection that was torn down and rebuilt on every +# query, which was not concurrency-safe: one request's _apply_s3_config call +# could swap the connection out from under another request mid-query. +# +# Extension binaries are cached on disk by DuckDB after the first download, +# so repeated INSTALL/LOAD is a cheap no-op (ms-level) for subsequent +# connections. - def __init__(self): - self.connection = None - self._setup_duckdb() - def _setup_duckdb(self): - """Initialize DuckDB with required extensions.""" - try: - logger.info("Setting up DuckDB...") - self.connection = duckdb.connect(":memory:") - - # Install extensions - self.connection.execute("INSTALL httpfs") - self.connection.execute("LOAD httpfs") - - # Iceberg extension is mandatory for correct data reads - try: - self.connection.execute("INSTALL iceberg") - self.connection.execute("LOAD iceberg") - # DuckDB 1.4+ requires explicit version handling - self.connection.execute("SET unsafe_enable_version_guessing=true") - logger.info("Iceberg extension loaded") - except Exception as e: - logger.error(f"FATAL: Iceberg extension required but not available: {e}") - raise RuntimeError("Iceberg extension is required for correct data reads") - - # Set default settings - self.connection.execute("SET memory_limit='2GB'") - self.connection.execute("SET threads=4") - - logger.info("DuckDB initialized successfully") +def _apply_s3_config(conn: duckdb.DuckDBPyConnection, config: ConnectionConfig) -> None: + """Apply storage-specific S3 settings to ``conn``. - except Exception as e: - logger.error(f"Failed to setup DuckDB: {e}") - raise + All user-supplied values are sent via DuckDB's ``?`` parameter binding + (the Pydantic validators on :class:`ConnectionConfig` are a second line + of defence, not the only one). + """ + logger.info(f"Applying S3 config: {config.storageType}, endpoint: {config.endpoint}") + + if config.storageType == "minio": + endpoint = config.endpoint + if "localhost" in endpoint: + # Replace localhost with container name for internal access + endpoint = endpoint.replace("localhost", "minio") + endpoint = endpoint.replace("http://", "").replace("https://", "") + logger.info(f"Final MinIO endpoint: {endpoint}") + conn.execute("SET s3_endpoint=?", [endpoint]) + conn.execute("SET s3_url_style='path'") + conn.execute("SET s3_use_ssl=false") + # MinIO requires AWS signature v4 + conn.execute("SET s3_region='us-east-1'") + elif config.storageType == "r2": + endpoint = config.endpoint.replace("https://", "") + conn.execute("SET s3_endpoint=?", [endpoint]) + conn.execute("SET s3_url_style='path'") + conn.execute("SET s3_use_ssl=true") + else: + logger.info(f"Setting S3 region: {config.region}") + conn.execute("SET s3_region=?", [config.region]) + conn.execute("SET s3_use_ssl=true") + + logger.info( + f"Setting S3 credentials - Access Key starts with: {config.accessKey[:8] if config.accessKey else 'EMPTY'}..." + ) + conn.execute("SET s3_access_key_id=?", [config.accessKey]) + conn.execute("SET s3_secret_access_key=?", [config.secretKey]) - def _apply_s3_config(self, config: ConnectionConfig): - """Apply S3 configuration to DuckDB session.""" - try: - logger.info(f"Applying S3 config: {config.storageType}, endpoint: {config.endpoint}") - - # Create a fresh connection to avoid state issues - if self.connection: - self.connection.close() - self.connection = duckdb.connect(":memory:") - - # Reinstall extensions for new connection - self.connection.execute("INSTALL httpfs") - self.connection.execute("LOAD httpfs") - self.connection.execute("INSTALL iceberg") - self.connection.execute("LOAD iceberg") - self.connection.execute("SET unsafe_enable_version_guessing=true") - - # Set default settings - self.connection.execute("SET memory_limit='2GB'") - self.connection.execute("SET threads=4") - - # Apply new settings based on storage type. All user-supplied - # values are sent through DuckDB parameter binding (?). - if config.storageType == "minio": - endpoint = config.endpoint - if "localhost" in endpoint: - # Replace localhost with container name for internal access - endpoint = endpoint.replace("localhost", "minio") - endpoint = endpoint.replace("http://", "").replace("https://", "") - logger.info(f"Final MinIO endpoint: {endpoint}") - self.connection.execute("SET s3_endpoint=?", [endpoint]) - self.connection.execute("SET s3_url_style='path'") - self.connection.execute("SET s3_use_ssl=false") - # MinIO requires AWS signature v4 - self.connection.execute("SET s3_region='us-east-1'") # MinIO default - elif config.storageType == "r2": - endpoint = config.endpoint.replace("https://", "") - self.connection.execute("SET s3_endpoint=?", [endpoint]) - self.connection.execute("SET s3_url_style='path'") - self.connection.execute("SET s3_use_ssl=true") - else: - logger.info(f"Setting S3 region: {config.region}") - self.connection.execute("SET s3_region=?", [config.region]) - self.connection.execute("SET s3_use_ssl=true") + if config.sessionToken: + conn.execute("SET s3_session_token=?", [config.sessionToken]) - # Set credentials (always via parameter binding). - logger.info( - f"Setting S3 credentials - Access Key starts with: {config.accessKey[:8] if config.accessKey else 'EMPTY'}..." - ) - self.connection.execute("SET s3_access_key_id=?", [config.accessKey]) - self.connection.execute("SET s3_secret_access_key=?", [config.secretKey]) + logger.info(f"Applied {config.storageType} configuration") - if config.sessionToken: - self.connection.execute("SET s3_session_token=?", [config.sessionToken]) - logger.info(f"Applied {config.storageType} configuration") +def _attach_iceberg_catalog(conn: duckdb.DuckDBPyConnection, config: ConnectionConfig) -> None: + """Attach an Iceberg REST catalog on ``conn`` if the config requests one.""" + if config.catalogType != "rest": + return - # Attach Iceberg catalog if configured - self._attach_iceberg_catalog(config) + if not config.catalogEndpoint: + raise HTTPException( + status_code=400, + detail="catalogEndpoint required for REST catalog", + ) + if not config.namespace: + raise HTTPException( + status_code=400, + detail="namespace required for REST catalog", + ) - except Exception as e: - logger.error(f"Failed to apply S3 config: {e}") - raise HTTPException(status_code=400, detail=f"Invalid S3 configuration: {e}") + logger.info(f"Attaching Iceberg REST catalog: {config.catalogEndpoint}") - def _attach_iceberg_catalog(self, config: ConnectionConfig): - """Attach Iceberg REST catalog if configured.""" - if config.catalogType != "rest": - return + # CREATE SECRET and ATTACH do not support prepared-statement placeholders + # for their option values, so we interpolate after (a) Pydantic-level + # regex validation on namespace/catalogEndpoint and (b) escaping through + # _sql_string_literal. + token_literal = _sql_string_literal(f"{config.accessKey}:{config.secretKey}") + namespace_literal = _sql_string_literal(config.namespace) + endpoint_literal = _sql_string_literal(config.catalogEndpoint) - if not config.catalogEndpoint: - raise HTTPException( - status_code=400, - detail="catalogEndpoint required for REST catalog", - ) - if not config.namespace: - raise HTTPException( - status_code=400, - detail="namespace required for REST catalog", - ) + conn.execute(f""" + CREATE SECRET iceberg_catalog_secret ( + TYPE iceberg, + TOKEN {token_literal} + ) + """) - logger.info(f"Attaching Iceberg REST catalog: {config.catalogEndpoint}") - - # CREATE SECRET and ATTACH do not support prepared-statement - # placeholders for their option values, so we interpolate after - # (a) Pydantic-level regex validation on namespace/catalogEndpoint and - # (b) escaping through _sql_string_literal, which doubles quotes and - # rejects any control characters. - token_literal = _sql_string_literal(f"{config.accessKey}:{config.secretKey}") - namespace_literal = _sql_string_literal(config.namespace) - endpoint_literal = _sql_string_literal(config.catalogEndpoint) - - self.connection.execute(f""" - CREATE SECRET iceberg_catalog_secret ( - TYPE iceberg, - TOKEN {token_literal} - ) - """) + conn.execute(f""" + ATTACH {namespace_literal} AS iceberg_catalog ( + TYPE iceberg, + SECRET iceberg_catalog_secret, + ENDPOINT {endpoint_literal} + ) + """) - self.connection.execute(f""" - ATTACH {namespace_literal} AS iceberg_catalog ( - TYPE iceberg, - SECRET iceberg_catalog_secret, - ENDPOINT {endpoint_literal} - ) - """) + logger.info("Iceberg catalog attached") - logger.info("Iceberg catalog attached") - def _validate_iceberg_table(self, table_path: str) -> dict: - """ - Validate Iceberg table compatibility (v1/v2 only, no deletes). +@contextmanager +def _duckdb_connection(config: ConnectionConfig): + """Yield a fresh DuckDB in-memory connection configured for ``config``. - Returns dict with validation results or raises HTTPException. - """ + One connection per caller means concurrent requests can't corrupt each + other's session state. The connection is always closed on exit. + """ + conn = duckdb.connect(":memory:") + try: + conn.execute("INSTALL httpfs") + conn.execute("LOAD httpfs") + conn.execute("INSTALL iceberg") + conn.execute("LOAD iceberg") + # DuckDB 1.4+ requires explicit version handling for Iceberg + conn.execute("SET unsafe_enable_version_guessing=true") + conn.execute("SET memory_limit='2GB'") + conn.execute("SET threads=4") try: - # Read metadata to check for delete files - metadata = self.connection.execute( - "SELECT * FROM iceberg_metadata(?)", [table_path] - ).fetchdf() - - # Check for delete files in manifests - has_deletes = any('DELETE' in str(v).upper() for v in metadata['manifest_content'].unique()) - - if has_deletes: - raise HTTPException( - status_code=400, - detail=( - "Table contains row-level deletes which are not supported. " - "This application only supports append-only Iceberg v1/v2 tables. " - "Reading this table may return incorrect data." - ) - ) - - logger.info(f"Table validation passed: {table_path}") - return {"valid": True, "warnings": []} - + _apply_s3_config(conn, config) + _attach_iceberg_catalog(conn, config) except HTTPException: raise except Exception as e: - logger.warning(f"Could not validate table (proceeding with caution): {e}") - return {"valid": True, "warnings": [f"Validation incomplete: {e}"]} + logger.error(f"Failed to apply S3 config: {e}") + raise HTTPException(status_code=400, detail=f"Invalid S3 configuration: {e}") + yield conn + finally: + conn.close() - def _convert_to_iceberg_query(self, sql: str, config: ConnectionConfig) -> str: - """ - Convert read_parquet() calls to iceberg_scan() calls. - Detects: - - read_parquet('s3://bucket/path/**/*.parquet') +def _validate_iceberg_table(conn: duckdb.DuckDBPyConnection, table_path: str) -> dict: + """Validate Iceberg table compatibility (v1/v2 only, no deletes).""" + try: + metadata = conn.execute( + "SELECT * FROM iceberg_metadata(?)", [table_path] + ).fetchdf() - Converts to: - - iceberg_scan('s3://bucket/path') OR - - SELECT * FROM iceberg_catalog.namespace.table - """ - import re + has_deletes = any( + 'DELETE' in str(v).upper() for v in metadata['manifest_content'].unique() + ) - # Pattern: read_parquet('s3://...**/*.parquet') - parquet_pattern = r"read_parquet\(['\"]s3://([^/]+)/([^'\"]+?)/?\*?\*?/?\*?\.parquet['\"]\)" + if has_deletes: + raise HTTPException( + status_code=400, + detail=( + "Table contains row-level deletes which are not supported. " + "This application only supports append-only Iceberg v1/v2 tables. " + "Reading this table may return incorrect data." + ), + ) - def replace_with_iceberg(match): - bucket = match.group(1) - path = match.group(2).rstrip('/*') + logger.info(f"Table validation passed: {table_path}") + return {"valid": True, "warnings": []} - if config.catalogType == "rest": - # Use catalog table reference - # Assumes table name is last path component - table_name = path.split('/')[-1] - return f"iceberg_catalog.{config.namespace}.{table_name}" - else: - # Direct iceberg_scan - iceberg_path = f"s3://{bucket}/{path}" - return f"iceberg_scan('{iceberg_path}')" + except HTTPException: + raise + except Exception as e: + logger.warning(f"Could not validate table (proceeding with caution): {e}") + return {"valid": True, "warnings": [f"Validation incomplete: {e}"]} - converted_sql = re.sub(parquet_pattern, replace_with_iceberg, sql, flags=re.IGNORECASE) - if converted_sql != sql: - logger.info("Converted query from read_parquet to Iceberg:") - logger.info(f" Original: {sql[:100]}...") - logger.info(f" Converted: {converted_sql[:100]}...") +def _convert_to_iceberg_query(sql: str, config: ConnectionConfig) -> str: + """Convert ``read_parquet('s3://...**/*.parquet')`` calls to + ``iceberg_scan()`` (or a catalog table reference when a REST catalog is + configured).""" + parquet_pattern = r"read_parquet\(['\"]s3://([^/]+)/([^'\"]+?)/?\*?\*?/?\*?\.parquet['\"]\)" - return converted_sql + def replace_with_iceberg(match): + bucket = match.group(1) + path = match.group(2).rstrip('/*') - def test_connection(self, config: ConnectionConfig) -> bool: - """Test if the connection configuration works.""" - try: - # tablePath is normalised (trailing / and /metadata stripped) and - # regex-validated by ConnectionConfig, so no further cleanup here. - self._apply_s3_config(config) + if config.catalogType == "rest": + table_name = path.split('/')[-1] + return f"iceberg_catalog.{config.namespace}.{table_name}" + iceberg_path = f"s3://{bucket}/{path}" + return f"iceberg_scan('{iceberg_path}')" + + converted_sql = re.sub(parquet_pattern, replace_with_iceberg, sql, flags=re.IGNORECASE) + if converted_sql != sql: + logger.info("Converted query from read_parquet to Iceberg:") + logger.info(f" Original: {sql[:100]}...") + logger.info(f" Converted: {converted_sql[:100]}...") + + return converted_sql + + +def run_connection_test(config: ConnectionConfig) -> bool: + """Open a fresh connection, apply config, and try a single probe query.""" + try: + with _duckdb_connection(config) as conn: if config.catalogType == "rest": # namespace is validated at ingress against a SQL identifier # pattern, so this interpolation is safe. - result = self.connection.execute( + result = conn.execute( f"SHOW TABLES FROM iceberg_catalog.{config.namespace}" ).fetchone() elif config.tablePath: - # Read version hint first to help DuckDB try: - version_hint = self.connection.execute( + version_hint = conn.execute( "SELECT * FROM read_text(?)", [f"{config.tablePath}/metadata/version-hint.text"], ).fetchone() @@ -483,104 +446,103 @@ def test_connection(self, config: ConnectionConfig) -> bool: except Exception as e: logger.warning(f"Could not read version-hint.text: {e}") - result = self.connection.execute( + result = conn.execute( "SELECT COUNT(*) FROM iceberg_scan(?) LIMIT 1", [config.tablePath], ).fetchone() else: # Demo MinIO setup (path is hardcoded, not user input) - result = self.connection.execute( + result = conn.execute( "SELECT COUNT(*) FROM iceberg_scan('s3://movies/warehouse/demo/movies') LIMIT 1" ).fetchone() logger.info(f"Connection test successful: {result}") return True - except Exception as e: - logger.warning(f"Connection test failed: {e}") - return False + except HTTPException: + raise + except Exception as e: + logger.warning(f"Connection test failed: {e}") + return False - def execute_query(self, sql: str, config: ConnectionConfig, row_limit: int = 1000) -> QueryResponse: - """Execute SQL query and return results.""" - start_time = time.time() - try: - # Apply S3 configuration - self._apply_s3_config(config) +def run_query( + sql: str, config: ConnectionConfig, row_limit: int = 1000 +) -> QueryResponse: + """Execute ``sql`` against a fresh DuckDB session built from ``config``.""" + start_time = time.time() - # Validate table if using direct path + try: + with _duckdb_connection(config) as conn: if config.tablePath: - self._validate_iceberg_table(config.tablePath) - - # Enable profiling to get execution stats - self.connection.execute("PRAGMA enable_profiling=json") - self.connection.execute("PRAGMA profiling_output='query_profile.json'") + _validate_iceberg_table(conn, config.tablePath) # Convert any legacy read_parquet() calls to iceberg_scan() first, # then validate + LIMIT-inject the resulting SQL with sqlglot. - converted_sql = self._convert_to_iceberg_query(sql, config) + converted_sql = _convert_to_iceberg_query(sql, config) final_sql = _validate_and_limit_sql(converted_sql, row_limit) logger.info(f"Executing full query: {final_sql}") - logger.info(f"Connection config: {config.storageType}, endpoint: {config.endpoint}") + logger.info( + f"Connection config: {config.storageType}, endpoint: {config.endpoint}" + ) - # Execute query - result = self.connection.execute(final_sql) + result = conn.execute(final_sql) columns = [desc[0] for desc in result.description] rows = result.fetchall() - execution_time = int((time.time() - start_time) * 1000) + execution_time = int((time.time() - start_time) * 1000) - # Calculate stats (simplified for now) - bytes_scanned = len(str(rows)) * 2 # Rough estimate - - stats = QueryStats( - executionTimeMs=execution_time, - bytesScanned=bytes_scanned, - rowsReturned=len(rows) - ) + # Rough estimate; replaced when DuckDB profiling instrumentation lands. + bytes_scanned = len(str(rows)) * 2 - # Check if results were truncated - truncated = len(rows) >= row_limit - - logger.info(f"Query completed: {len(rows)} rows in {execution_time}ms") + stats = QueryStats( + executionTimeMs=execution_time, + bytesScanned=bytes_scanned, + rowsReturned=len(rows), + ) + truncated = len(rows) >= row_limit - return QueryResponse( - columns=columns, - rows=rows, - stats=stats, - truncated=truncated - ) + logger.info(f"Query completed: {len(rows)} rows in {execution_time}ms") - except HTTPException: - # Validation and other deliberate 400s already carry a useful - # detail string — re-raise unchanged rather than wrapping. - raise - except Exception as e: - execution_time = int((time.time() - start_time) * 1000) - logger.error(f"Query failed after {execution_time}ms: {e}") - raise HTTPException(status_code=400, detail=f"Query execution failed: {e}") + return QueryResponse( + columns=columns, + rows=rows, + stats=stats, + truncated=truncated, + ) -# Global DuckDB manager -db_manager = None + except HTTPException: + # Validation and other deliberate 400s already carry a useful detail + # string — re-raise unchanged rather than wrapping. + raise + except Exception as e: + execution_time = int((time.time() - start_time) * 1000) + logger.error(f"Query failed after {execution_time}ms: {e}") + raise HTTPException(status_code=400, detail=f"Query execution failed: {e}") @asynccontextmanager async def lifespan(app: FastAPI): """Application lifespan management.""" - global db_manager - - # Startup logger.info("Starting Cloudfloe backend...") - db_manager = DuckDBManager() - logger.info("Backend ready!") - yield + # Pre-install extensions at startup so the first real request doesn't pay + # the download cost. DuckDB caches the binaries on disk, so subsequent + # per-request connections just LOAD them. + warmup = duckdb.connect(":memory:") + try: + warmup.execute("INSTALL httpfs") + warmup.execute("INSTALL iceberg") + logger.info("DuckDB extensions warmed up") + except Exception as e: + logger.warning(f"Extension warmup failed (will retry on first request): {e}") + finally: + warmup.close() - # Shutdown + logger.info("Backend ready!") + yield logger.info("Shutting down Cloudfloe backend...") - if db_manager and db_manager.connection: - db_manager.connection.close() # Create FastAPI app @@ -625,24 +587,22 @@ async def health(): async def test_connection(request: TestConnectionRequest): """Test connection to data source.""" try: - success = db_manager.test_connection(request.connection) + success = run_connection_test(request.connection) if success: - # Get table info to show in the UI table_info = None if request.connection.tablePath: table_info = { "path": request.connection.tablePath, - "suggestedQuery": f"SELECT * FROM iceberg_scan('{request.connection.tablePath}') LIMIT 10" + "suggestedQuery": f"SELECT * FROM iceberg_scan('{request.connection.tablePath}') LIMIT 10", } return { "status": "success", "message": "Connection successful", - "tableInfo": table_info + "tableInfo": table_info, } - else: - raise HTTPException(status_code=400, detail="Connection test failed") + raise HTTPException(status_code=400, detail="Connection test failed") except HTTPException: raise @@ -655,14 +615,7 @@ async def test_connection(request: TestConnectionRequest): async def execute_query(request: QueryRequest): """Execute SQL query against data source.""" try: - # SQL is validated + LIMIT-injected inside execute_query via - # _validate_and_limit_sql. Keep the route thin. - result = db_manager.execute_query( - request.sql, - request.connection, - request.rowLimit, - ) - return result + return run_query(request.sql, request.connection, request.rowLimit) except HTTPException: raise @@ -727,4 +680,4 @@ async def get_demo_queries(): port=port, reload=True, log_level="info" - ) \ No newline at end of file + ) From 15e2ee0572f5021f5e42af53c05df6b64195a538 Mon Sep 17 00:00:00 2001 From: Gordon Murray Date: Wed, 15 Apr 2026 19:40:08 +0100 Subject: [PATCH 2/2] test(backend): add pytest suite for validators and converters Covers the two guard functions exposed by the per-request refactor (_validate_and_limit_sql, _convert_to_iceberg_query) plus the two SQL-safety helpers at the API boundary (_sql_string_literal, ConnectionConfig validators). Unit-only: no DuckDB engine, no network, ~0.5s to run. Wired into CI as a new backend-tests job so the suite runs on every push and PR. - conftest.py: adds backend/ to sys.path so tests import cleanly - pytest.ini: fixes test discovery to backend/tests - requirements.txt: pins pytest 8.1.1 and httpx 0.27.0 - ci.yml: adds backend-tests job (setup-python, pip cache, pytest) Refs #10 --- .github/workflows/ci.yml | 19 ++++ backend/pytest.ini | 4 + backend/requirements.txt | 2 + backend/tests/conftest.py | 9 ++ backend/tests/test_connection_config.py | 110 +++++++++++++++++++++++ backend/tests/test_query_conversion.py | 37 ++++++++ backend/tests/test_sql_string_literal.py | 30 +++++++ backend/tests/test_sql_validation.py | 44 +++++++++ 8 files changed, 255 insertions(+) create mode 100644 backend/pytest.ini create mode 100644 backend/tests/conftest.py create mode 100644 backend/tests/test_connection_config.py create mode 100644 backend/tests/test_query_conversion.py create mode 100644 backend/tests/test_sql_string_literal.py create mode 100644 backend/tests/test_sql_validation.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 0dadce7..e7e7bfb 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -17,6 +17,25 @@ jobs: src: "./backend ./scripts" args: "check" + backend-tests: + name: Backend tests (pytest) + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + cache: pip + cache-dependency-path: backend/requirements.txt + - name: Install backend dependencies + working-directory: backend + run: | + python -m pip install --upgrade pip + pip install -r requirements.txt + - name: Run pytest + working-directory: backend + run: pytest -v + backend-docker: name: Build backend image runs-on: ubuntu-latest diff --git a/backend/pytest.ini b/backend/pytest.ini new file mode 100644 index 0000000..6605410 --- /dev/null +++ b/backend/pytest.ini @@ -0,0 +1,4 @@ +[pytest] +testpaths = tests +python_files = test_*.py +python_functions = test_* diff --git a/backend/requirements.txt b/backend/requirements.txt index b065464..9b56929 100644 --- a/backend/requirements.txt +++ b/backend/requirements.txt @@ -5,3 +5,5 @@ uvicorn[standard]==0.37.0 duckdb==1.4.1 pydantic==2.12.0 sqlglot==30.4.3 +pytest==8.1.1 +httpx==0.27.0 diff --git a/backend/tests/conftest.py b/backend/tests/conftest.py new file mode 100644 index 0000000..816cde5 --- /dev/null +++ b/backend/tests/conftest.py @@ -0,0 +1,9 @@ +"""Make ``backend/`` importable so tests can ``from main import ...`` without +requiring the caller to set ``PYTHONPATH``.""" + +import sys +from pathlib import Path + +BACKEND_DIR = Path(__file__).resolve().parent.parent +if str(BACKEND_DIR) not in sys.path: + sys.path.insert(0, str(BACKEND_DIR)) diff --git a/backend/tests/test_connection_config.py b/backend/tests/test_connection_config.py new file mode 100644 index 0000000..14e2649 --- /dev/null +++ b/backend/tests/test_connection_config.py @@ -0,0 +1,110 @@ +"""Tests for ``ConnectionConfig`` Pydantic validators. + +These guard the API boundary: any user-supplied config value that would +break SQL (quotes, semicolons, control chars) must be rejected before it +can reach DuckDB. +""" + +import pytest +from pydantic import ValidationError + +from main import ConnectionConfig + + +def _base(**overrides): + defaults = dict( + storageType="s3", + endpoint="s3.amazonaws.com", + accessKey="AKIAFAKE", + secretKey="secret", + ) + defaults.update(overrides) + return defaults + + +def test_valid_minimal_config(): + cfg = ConnectionConfig(**_base()) + assert cfg.storageType == "s3" + assert cfg.catalogType == "none" + assert cfg.region == "us-east-1" + + +def test_storage_type_rejects_unknown(): + with pytest.raises(ValidationError): + ConnectionConfig(**_base(storageType="gcs")) + + +def test_catalog_type_rejects_unknown(): + with pytest.raises(ValidationError): + ConnectionConfig(**_base(catalogType="hive")) + + +@pytest.mark.parametrize( + "bad_path", + [ + "s3://bucket/warehouse; DROP TABLE foo", + "s3://bucket/warehouse' OR '1'='1", + "s3://bucket/warehouse\nmalicious", + "http://bucket/warehouse", # wrong scheme + "bucket/warehouse", # missing scheme + ], +) +def test_table_path_rejects_injection(bad_path): + with pytest.raises(ValidationError): + ConnectionConfig(**_base(tablePath=bad_path)) + + +def test_table_path_strips_trailing_slash_and_metadata(): + cfg = ConnectionConfig(**_base(tablePath="s3://bucket/warehouse/db/table/metadata")) + assert cfg.tablePath == "s3://bucket/warehouse/db/table" + + cfg = ConnectionConfig(**_base(tablePath="s3://bucket/warehouse/db/table/")) + assert cfg.tablePath == "s3://bucket/warehouse/db/table" + + +@pytest.mark.parametrize( + "bad_ns", + [ + "default; DROP", + "db with space", + "db'name", + "123leading_digit", + ], +) +def test_namespace_rejects_non_identifier(bad_ns): + with pytest.raises(ValidationError): + ConnectionConfig(**_base(catalogType="rest", catalogEndpoint="http://cat:8181", namespace=bad_ns)) + + +def test_namespace_accepts_identifier(): + cfg = ConnectionConfig(**_base( + catalogType="rest", + catalogEndpoint="http://cat:8181", + namespace="my_db_1", + )) + assert cfg.namespace == "my_db_1" + + +@pytest.mark.parametrize( + "bad_region", + ["eu west 1", "eu-west-1;", "eu-west-1'"], +) +def test_region_rejects_bad_chars(bad_region): + with pytest.raises(ValidationError): + ConnectionConfig(**_base(region=bad_region)) + + +def test_credentials_reject_newline_and_null(): + with pytest.raises(ValidationError): + ConnectionConfig(**_base(accessKey="AKIA\nBAD")) + with pytest.raises(ValidationError): + ConnectionConfig(**_base(secretKey="sec\x00ret")) + + +def test_catalog_endpoint_requires_http_scheme(): + with pytest.raises(ValidationError): + ConnectionConfig(**_base( + catalogType="rest", + catalogEndpoint="ftp://cat:8181", + namespace="db", + )) diff --git a/backend/tests/test_query_conversion.py b/backend/tests/test_query_conversion.py new file mode 100644 index 0000000..8586002 --- /dev/null +++ b/backend/tests/test_query_conversion.py @@ -0,0 +1,37 @@ +from main import _convert_to_iceberg_query, ConnectionConfig + +def test_parquet_to_iceberg_scan(): + config = ConnectionConfig( + storageType="s3", + endpoint="s3.amazonaws.com", + accessKey="key", + secretKey="secret" + ) + sql = "SELECT * FROM read_parquet('s3://bucket/path/to/table/*.parquet')" + converted = _convert_to_iceberg_query(sql, config) + assert "iceberg_scan('s3://bucket/path/to/table')" in converted + +def test_parquet_to_rest_catalog(): + config = ConnectionConfig( + storageType="s3", + endpoint="s3.amazonaws.com", + accessKey="key", + secretKey="secret", + catalogType="rest", + catalogEndpoint="http://localhost:8181", + namespace="db" + ) + sql = "SELECT * FROM read_parquet('s3://bucket/path/to/my_table/*.parquet')" + converted = _convert_to_iceberg_query(sql, config) + assert "iceberg_catalog.db.my_table" in converted + +def test_non_parquet_query_unchanged(): + config = ConnectionConfig( + storageType="s3", + endpoint="s3.amazonaws.com", + accessKey="key", + secretKey="secret" + ) + sql = "SELECT * FROM my_table" + converted = _convert_to_iceberg_query(sql, config) + assert converted == sql diff --git a/backend/tests/test_sql_string_literal.py b/backend/tests/test_sql_string_literal.py new file mode 100644 index 0000000..126380c --- /dev/null +++ b/backend/tests/test_sql_string_literal.py @@ -0,0 +1,30 @@ +"""Tests for ``_sql_string_literal``. + +It's the last line of defence when we have to interpolate a pre-validated +value into a DuckDB statement that doesn't accept ``?`` placeholders +(CREATE SECRET, ATTACH). It must double any embedded single quote and +reject control characters. +""" + +import pytest + +from main import _sql_string_literal + + +def test_plain_value_quoted(): + assert _sql_string_literal("plain") == "'plain'" + + +def test_single_quote_doubled(): + assert _sql_string_literal("O'Brien") == "'O''Brien'" + assert _sql_string_literal("a';--") == "'a'';--'" + + +def test_tab_allowed(): + assert _sql_string_literal("a\tb") == "'a\tb'" + + +@pytest.mark.parametrize("bad", ["has\x00null", "has\nnewline", "has\rreturn", "\x01"]) +def test_control_chars_rejected(bad): + with pytest.raises(ValueError): + _sql_string_literal(bad) diff --git a/backend/tests/test_sql_validation.py b/backend/tests/test_sql_validation.py new file mode 100644 index 0000000..aee58a9 --- /dev/null +++ b/backend/tests/test_sql_validation.py @@ -0,0 +1,44 @@ +import pytest +from main import _validate_and_limit_sql +from fastapi import HTTPException + +def test_select_allowed(): + sql = "SELECT * FROM my_table" + validated = _validate_and_limit_sql(sql, 100) + assert "LIMIT 100" in validated + +def test_limit_not_overwritten(): + sql = "SELECT * FROM my_table LIMIT 10" + validated = _validate_and_limit_sql(sql, 100) + assert "LIMIT 10" in validated + assert "LIMIT 100" not in validated + +def test_forbidden_statements(): + forbidden = [ + "INSERT INTO table VALUES (1)", + "DROP TABLE my_table", + "DELETE FROM my_table WHERE id = 1", + "UPDATE my_table SET val = 1", + "CREATE TABLE oops (id INT)", + "ATTACH 'db.db' AS db", + ] + for sql in forbidden: + with pytest.raises(HTTPException) as exc: + _validate_and_limit_sql(sql, 100) + assert exc.value.status_code == 400 + +def test_multi_statement_rejected(): + sql = "SELECT 1; SELECT 2;" + with pytest.raises(HTTPException) as exc: + _validate_and_limit_sql(sql, 100) + assert exc.value.status_code == 400 + +def test_with_clause_allowed(): + sql = "WITH cte AS (SELECT 1) SELECT * FROM cte" + validated = _validate_and_limit_sql(sql, 100) + assert "LIMIT 100" in validated + +def test_union_allowed(): + sql = "SELECT 1 UNION SELECT 2" + validated = _validate_and_limit_sql(sql, 100) + assert "LIMIT 100" in validated