from typing import Any from fastapi import HTTPException, status from sqlalchemy import text from sqlalchemy.engine import Engine from sqlalchemy.exc import SQLAlchemyError from app.core.config import get_settings from app.core.constants import ApiStatus from app.core.database import legacy_engine from app.core.http.pagination import bounded_limit from app.modules.legacy_mysql.constants import ( LEGACY_HEALTH_SQL, LEGACY_LIMIT_CLAUSE, LEGACY_LIMIT_MARKER, LEGACY_MYSQL_CONNECTION_FAILED_TEMPLATE, LEGACY_SELECT_PREFIX, LEGACY_SQL_TRAILING_TERMINATOR, LegacyQueryError, LegacyQueryName, LegacyResponseKey, ) from app.modules.legacy_mysql.services.common import FORBIDDEN_SQL_TOKENS, _normalize_sql, _query_name_text, _row_to_dict class LegacyQueryMixin: @staticmethod def _ensure_engine() -> Engine: if legacy_engine is None: raise HTTPException( status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail=LegacyQueryError.DATABASE_NOT_CONFIGURED, ) return legacy_engine @staticmethod def _ensure_readonly(sql: str) -> None: stripped = sql.strip().lower() if not stripped.startswith(LEGACY_SELECT_PREFIX): raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail=LegacyQueryError.ONLY_SELECT_ALLOWED, ) tokens = {token.strip(" ;,\n\t") for token in stripped.replace("(", " ").split()} if tokens & FORBIDDEN_SQL_TOKENS: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail=LegacyQueryError.FORBIDDEN_SQL_TOKEN, ) @staticmethod def _allowed_queries() -> dict[str, str]: settings = get_settings() queries = { _query_name_text(name): sql for name, sql in settings.legacy_allowed_queries.items() } if settings.legacy_project_query: queries.setdefault(LegacyQueryName.PROJECTS.value, settings.legacy_project_query) if settings.legacy_task_query: queries.setdefault(LegacyQueryName.TASKS.value, settings.legacy_task_query) return queries def health(self) -> dict[str, str]: engine = self._ensure_engine() try: with engine.connect() as conn: conn.execute(text(LEGACY_HEALTH_SQL)) except SQLAlchemyError as exc: raise HTTPException( status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail=LEGACY_MYSQL_CONNECTION_FAILED_TEMPLATE.format(error=exc), ) from exc return {LegacyResponseKey.STATUS: ApiStatus.OK} def execute_readonly( self, sql: str, params: dict[str, Any] | None = None, limit: int = 100, ) -> dict[str, Any]: """Execute a readonly SQL statement only when it matches the allowlist.""" normalized_sql = _normalize_sql(sql) for allowed_sql in self._allowed_queries().values(): if _normalize_sql(allowed_sql) == normalized_sql: return self._execute_readonly_sql(allowed_sql, params, limit) raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, detail=LegacyQueryError.QUERY_NOT_ALLOWED, ) def execute_allowed_query( self, query_name: str | None, params: dict[str, Any] | None = None, limit: int = 100, ) -> dict[str, Any]: queries = self._allowed_queries() normalized_name = _query_name_text(query_name) sql = queries.get(normalized_name) if not sql: raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, detail=LegacyQueryError.QUERY_NOT_ALLOWED, ) return self._execute_readonly_sql(sql, params, limit) def _execute_readonly_sql( self, sql: str, params: dict[str, Any] | None = None, limit: int = 100, ) -> dict[str, Any]: self._ensure_readonly(sql) engine = self._ensure_engine() params = dict(params or {}) try: params[LegacyResponseKey.LIMIT] = bounded_limit( params.get(LegacyResponseKey.LIMIT, limit) ) except (TypeError, ValueError) as exc: raise HTTPException( status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, detail=LegacyQueryError.INVALID_LIMIT, ) from exc limited_sql = sql if LEGACY_LIMIT_MARKER not in sql.lower(): limited_sql = f"{sql.rstrip(LEGACY_SQL_TRAILING_TERMINATOR)}{LEGACY_LIMIT_CLAUSE}" with engine.connect() as conn: result = conn.execute(text(limited_sql), params) rows = [_row_to_dict(row) for row in result.mappings().all()] columns = list(rows[0].keys()) if rows else [] return { LegacyResponseKey.COLUMNS: columns, LegacyResponseKey.ROWS: rows, LegacyResponseKey.ROW_COUNT: len(rows), } def fetch_default_projects(self, limit: int = 100) -> dict[str, Any]: settings = get_settings() if not settings.legacy_project_query: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail=LegacyQueryError.PROJECT_QUERY_NOT_CONFIGURED, ) return self.execute_allowed_query( LegacyQueryName.PROJECTS, {LegacyResponseKey.LIMIT: limit}, limit=limit, ) def fetch_default_tasks(self, limit: int = 100) -> dict[str, Any]: settings = get_settings() if not settings.legacy_task_query: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail=LegacyQueryError.TASK_QUERY_NOT_CONFIGURED, ) return self.execute_allowed_query( LegacyQueryName.TASKS, {LegacyResponseKey.LIMIT: limit}, limit=limit, )