backend/database.py

from __future__ import annotations

import datetime as dt
import decimal
import math
import re
import threading
from pathlib import Path
from typing import Any

import duckdb
from sqlglot import exp, parse_one
from sqlglot.errors import ParseError


PROJECT_ROOT = Path(__file__).resolve().parent.parent
DATA_DIR = PROJECT_ROOT / "data"
ALLOWED_TABLES = {"air_quality", "states", "ncap_funding"}

DISALLOWED_SQL = re.compile(
    r"\b("
    r"attach|copy|export|import|install|load|pragma|call|set|create|drop|alter|"
    r"insert|update|delete|merge|truncate|grant|revoke|vacuum|checkpoint|"
    r"read_csv|read_json|read_parquet|read_text|glob|httpfs|sqlite_scan|"
    r"postgres_scan|mysql_scan|getenv|current_setting"
    r")\b",
    re.IGNORECASE,
)
DISALLOWED_FUNCTION = re.compile(
    r"\b(?:"
    r"read|scan|glob|sniff|query_table|duckdb|pragma|current_setting|"
    r"getenv|secret|shell|system|parquet|csv|json|sqlite|postgres|mysql"
    r")\w*\s*\(",
    re.IGNORECASE,
)
class AirQualityDatabase:
    def __init__(self) -> None:
        self._connection: duckdb.DuckDBPyConnection | None = None
        self._lock = threading.Lock()

    def initialize(self) -> None:
        required = [
            DATA_DIR / "AQ_met_data.csv",
            DATA_DIR / "states_data.csv",
            DATA_DIR / "ncap_funding_data.csv",
        ]
        missing = [str(path) for path in required if not path.exists()]
        if missing:
            raise RuntimeError(f"Missing data files: {', '.join(missing)}")

        connection = duckdb.connect(database=":memory:")
        connection.execute("SET threads = 2")
        connection.execute(
            """
            CREATE TABLE air_quality AS
            SELECT
                TRY_CAST("Timestamp" AS DATE) AS timestamp,
                "State" AS state,
                "City" AS city,
                "Station" AS station,
                "site_id" AS site_id,
                TRY_CAST("Year" AS INTEGER) AS year,
                TRY_CAST("PM2.5 (µg/m³)" AS DOUBLE) AS pm25,
                TRY_CAST("PM10 (µg/m³)" AS DOUBLE) AS pm10,
                TRY_CAST("NO (µg/m³)" AS DOUBLE) AS no,
                TRY_CAST("NO2 (µg/m³)" AS DOUBLE) AS no2,
                TRY_CAST("NOx (ppb)" AS DOUBLE) AS nox,
                TRY_CAST("NH3 (µg/m³)" AS DOUBLE) AS nh3,
                TRY_CAST("SO2 (µg/m³)" AS DOUBLE) AS so2,
                TRY_CAST("CO (mg/m³)" AS DOUBLE) AS co,
                TRY_CAST("Ozone (µg/m³)" AS DOUBLE) AS ozone,
                TRY_CAST("AT (°C)" AS DOUBLE) AS temperature,
                TRY_CAST("RH (%)" AS DOUBLE) AS humidity,
                TRY_CAST("WS (m/s)" AS DOUBLE) AS wind_speed,
                TRY_CAST("WD (deg)" AS DOUBLE) AS wind_direction,
                TRY_CAST("RF (mm)" AS DOUBLE) AS rainfall,
                TRY_CAST("TOT-RF (mm)" AS DOUBLE) AS total_rainfall,
                TRY_CAST("SR (W/mt2)" AS DOUBLE) AS solar_radiation,
                TRY_CAST("BP (mmHg)" AS DOUBLE) AS pressure,
                TRY_CAST("VWS (m/s)" AS DOUBLE) AS vertical_wind_speed
            FROM read_csv_auto(?, header = true, ignore_errors = true)
            """,
            [str(required[0])],
        )
        connection.execute(
            """
            CREATE TABLE states AS
            SELECT
                state,
                TRY_CAST(population AS BIGINT) AS population,
                TRY_CAST("area (km2)" AS DOUBLE) AS area_km2,
                TRY_CAST(isUnionTerritory AS BOOLEAN) AS is_union_territory
            FROM read_csv_auto(?, header = true, ignore_errors = true)
            """,
            [str(required[1])],
        )
        connection.execute(
            """
            CREATE TABLE ncap_funding AS
            SELECT
                state,
                city,
                TRY_CAST("Amount released during FY 2019-20" AS DOUBLE) AS fy_2019_20,
                TRY_CAST("Amount released during FY 2020-21" AS DOUBLE) AS fy_2020_21,
                TRY_CAST("Amount released during FY 2021-22" AS DOUBLE) AS fy_2021_22,
                TRY_CAST("Total fund released" AS DOUBLE) AS total_fund_released,
                TRY_CAST("Utilisation as on June 2022" AS DOUBLE) AS utilisation_june_2022
            FROM read_csv_auto(?, header = true, ignore_errors = true)
            """,
            [str(required[2])],
        )
        connection.execute("ANALYZE")
        self._connection = connection

    @staticmethod
    def validate_sql(sql: str) -> str:
        cleaned = sql.strip()
        if cleaned.startswith("```"):
            cleaned = re.sub(r"^```(?:sql)?\s*", "", cleaned, flags=re.IGNORECASE)
            cleaned = re.sub(r"\s*```$", "", cleaned)
        cleaned = cleaned.rstrip(";").strip()

        if not cleaned or not re.match(r"^(select|with)\b", cleaned, re.IGNORECASE):
            raise ValueError("Only read-only SELECT queries are allowed.")
        if ";" in cleaned or "--" in cleaned or "/*" in cleaned or "*/" in cleaned:
            raise ValueError("Multiple statements and SQL comments are not allowed.")
        if DISALLOWED_SQL.search(cleaned):
            raise ValueError("The generated query contains a blocked operation.")
        if DISALLOWED_FUNCTION.search(cleaned):
            raise ValueError("The generated query contains a blocked function.")

        try:
            expression = parse_one(cleaned, read="duckdb")
        except ParseError as exc:
            raise ValueError("The generated query is not valid DuckDB SQL.") from exc
        if not isinstance(expression, (exp.Select, exp.Union, exp.Intersect, exp.Except)):
            raise ValueError("Only read-only SELECT queries are allowed.")

        ctes = {
            cte.alias_or_name.lower()
            for cte in expression.find_all(exp.CTE)
            if cte.alias_or_name
        }
        referenced = {
            table.name.lower()
            for table in expression.find_all(exp.Table)
            if table.name
        }
        invalid = referenced - ALLOWED_TABLES - ctes
        if invalid:
            raise ValueError(
                f"Query referenced an unavailable table: {', '.join(sorted(invalid))}."
            )
        if not referenced.intersection(ALLOWED_TABLES):
            raise ValueError("Query must use one of the VayuChat data tables.")
        return cleaned

    def execute(self, sql: str, max_rows: int = 200) -> tuple[list[str], list[dict], bool]:
        if self._connection is None:
            raise RuntimeError("Database has not been initialized.")
        safe_sql = self.validate_sql(sql)
        wrapped = f"SELECT * FROM ({safe_sql}) AS vayuchat_result LIMIT {max_rows + 1}"
        with self._lock:
            cursor = self._connection.execute(wrapped)
            columns = [column[0] for column in cursor.description]
            values = cursor.fetchall()
        truncated = len(values) > max_rows
        rows = [
            {
                column: self._json_safe(value)
                for column, value in zip(columns, row, strict=True)
            }
            for row in values[:max_rows]
        ]
        return columns, rows, truncated

    def stats(self) -> dict[str, Any]:
        if self._connection is None:
            return {"ready": False}
        with self._lock:
            row = self._connection.execute(
                """
                SELECT
                    MIN(timestamp),
                    MAX(timestamp),
                    COUNT(*),
                    COUNT(DISTINCT city)
                FROM air_quality
                """
            ).fetchone()
        return {
            "ready": True,
            "first_date": self._json_safe(row[0]),
            "last_date": self._json_safe(row[1]),
            "records": row[2],
            "cities": row[3],
        }

    @staticmethod
    def _json_safe(value: Any) -> Any:
        if value is None:
            return None
        if isinstance(value, (dt.date, dt.datetime, dt.time)):
            return value.isoformat()
        if isinstance(value, decimal.Decimal):
            return float(value)
        if isinstance(value, float) and (math.isnan(value) or math.isinf(value)):
            return None
        return value


database = AirQualityDatabase()