Update app.py

app.py

@@ -4,6 +4,8 @@ from datetime import datetime
 from pathlib import Path
 from typing import Dict, Any, Tuple, Optional
 from contextlib import contextmanager
+import signal
+from functools import wraps
 
 import duckdb
 import pandas as pd
@@ -36,6 +38,52 @@ EXPORT_DIR.mkdir(exist_ok=True)
 # Query timeout in seconds
 QUERY_TIMEOUT_SECONDS = 30
 
+# -------------------------------------------------------------------
+# Query Timeout Handler (since DuckDB doesn't support statement_timeout)
+# -------------------------------------------------------------------
+class QueryTimeoutError(Exception):
+    """Raised when a query exceeds the timeout limit"""
+    pass
+
+
+def with_timeout(timeout_seconds: int = QUERY_TIMEOUT_SECONDS):
+    """
+    Decorator to add timeout to query functions.
+    Note: This is a simplified version. For production, consider using
+    concurrent.futures or multiprocessing for better timeout handling.
+    
+    Args:
+        timeout_seconds: Maximum execution time in seconds
+    """
+    def decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            # For now, we'll log the timeout but not enforce it
+            # since signal.alarm doesn't work well with threads
+            # and DuckDB doesn't support native query timeouts
+            logger.debug(f"Starting {func.__name__} with {timeout_seconds}s timeout")
+            start_time = datetime.now()
+            
+            try:
+                result = func(*args, **kwargs)
+                elapsed = (datetime.now() - start_time).total_seconds()
+                
+                if elapsed > timeout_seconds:
+                    logger.warning(
+                        f"{func.__name__} exceeded timeout: {elapsed:.2f}s > {timeout_seconds}s"
+                    )
+                else:
+                    logger.debug(f"{func.__name__} completed in {elapsed:.2f}s")
+                
+                return result
+            except Exception as e:
+                elapsed = (datetime.now() - start_time).total_seconds()
+                logger.error(f"{func.__name__} failed after {elapsed:.2f}s: {str(e)}")
+                raise
+        
+        return wrapper
+    return decorator
+
 # -------------------------------------------------------------------
 # Data Models
 # -------------------------------------------------------------------
@@ -81,8 +129,8 @@ def get_db_connection():
         logger.info("Establishing MotherDuck connection")
         conn = duckdb.connect(f"md:?motherduck_token={token}")
         
-        # Set query timeout
-        conn.execute(f"SET statement_timeout='{QUERY_TIMEOUT_SECONDS}s'")
+        # Note: DuckDB/MotherDuck doesn't support statement_timeout like PostgreSQL
+        # Query timeouts should be handled at application level with threading/async
         
         yield conn
         
@@ -114,13 +162,21 @@ def execute_query(conn: duckdb.DuckDBPyConnection, query: str,
     Raises:
         Exception: For query execution errors
     """
+    start_time = datetime.now()
     try:
         logger.info(f"Executing {description}")
         result = conn.execute(query).df()
-        logger.info(f"{description} completed: {len(result)} rows returned")
+        elapsed = (datetime.now() - start_time).total_seconds()
+        logger.info(f"{description} completed: {len(result)} rows in {elapsed:.2f}s")
+        
+        # Warn if query is slow
+        if elapsed > QUERY_TIMEOUT_SECONDS:
+            logger.warning(f"{description} exceeded timeout threshold: {elapsed:.2f}s")
+        
         return result
     except Exception as e:
-        logger.error(f"{description} failed: {str(e)}")
+        elapsed = (datetime.now() - start_time).total_seconds()
+        logger.error(f"{description} failed after {elapsed:.2f}s: {str(e)}")
         raise Exception(f"Query execution failed for {description}: {str(e)}")