opsxsql/backend/app/services/query_service.py

"""统一多数据库查询服务 — 支持 MySQL、PostgreSQL、Redis、MongoDB 的只读查询。

核心功能:
1. 统一入口: 根据 db_type 自动路由到对应查询处理器
2. 安全验证:
   - SQL 数据库: 使用 SQLAnalyzer 验证只读
   - Redis: 使用 READONLY_COMMANDS 白名单验证
   - MongoDB: 使用 READONLY_OPERATIONS 白名单验证
3. 结果脱敏: 集成 DataMasker 对敏感字段脱敏
4. 查询限制: 超时控制、行数限制

使用方式:
    result = await QueryService.execute_query(instance, database, query, user_id)
"""

import asyncio
import json
import re
import time
from abc import ABC, abstractmethod
from typing import Any, Dict, List, Optional, Tuple

from app.models.instance import Instance
from app.schemas.query import UnifiedQueryResult
from app.services.db_connector import DBConnector
from app.services.redis_connector import RedisConnector
from app.services.sql_analyzer import SQLAnalyzer

MAX_ROWS = 10000
QUERY_TIMEOUT = 30


# ===== Redis 只读命令白名单 =====

REDIS_READONLY_COMMANDS = {
    # 字符串操作
    "GET", "MGET", "STRLEN", "GETRANGE", "GETBIT",
    # 哈希操作
    "HGET", "HGETALL", "HKEYS", "HVALS", "HLEN", "HEXISTS", "HMGET",
    # 列表操作
    "LRANGE", "LLEN", "LINDEX",
    # 集合操作
    "SMEMBERS", "SCARD", "SISMEMBER", "SRANDMEMBER", "SUNION", "SINTER", "SDIFF",
    # 有序集合操作
    "ZRANGE", "ZREVRANGE", "ZRANGEBYSCORE", "ZREVRANGEBYSCORE",
    "ZCARD", "ZSCORE", "ZRANK", "ZREVRANK", "ZCOUNT", "ZLEXCOUNT",
    # 通用操作
    "EXISTS", "TTL", "PTTL", "TYPE", "DBSIZE",
    "SCAN", "HSCAN", "SSCAN", "ZSCAN", "KEYS",
    "INFO", "CONFIG", "MEMORY", "OBJECT",
    # 其他
    "PING", "ECHO", "TIME",
}

# 需要特殊处理的子命令: CONFIG GET / MEMORY USAGE / OBJECT ENCODING 等
REDIS_SUBCOMMAND_ALLOWLIST = {
    ("CONFIG", "GET"),
    ("MEMORY", "USAGE"),
    ("MEMORY", "STATS"),
    ("OBJECT", "ENCODING"),
    ("OBJECT", "IDLETIME"),
    ("OBJECT", "REFCOUNT"),
}

# ===== MongoDB 只读操作白名单 =====

MONGODB_READONLY_OPS = {
    "find", "findOne", "countDocuments", "estimatedDocumentCount",
    "distinct", "aggregate",
}

MONGODB_WRITE_OPS = {
    "insertOne", "insertMany", "updateOne", "updateMany",
    "replaceOne", "deleteOne", "deleteMany", "drop",
    "createIndex", "dropIndex", "renameCollection",
    "findOneAndUpdate", "findOneAndReplace", "findOneAndDelete",
    "bulkWrite",
}


# ===== 查询处理器基类 =====


class BaseQueryHandler(ABC):
    """查询处理器基类。

    所有数据库类型的查询处理器必须继承此基类，
    并实现 execute 和 validate_readonly 方法。
    """

    @abstractmethod
    async def execute(
        self,
        instance: Instance,
        database: Optional[str],
        query: str,
        user_id: int,
        options: Optional[Dict[str, Any]] = None,
    ) -> UnifiedQueryResult:
        """执行查询并返回统一格式的结果。

        Args:
            instance: 数据源实例配置
            database: 目标数据库名/Redis DB号/MongoDB collection
            query: 查询语句（SQL/Redis命令/MongoDB查询JSON）
            user_id: 查询用户ID
            options: 额外选项（如 limit, offset 等）

        Returns:
            UnifiedQueryResult: 统一格式的查询结果
        """
        pass

    @abstractmethod
    def validate_readonly(self, query: str) -> Tuple[bool, str]:
        """验证查询是否为只读。

        Args:
            query: 查询语句

        Returns:
            (is_safe, error_message): 是否安全及错误信息
        """
        pass


# ===== SQL 查询处理器（MySQL / PostgreSQL 共用基类） =====


class SQLQueryHandlerBase(BaseQueryHandler):
    """SQL 查询处理器基类（MySQL 和 PostgreSQL 共用逻辑）。"""

    _analyzer = SQLAnalyzer()

    def validate_readonly(self, query: str) -> Tuple[bool, str]:
        """使用 SQLAnalyzer 验证 SQL 是否为只读。

        Args:
            query: SQL 查询语句

        Returns:
            (is_safe, error_message)
        """
        is_safe, error_msg = self._analyzer.validate_readonly(query)
        return is_safe, error_msg

    def _auto_add_limit(self, sql: str, max_rows: int = MAX_ROWS) -> str:
        """如果 SQL 没有 LIMIT 则自动添加。

        Args:
            sql: 原始 SQL
            max_rows: 最大返回行数

        Returns:
            添加了 LIMIT 的 SQL
        """
        upper_sql = sql.strip().upper()
        if not re.search(r"LIMIT\s+\d+", upper_sql, re.IGNORECASE):
            sql = sql.rstrip(";") + f" LIMIT {max_rows};"
        return sql


class MySQLQueryHandler(SQLQueryHandlerBase):
    """MySQL 查询处理器。

    使用 DBConnector._query_mysql 执行查询，
    SQLAnalyzer 验证只读性，自动添加 LIMIT，30秒超时。
    """

    async def execute(
        self,
        instance: Instance,
        database: Optional[str],
        query: str,
        user_id: int,
        options: Optional[Dict[str, Any]] = None,
    ) -> UnifiedQueryResult:
        """执行 MySQL 只读查询。

        Args:
            instance: MySQL 实例配置
            database: 目标数据库名
            query: SQL 查询语句
            user_id: 查询用户ID
            options: 额外选项

        Returns:
            UnifiedQueryResult: 统一格式的查询结果

        Raises:
            ValueError: 查询不安全或执行失败
        """
        # 1. 验证只读
        is_safe, error_msg = self.validate_readonly(query)
        if not is_safe:
            raise ValueError(error_msg)

        # 2. 自动添加 LIMIT
        query = self._auto_add_limit(query)

        # 3. 执行查询
        start = time.time()
        try:
            rows, columns = await DBConnector._query_mysql(instance, database, query)
        except asyncio.TimeoutError:
            raise ValueError("查询超时（30秒）")

        elapsed = int((time.time() - start) * 1000)

        return UnifiedQueryResult(
            columns=columns,
            rows=rows if rows else [],
            raw_result=None,
            total=len(rows) if rows else 0,
            elapsed_ms=elapsed,
            truncated=(len(rows) if rows else 0) >= MAX_ROWS,
            db_type="mysql",
            query_type="SELECT",
        )


class PostgresQueryHandler(SQLQueryHandlerBase):
    """PostgreSQL 查询处理器。

    使用 DBConnector._query_postgres 执行查询，
    SQLAnalyzer 验证只读性（SQL 标准兼容），自动添加 LIMIT，30秒超时。
    """

    async def execute(
        self,
        instance: Instance,
        database: Optional[str],
        query: str,
        user_id: int,
        options: Optional[Dict[str, Any]] = None,
    ) -> UnifiedQueryResult:
        """执行 PostgreSQL 只读查询。

        Args:
            instance: PostgreSQL 实例配置
            database: 目标数据库名
            query: SQL 查询语句
            user_id: 查询用户ID
            options: 额外选项

        Returns:
            UnifiedQueryResult: 统一格式的查询结果

        Raises:
            ValueError: 查询不安全或执行失败
        """
        # 1. 验证只读
        is_safe, error_msg = self.validate_readonly(query)
        if not is_safe:
            raise ValueError(error_msg)

        # 2. 自动添加 LIMIT
        query = self._auto_add_limit(query)

        # 3. 执行查询
        start = time.time()
        try:
            rows, columns = await DBConnector._query_postgres(instance, database, query)
        except asyncio.TimeoutError:
            raise ValueError("查询超时（30秒）")

        elapsed = int((time.time() - start) * 1000)

        return UnifiedQueryResult(
            columns=columns,
            rows=rows if rows else [],
            raw_result=None,
            total=len(rows) if rows else 0,
            elapsed_ms=elapsed,
            truncated=(len(rows) if rows else 0) >= MAX_ROWS,
            db_type="postgres",
            query_type="SELECT",
        )


# ===== Redis 查询处理器 =====


class RedisQueryHandler(BaseQueryHandler):
    """Redis 查询处理器。

    使用 READONLY_COMMANDS 白名单验证命令安全性，
    RedisConnector.execute_command 执行命令，
    返回格式化为表格/JSON 的结果。
    """

    def validate_readonly(self, query: str) -> Tuple[bool, str]:
        """验证 Redis 命令是否在只读白名单中。

        Args:
            query: Redis 命令字符串，如 "GET mykey"

        Returns:
            (is_safe, error_message)
        """
        parts = query.strip().split()
        if not parts:
            return False, "Redis 命令不能为空"

        cmd = parts[0].upper()

        # 检查子命令（如 CONFIG GET, MEMORY USAGE）
        if len(parts) >= 2:
            sub = parts[1].upper()
            subcmd_key = (cmd, sub)
            # 如果是已知的多词命令，检查是否在允许列表中
            if cmd in ("CONFIG", "MEMORY", "OBJECT"):
                if subcmd_key not in REDIS_SUBCOMMAND_ALLOWLIST:
                    return False, f"Redis 命令 {cmd} {sub} 不在只读白名单中"
                return True, ""

        if cmd not in REDIS_READONLY_COMMANDS:
            return False, f"Redis 命令 '{cmd}' 不在只读白名单中，禁止执行"

        return True, ""

    async def execute(
        self,
        instance: Instance,
        database: Optional[str],
        query: str,
        user_id: int,
        options: Optional[Dict[str, Any]] = None,
    ) -> UnifiedQueryResult:
        """执行 Redis 只读命令。

        Args:
            instance: Redis 实例配置
            database: Redis DB 号（字符串，如 "0"）
            query: Redis 命令字符串
            user_id: 查询用户ID
            options: 额外选项

        Returns:
            UnifiedQueryResult: 统一格式的查询结果

        Raises:
            ValueError: 命令不安全或执行失败
        """
        # 1. 验证只读
        is_safe, error_msg = self.validate_readonly(query)
        if not is_safe:
            raise ValueError(error_msg)

        # 2. 解析 DB 号
        try:
            db = int(database) if database else 0
        except (ValueError, TypeError):
            db = 0

        # 3. 执行命令
        start = time.time()
        try:
            result_str = await RedisConnector.execute_command(instance, db, query)
        except Exception as e:
            raise ValueError(f"Redis 命令执行失败: {str(e)}")

        elapsed = int((time.time() - start) * 1000)

        # 4. 格式化结果为表格形式
        parts = query.strip().split()
        cmd = parts[0].upper() if parts else ""

        columns, rows = self._format_result(cmd, result_str)

        return UnifiedQueryResult(
            columns=columns,
            rows=rows,
            raw_result=result_str,
            total=len(rows),
            elapsed_ms=elapsed,
            truncated=False,
            db_type="redis",
            query_type=cmd,
        )

    def _format_result(self, cmd: str, result: str) -> Tuple[List[str], List[Dict[str, Any]]]:
        """将 Redis 原始结果格式化为表格。

        Args:
            cmd: Redis 命令名
            result: 原始结果字符串

        Returns:
            (columns, rows): 列名列表和行数据列表
        """
        if not result or result == "(nil)":
            return ["result"], [{"result": "(nil)"}]

        lines = result.split("\n")

        # 对于返回列表的命令（KEYS, HKEYS, HVALS, SMEMBERS 等）
        list_commands = {"KEYS", "HKEYS", "HVALS", "SMEMBERS", "LRANGE", "SINTER", "SUNION", "SDIFF"}
        if cmd in list_commands and len(lines) > 1:
            return ["value"], [{"value": line} for line in lines if line.strip()]

        # 对于 INFO 命令，返回 key-value 表格
        if cmd == "INFO":
            rows = []
            for line in lines:
                line = line.strip()
                if line and not line.startswith("#") and ":" in line:
                    k, v = line.split(":", 1)
                    rows.append({"key": k.strip(), "value": v.strip()})
            if rows:
                return ["key", "value"], rows

        # 对于 CONFIG GET，返回 key-value 表格
        if cmd == "CONFIG":
            rows = []
            for i in range(0, len(lines), 2):
                if i + 1 < len(lines):
                    rows.append({"key": lines[i].strip(), "value": lines[i + 1].strip()})
                elif lines[i].strip():
                    rows.append({"key": lines[i].strip(), "value": ""})
            if rows:
                return ["key", "value"], rows

        # 对于 HGETALL，返回 field-value 表格
        if cmd == "HGETALL":
            rows = []
            for i in range(0, len(lines), 2):
                if i + 1 < len(lines):
                    rows.append({"field": lines[i].strip(), "value": lines[i + 1].strip()})
            if rows:
                return ["field", "value"], rows

        # 默认: 返回单行结果
        return ["result"], [{"result": result}]


# ===== MongoDB 查询处理器 =====


class MongoQueryHandler(BaseQueryHandler):
    """MongoDB 查询处理器。

    使用 READONLY_OPERATIONS 白名单验证操作安全性，
    motor 异步驱动执行查询，
    查询格式为 JSON 字符串。
    """

    def validate_readonly(self, query: str) -> Tuple[bool, str]:
        """验证 MongoDB 操作是否在只读白名单中。

        Args:
            query: MongoDB 查询 JSON 字符串

        Returns:
            (is_safe, error_message)
        """
        try:
            query_obj = json.loads(query)
        except json.JSONDecodeError as e:
            return False, f"MongoDB 查询 JSON 格式错误: {str(e)}"

        operation = query_obj.get("operation", "")
        if not operation:
            return False, "MongoDB 查询必须指定 'operation' 字段"

        if operation in MONGODB_WRITE_OPS:
            return False, f"MongoDB 操作 '{operation}' 是写操作，禁止执行"

        if operation not in MONGODB_READONLY_OPS:
            return False, f"MongoDB 操作 '{operation}' 不在允许的操作列表中"

        # 对 aggregate 进行额外检查: 禁止 $out 和 $merge 阶段
        if operation == "aggregate":
            pipeline = query_obj.get("pipeline", [])
            if not isinstance(pipeline, list):
                return False, "aggregate 的 pipeline 必须是数组"
            for stage in pipeline:
                if isinstance(stage, dict):
                    stage_name = list(stage.keys())[0] if stage else ""
                    if stage_name in ("$out", "$merge"):
                        return False, f"aggregate 管道中的 '{stage_name}' 阶段是写操作，禁止执行"

        return True, ""

    async def execute(
        self,
        instance: Instance,
        database: Optional[str],
        query: str,
        user_id: int,
        options: Optional[Dict[str, Any]] = None,
    ) -> UnifiedQueryResult:
        """执行 MongoDB 只读查询。

        Args:
            instance: MongoDB 实例配置
            database: MongoDB 数据库名
            query: MongoDB 查询 JSON 字符串，格式如:
                {"collection": "users", "operation": "find",
                 "filter": {"age": {"$gt": 18}},
                 "projection": {"name": 1, "age": 1}}
            user_id: 查询用户ID
            options: 额外选项

        Returns:
            UnifiedQueryResult: 统一格式的查询结果

        Raises:
            ValueError: 操作不安全或执行失败
        """
        # 1. 验证只读
        is_safe, error_msg = self.validate_readonly(query)
        if not is_safe:
            raise ValueError(error_msg)

        # 2. 解析查询 JSON
        try:
            query_obj = json.loads(query)
        except json.JSONDecodeError as e:
            raise ValueError(f"MongoDB 查询 JSON 格式错误: {str(e)}")

        collection_name = query_obj.get("collection", database or "")
        operation = query_obj.get("operation", "")
        filter_doc = query_obj.get("filter", {})
        projection = query_obj.get("projection", None)
        limit_val = query_obj.get("limit", MAX_ROWS)
        skip_val = query_obj.get("skip", 0)
        sort_doc = query_obj.get("sort", None)
        pipeline = query_obj.get("pipeline", [])

        if not collection_name:
            raise ValueError("MongoDB 查询必须指定 'collection' 字段或在 database 参数中提供")

        # 3. 导入 MongoConnector（延迟导入避免循环依赖）
        from app.services.mongo_connector import MongoConnector

        start = time.time()

        try:
            if operation == "find":
                docs, total = await MongoConnector.execute_find(
                    instance, database or instance.database or "",
                    collection_name, filter_doc, projection,
                    limit_val, skip_val, sort_doc,
                )
                columns = list(docs[0].keys()) if docs else []
                # 将 _id 转换为字符串以便 JSON 序列化
                for doc in docs:
                    if "_id" in doc:
                        doc["_id"] = str(doc["_id"])
                rows = docs

            elif operation == "findOne":
                doc = await MongoConnector.execute_find_one(
                    instance, database or instance.database or "",
                    collection_name, filter_doc, projection,
                )
                if doc:
                    if "_id" in doc:
                        doc["_id"] = str(doc["_id"])
                    columns = list(doc.keys())
                    rows = [doc]
                    total = 1
                else:
                    columns, rows, total = [], [], 0

            elif operation == "countDocuments":
                count = await MongoConnector.execute_count(
                    instance, database or instance.database or "",
                    collection_name, filter_doc,
                )
                columns = ["count"]
                rows = [{"count": count}]
                total = 1

            elif operation == "estimatedDocumentCount":
                count = await MongoConnector.execute_estimated_count(
                    instance, database or instance.database or "",
                    collection_name,
                )
                columns = ["count"]
                rows = [{"count": count}]
                total = 1

            elif operation == "distinct":
                field = query_obj.get("field", "")
                if not field:
                    raise ValueError("distinct 操作必须指定 'field' 字段")
                values = await MongoConnector.execute_distinct(
                    instance, database or instance.database or "",
                    collection_name, field, filter_doc,
                )
                columns = ["value"]
                rows = [{"value": str(v)} for v in values]
                total = len(values)

            elif operation == "aggregate":
                docs = await MongoConnector.execute_aggregate(
                    instance, database or instance.database or "",
                    collection_name, pipeline,
                )
                columns = list(docs[0].keys()) if docs else []
                for doc in docs:
                    if "_id" in doc:
                        doc["_id"] = str(doc["_id"])
                rows = docs
                total = len(docs)

            else:
                raise ValueError(f"不支持的 MongoDB 操作: {operation}")

        except ValueError:
            raise
        except asyncio.TimeoutError:
            raise ValueError("查询超时（30秒）")
        except Exception as e:
            raise ValueError(f"MongoDB 查询执行失败: {str(e)}")

        elapsed = int((time.time() - start) * 1000)

        return UnifiedQueryResult(
            columns=columns,
            rows=rows,
            raw_result={"operation": operation, "collection": collection_name},
            total=total,
            elapsed_ms=elapsed,
            truncated=total >= MAX_ROWS,
            db_type="mongo",
            query_type=operation,
        )


# ===== 统一查询服务入口 =====


class QueryService:
    """统一多数据库查询服务入口。

    根据实例的 db_type 自动路由到对应的查询处理器，
    支持 MySQL、PostgreSQL、Redis、MongoDB 四种数据库。

    Examples:
        >>> result = await QueryService.execute_query(instance, database, query, user_id)
        >>> print(result.rows)
    """

    _handlers: Dict[str, BaseQueryHandler] = {
        "mysql": MySQLQueryHandler(),
        "postgres": PostgresQueryHandler(),
        "postgresql": PostgresQueryHandler(),
        "redis": RedisQueryHandler(),
        "mongo": MongoQueryHandler(),
    }

    @staticmethod
    async def execute_query(
        instance: Instance,
        database: Optional[str],
        query: str,
        user_id: int,
        options: Optional[Dict[str, Any]] = None,
    ) -> UnifiedQueryResult:
        """统一查询入口。

        根据 instance.db_type 自动路由到对应查询处理器，
        执行安全验证后返回统一格式的结果。

        Args:
            instance: 数据源实例配置
            database: 目标数据库/Redis DB号/MongoDB collection
            query: 查询语句（SQL/Redis命令/MongoDB查询JSON）
            user_id: 查询用户ID
            options: 额外选项（如 limit, offset 等）

        Returns:
            UnifiedQueryResult: 统一格式的查询结果

        Raises:
            ValueError: 不支持的数据库类型或查询不安全
        """
        db_type = instance.db_type.lower()
        handler = QueryService._handlers.get(db_type)

        if handler is None:
            raise ValueError(f"暂不支持 {instance.db_type} 查询")

        return await handler.execute(instance, database, query, user_id, options)

    @staticmethod
    async def execute_dml(
        instance: Instance, database: str, sql: str
    ) -> str:
        """执行 DML 语句（INSERT/UPDATE/DELETE）。

        目前仅支持 SQL 数据库（MySQL、PostgreSQL）。
        Redis 和 MongoDB 的变更操作通过各自连接器处理。

        Args:
            instance: 数据源实例配置
            database: 目标数据库名
            sql: DML SQL 语句

        Returns:
            str: 执行结果描述

        Raises:
            ValueError: 不支持的数据库类型或执行失败
        """
        return await DBConnector.execute_dml(instance, database, sql)