company-ai-platform/app/modules/legacy_mysql/services/query.py

import re
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

_SQL_COMMENT_MARKERS = ("--", "#", "/*", "*/")
_SQL_QUOTED_CONTENT_PATTERN = re.compile(
    r"""'(?:''|\\.|[^'])*'|"(?:""|\\.|[^"])*"|`(?:``|[^`])*`""",
    flags=re.DOTALL,
)
_SQL_WORD_PATTERN = re.compile(r"[a-z_]+")


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()
        scrubbed = _SQL_QUOTED_CONTENT_PATTERN.sub(" ", stripped)
        if any(marker in scrubbed for marker in _SQL_COMMENT_MARKERS):
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=LegacyQueryError.SQL_COMMENTS_NOT_ALLOWED,
            )
        statement = scrubbed.rstrip()
        if statement.endswith(LEGACY_SQL_TRAILING_TERMINATOR):
            statement = statement[:-1].rstrip()
        if LEGACY_SQL_TRAILING_TERMINATOR in statement:
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=LegacyQueryError.SINGLE_STATEMENT_REQUIRED,
            )
        if not re.match(rf"^{LEGACY_SELECT_PREFIX}\b", statement, flags=re.IGNORECASE):
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=LegacyQueryError.ONLY_SELECT_ALLOWED,
            )
        tokens = set(_SQL_WORD_PATTERN.findall(statement.lower()))
        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:
            limit_value = bounded_limit(
                params.get(LegacyResponseKey.LIMIT, limit)
            )
            params[LegacyResponseKey.LIMIT] = limit_value
        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().fetchmany(limit_value)
            ]
        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,
        )