testing/test_runner.py

"""
SQL Test Runner using Modal (SECONDARY SPONSOR)
Provides serverless SQL execution for testing queries
"""

import os
from typing import Dict, Any

# Try to import Modal, but don't fail if not available
try:
    import modal
    MODAL_AVAILABLE = True
except ImportError:
    MODAL_AVAILABLE = False
    print("⚠️ Modal not available. Install with: pip install modal")

# Try to import DuckDB for local fallback
try:
    import duckdb
    DUCKDB_AVAILABLE = True
except ImportError:
    DUCKDB_AVAILABLE = False
    print("⚠️ DuckDB not available for local testing fallback")


class SQLTestRunner:
    """
    Test runner for SQL queries
    Uses Modal for serverless execution, with DuckDB fallback
    """
    
    def __init__(self):
        self.modal_available = MODAL_AVAILABLE
        self.duckdb_available = DUCKDB_AVAILABLE
        
        # Initialize Modal app if available
        if self.modal_available:
            try:
                self.app = modal.App("codeflow-sql-tester")
                self._setup_modal()
            except Exception as e:
                print(f"⚠️ Modal setup failed: {e}")
                self.modal_available = False
    
    def _setup_modal(self):
        """Setup Modal function for SQL execution"""
        if not self.modal_available:
            return
        
        try:
            # Define Modal image with DuckDB
            image = modal.Image.debian_slim().pip_install("duckdb==0.9.2")
            
            # Define Modal function
            @self.app.function(image=image, timeout=60)
            def execute_sql_modal(sql: str, sample_data: dict = None) -> dict:
                """Execute SQL in Modal sandbox"""
                import duckdb
                
                try:
                    # Create in-memory DuckDB database
                    conn = duckdb.connect(':memory:')
                    
                    # If sample data provided, create tables
                    if sample_data:
                        for table_name, data in sample_data.items():
                            conn.execute(f"CREATE TABLE {table_name} AS SELECT * FROM ?", [data])
                    
                    # Execute the SQL
                    result = conn.execute(sql).fetchall()
                    columns = [desc[0] for desc in conn.description] if conn.description else []
                    
                    conn.close()
                    
                    return {
                        "success": True,
                        "rows": result[:100],  # Limit to 100 rows
                        "row_count": len(result),
                        "columns": columns,
                        "error": None
                    }
                    
                except Exception as e:
                    return {
                        "success": False,
                        "rows": [],
                        "row_count": 0,
                        "columns": [],
                        "error": str(e)
                    }
            
            self.execute_sql_modal = execute_sql_modal
            
        except Exception as e:
            print(f"⚠️ Modal function setup failed: {e}")
            self.modal_available = False
    
    def test_sql(self, sql: str, sample_data: Dict = None) -> Dict[str, Any]:
        """
        Test SQL query execution
        
        Args:
            sql: SQL query to test
            sample_data: Optional sample data as dict of table_name -> rows
            
        Returns:
            Dict with test results
        """
        
        # Try Modal first if available
        if self.modal_available:
            try:
                result = self.execute_sql_modal.remote(sql, sample_data)
                result["execution_method"] = "Modal (Serverless)"
                return result
            except Exception as e:
                print(f"⚠️ Modal execution failed, falling back to local: {e}")
        
        # Fallback to local DuckDB
        if self.duckdb_available:
            return self._test_sql_local(sql, sample_data)
        
        # No execution method available
        return {
            "success": False,
            "rows": [],
            "row_count": 0,
            "columns": [],
            "error": "No SQL execution method available. Install Modal or DuckDB.",
            "execution_method": "None"
        }
    
    def _test_sql_local(self, sql: str, sample_data: Dict = None) -> Dict[str, Any]:
        """Execute SQL locally using DuckDB"""
        try:
            conn = duckdb.connect(':memory:')
            
            # Create sample tables if provided
            if sample_data:
                for table_name, data in sample_data.items():
                    conn.execute(f"CREATE TABLE {table_name} AS SELECT * FROM ?", [data])
            
            # Execute SQL
            result = conn.execute(sql).fetchall()
            columns = [desc[0] for desc in conn.description] if conn.description else []
            
            conn.close()
            
            return {
                "success": True,
                "rows": result[:100],  # Limit to 100 rows
                "row_count": len(result),
                "columns": columns,
                "error": None,
                "execution_method": "DuckDB (Local)"
            }
            
        except Exception as e:
            return {
                "success": False,
                "rows": [],
                "row_count": 0,
                "columns": [],
                "error": str(e),
                "execution_method": "DuckDB (Local)"
            }
    
    def validate_syntax(self, sql: str) -> Dict[str, Any]:
        """
        Quick syntax validation without execution
        
        Args:
            sql: SQL query to validate
            
        Returns:
            Dict with validation results
        """
        if not self.duckdb_available:
            return {
                "valid": None,
                "error": "DuckDB not available for syntax validation"
            }
        
        try:
            conn = duckdb.connect(':memory:')
            # Try to prepare the query (doesn't execute, just validates syntax)
            conn.execute(f"EXPLAIN {sql}")
            conn.close()
            
            return {
                "valid": True,
                "error": None
            }
            
        except Exception as e:
            return {
                "valid": False,
                "error": str(e)
            }
    
    def get_execution_plan(self, sql: str) -> Dict[str, Any]:
        """
        Get query execution plan
        
        Args:
            sql: SQL query
            
        Returns:
            Dict with execution plan
        """
        if not self.duckdb_available:
            return {
                "success": False,
                "plan": None,
                "error": "DuckDB not available"
            }
        
        try:
            conn = duckdb.connect(':memory:')
            plan = conn.execute(f"EXPLAIN {sql}").fetchall()
            conn.close()
            
            return {
                "success": True,
                "plan": "\n".join([str(row[0]) for row in plan]),
                "error": None
            }
            
        except Exception as e:
            return {
                "success": False,
                "plan": None,
                "error": str(e)
            }


# Singleton instance
_test_runner = None

def get_test_runner() -> SQLTestRunner:
    """Get or create the global test runner instance"""
    global _test_runner
    if _test_runner is None:
        _test_runner = SQLTestRunner()
    return _test_runner