backend/services/sql_service.py

import os
import sqlite3
import re
import time
import pandas as pd
from typing import Dict, Any
from backend import config

class SQLService:
    def __init__(self):
        print("[SQLService] Initializing SQLService singleton...")

    def validate_sql(self, sql: str) -> bool:
        """
        Validates that the SQL query is a SELECT statement and does not contain
        forbidden modifying operations like INSERT, UPDATE, DELETE, DROP, CREATE, ALTER, EXEC.
        """
        print(f"[SQLService] Validating SQL: {sql}")
        try:
            if not sql or not isinstance(sql, str):
                return False

            sql_upper = sql.upper().strip()

            # The SQL must be a SELECT statement
            if "SELECT" not in sql_upper:
                print("[SQLService] Validation failed: SQL does not contain SELECT keyword.")
                return False

            # Forbidden keywords pattern with word boundaries to avoid false positives (e.g. column name like 'create_date')
            forbidden_pattern = r"\b(INSERT|UPDATE|DELETE|DROP|CREATE|ALTER|EXEC)\b"
            if re.search(forbidden_pattern, sql_upper):
                print("[SQLService] Validation failed: SQL contains forbidden modifying keywords.")
                return False

            return True
        except Exception as e:
            print(f"[SQLService] Error during SQL validation: {e}")
            return False

    def _clean_query_tsql_artifacts(self, sql: str) -> str:
        """
        Lightweight cleaner that removes accidental T-SQL syntax the LLM might
        hallucinate when generating queries for a schema DB that was originally T-SQL.
        Only applied to SELECT queries.
        """
        cleaned = sql

        # Convert SELECT TOP N ... to SELECT ... LIMIT N
        top_match = re.search(
            r'\bSELECT\s+TOP\s+(\d+)\b',
            cleaned, re.IGNORECASE
        )
        if top_match:
            limit_n = top_match.group(1)
            # Remove TOP N from SELECT
            cleaned = re.sub(r'\bSELECT\s+TOP\s+\d+\b', 'SELECT', cleaned, flags=re.IGNORECASE)
            
            # Check for semicolon at the end
            has_semicolon = cleaned.strip().endswith(';')
            cleaned_body = cleaned.strip()
            if has_semicolon:
                cleaned_body = cleaned_body[:-1].strip()
                
            # Remove any existing LIMIT to avoid duplication
            cleaned_body = re.sub(r'\bLIMIT\s+\d+\s*$', '', cleaned_body, flags=re.IGNORECASE).strip()
            
            # Append LIMIT at the end
            cleaned = f"{cleaned_body} LIMIT {limit_n}"
            if has_semicolon:
                cleaned += ";"

        # Remove WITH (NOLOCK) hints
        cleaned = re.sub(r'\bWITH\s*\(\s*NOLOCK\s*\)', '', cleaned, flags=re.IGNORECASE)

        # Remove square brackets around identifiers: [EMPNO] → EMPNO
        cleaned = re.sub(r'\[([^\]]+)\]', r'\1', cleaned)

        # Clean up any double spaces left behind
        cleaned = re.sub(r'  +', ' ', cleaned).strip()

        return cleaned

    def execute_query(self, db_name: str, sql: str, mode: str = "csv") -> Dict[str, Any]:
        """
        Executes the provided SQL query against the specified SQLite database.
        Validates the SQL first, limits output using pandas, measures execution time,
        and translates SQLite/pandas dtypes to JSON-serializable Python objects.
        """
        print(f"[SQLService] Executing SQL on {db_name}.db (mode={mode}): {sql}")
        try:
            # 1. Validate SQL safety
            if not self.validate_sql(sql):
                raise ValueError("Dangerous or invalid SQL query. Only SELECT queries are allowed.")

            # 1b. Clean any accidental T-SQL artifacts from LLM-generated queries
            if mode == "schema":
                sql = self._clean_query_tsql_artifacts(sql)
                print(f"[SQLService] After T-SQL cleanup: {sql}")

            if mode == "csv":
                db_path = os.path.join(config.DB_DIR, f"{db_name}.db")
            else:
                db_path = os.path.join(config.SCHEMA_DB_DIR, f"{db_name}.db")

            if not os.path.exists(db_path):
                raise FileNotFoundError(f"Database {db_name}.db does not exist.")


            conn = sqlite3.connect(db_path)
            
            # 2. Execute query and measure time
            start_time = time.perf_counter()
            try:
                df = pd.read_sql_query(sql, conn)
            except Exception as sql_err:
                conn.close()
                print(f"[SQLService] SQLite Error: {sql_err}")
                raise ValueError(f"SQLite error: {str(sql_err)}")
            
            execution_time_ms = (time.perf_counter() - start_time) * 1000.0
            conn.close()

            # 3. Limit rows to MAX_ROWS_RETURNED
            df_limited = df.head(config.MAX_ROWS_RETURNED)
            
            # 4. Extract columns and rows, converting numpy/pandas NaN and types to native types
            columns = df_limited.columns.tolist()
            rows = []
            for row in df_limited.itertuples(index=False):
                clean_row = []
                for val in row:
                    if pd.isna(val):
                        clean_row.append(None)
                    elif hasattr(val, "item") and not isinstance(val, (str, bytes)):
                        # Convert numpy/pandas numeric scalars to Python scalars
                        clean_row.append(val.item())
                    else:
                        clean_row.append(val)
                rows.append(clean_row)

            result = {
                "columns": columns,
                "rows": rows,
                "row_count": len(rows),
                "execution_time_ms": round(execution_time_ms, 2)
            }
            print(f"[SQLService] Query execution succeeded. Returned {result['row_count']} rows in {result['execution_time_ms']} ms.")
            return result

        except Exception as e:
            print(f"[SQLService] Error during SQL execution: {e}")
            raise ValueError(str(e))

# Singleton instance
sql_service = SQLService()