opsxsql/backend/tests/test_query_service.py

"""多数据库查询服务测试 — 验证 QueryService 统一查询路由和各处理器。

覆盖:
1. MySQL 查询处理器: 只读验证、自动 LIMIT、结果格式
2. PostgreSQL 查询处理器: 只读验证、自动 LIMIT
3. Redis 查询处理器: 白名单验证、子命令验证、结果格式化
4. MongoDB 查询处理器: JSON 解析、只读操作白名单、aggregate 管道安全
5. QueryService 路由: 按 db_type 分发、不支持的类型拒绝
"""

import json
import pytest
from unittest.mock import MagicMock

from app.services.query_service import (
    QueryService,
    MySQLQueryHandler,
    PostgresQueryHandler,
    RedisQueryHandler,
    MongoQueryHandler,
    REDIS_READONLY_COMMANDS,
    REDIS_SUBCOMMAND_ALLOWLIST,
    MONGODB_READONLY_OPS,
    MONGODB_WRITE_OPS,
)
from app.schemas.query import UnifiedQueryResult


# ===== MySQL 查询处理器 =====


class TestMySQLQueryHandler:
    """MySQL 查询处理器单元测试。"""

    def setup_method(self):
        self.handler = MySQLQueryHandler()

    def test_validate_readonly_select_pass(self):
        """SELECT 语句应通过只读验证。"""
        is_safe, msg = self.handler.validate_readonly("SELECT * FROM users LIMIT 10")
        assert is_safe is True
        assert msg == ""

    def test_validate_readonly_drop_rejected(self):
        """DROP 语句应被拒绝。"""
        is_safe, msg = self.handler.validate_readonly("DROP TABLE users")
        assert is_safe is False
        assert "DROP" in msg.upper() or "只读" in msg or "禁止" in msg

    def test_validate_readonly_insert_rejected(self):
        """INSERT 语句应被拒绝。"""
        is_safe, msg = self.handler.validate_readonly("INSERT INTO users (name) VALUES ('test')")
        assert is_safe is False

    def test_validate_readonly_update_rejected(self):
        """UPDATE 语句应被拒绝。"""
        is_safe, msg = self.handler.validate_readonly("UPDATE users SET name='x' WHERE id=1")
        assert is_safe is False

    def test_validate_readonly_delete_rejected(self):
        """DELETE 语句应被拒绝。"""
        is_safe, msg = self.handler.validate_readonly("DELETE FROM users WHERE id=1")
        assert is_safe is False

    def test_validate_readonly_alter_rejected(self):
        """ALTER 语句应被拒绝。"""
        is_safe, msg = self.handler.validate_readonly("ALTER TABLE users ADD COLUMN age INT")
        assert is_safe is False

    def test_auto_add_limit_no_existing(self):
        """没有 LIMIT 的 SQL 应自动添加 LIMIT。"""
        sql = "SELECT * FROM users"
        result = self.handler._auto_add_limit(sql)
        assert "LIMIT 10000" in result.upper()

    def test_auto_add_limit_already_has(self):
        """已有 LIMIT 的 SQL 不应重复添加。"""
        sql = "SELECT * FROM users LIMIT 5"
        result = self.handler._auto_add_limit(sql)
        assert "10000" not in result
        assert "LIMIT 5" in result.upper()

    @pytest.mark.asyncio
    async def test_execute_non_readonly_raises(self):
        """非只读查询应抛出 ValueError。"""
        instance = MagicMock()
        instance.db_type = "mysql"
        with pytest.raises(ValueError):
            await self.handler.execute(instance, "test_db", "DROP TABLE users", 1)


# ===== PostgreSQL 查询处理器 =====


class TestPostgresQueryHandler:
    """PostgreSQL 查询处理器单元测试。"""

    def setup_method(self):
        self.handler = PostgresQueryHandler()

    def test_validate_readonly_select_pass(self):
        """SELECT 语句应通过只读验证。"""
        is_safe, msg = self.handler.validate_readonly("SELECT * FROM users LIMIT 10")
        assert is_safe is True

    def test_validate_readonly_update_rejected(self):
        """UPDATE 语句应被拒绝。"""
        is_safe, msg = self.handler.validate_readonly("UPDATE users SET name='x' WHERE id=1")
        assert is_safe is False

    def test_validate_readonly_create_rejected(self):
        """CREATE 语句应被拒绝。"""
        is_safe, msg = self.handler.validate_readonly("CREATE TABLE test (id INT)")
        assert is_safe is False

    def test_auto_add_limit(self):
        """应自动添加 LIMIT。"""
        sql = "SELECT * FROM users"
        result = self.handler._auto_add_limit(sql)
        assert "LIMIT 10000" in result.upper()

    @pytest.mark.asyncio
    async def test_execute_non_readonly_raises(self):
        """非只读查询应抛出 ValueError。"""
        instance = MagicMock()
        instance.db_type = "postgres"
        with pytest.raises(ValueError):
            await self.handler.execute(instance, "test_db", "DELETE FROM users", 1)


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


class TestRedisQueryHandler:
    """Redis 查询处理器单元测试。"""

    def setup_method(self):
        self.handler = RedisQueryHandler()

    # --- 白名单验证 ---

    def test_validate_get_pass(self):
        """GET 命令应在白名单中。"""
        is_safe, msg = self.handler.validate_readonly("GET mykey")
        assert is_safe is True

    def test_validate_hgetall_pass(self):
        """HGETALL 命令应在白名单中。"""
        is_safe, msg = self.handler.validate_readonly("HGETALL myhash")
        assert is_safe is True

    def test_validate_keys_pass(self):
        """KEYS 命令应在白名单中。"""
        is_safe, msg = self.handler.validate_readonly("KEYS user:*")
        assert is_safe is True

    def test_validate_info_pass(self):
        """INFO 命令应在白名单中。"""
        is_safe, msg = self.handler.validate_readonly("INFO")
        assert is_safe is True

    def test_validate_scan_pass(self):
        """SCAN 命令应在白名单中。"""
        is_safe, msg = self.handler.validate_readonly("SCAN 0 MATCH user:* COUNT 100")
        assert is_safe is True

    def test_validate_ttl_pass(self):
        """TTL 命令应在白名单中。"""
        is_safe, msg = self.handler.validate_readonly("TTL mykey")
        assert is_safe is True

    def test_validate_exists_pass(self):
        """EXISTS 命令应在白名单中。"""
        is_safe, msg = self.handler.validate_readonly("EXISTS mykey")
        assert is_safe is True

    def test_validate_type_pass(self):
        """TYPE 命令应在白名单中。"""
        is_safe, msg = self.handler.validate_readonly("TYPE mykey")
        assert is_safe is True

    def test_validate_config_get_pass(self):
        """CONFIG GET 子命令应在白名单中。"""
        is_safe, msg = self.handler.validate_readonly("CONFIG GET maxmemory")
        assert is_safe is True

    def test_validate_memory_usage_pass(self):
        """MEMORY USAGE 子命令应在白名单中。"""
        is_safe, msg = self.handler.validate_readonly("MEMORY USAGE mykey")
        assert is_safe is True

    def test_validate_object_encoding_pass(self):
        """OBJECT ENCODING 子命令应在白名单中。"""
        is_safe, msg = self.handler.validate_readonly("OBJECT ENCODING mykey")
        assert is_safe is True

    # --- 写命令拒绝 ---

    def test_validate_set_rejected(self):
        """SET 命令应被拒绝（写操作）。"""
        is_safe, msg = self.handler.validate_readonly("SET mykey value")
        assert is_safe is False
        assert "SET" in msg

    def test_validate_del_rejected(self):
        """DEL 命令应被拒绝。"""
        is_safe, msg = self.handler.validate_readonly("DEL mykey")
        assert is_safe is False

    def test_validate_flushdb_rejected(self):
        """FLUSHDB 命令应被拒绝。"""
        is_safe, msg = self.handler.validate_readonly("FLUSHDB")
        assert is_safe is False

    def test_validate_flushall_rejected(self):
        """FLUSHALL 命令应被拒绝。"""
        is_safe, msg = self.handler.validate_readonly("FLUSHALL")
        assert is_safe is False

    def test_validate_hset_rejected(self):
        """HSET 命令应被拒绝。"""
        is_safe, msg = self.handler.validate_readonly("HSET myhash field value")
        assert is_safe is False

    def test_validate_lpush_rejected(self):
        """LPUSH 命令应被拒绝。"""
        is_safe, msg = self.handler.validate_readonly("LPUSH mylist value")
        assert is_safe is False

    def test_validate_sadd_rejected(self):
        """SADD 命令应被拒绝。"""
        is_safe, msg = self.handler.validate_readonly("SADD myset member")
        assert is_safe is False

    def test_validate_zadd_rejected(self):
        """ZADD 命令应被拒绝。"""
        is_safe, msg = self.handler.validate_readonly("ZADD myzset 1 member")
        assert is_safe is False

    def test_validate_config_set_rejected(self):
        """CONFIG SET 子命令应被拒绝。"""
        is_safe, msg = self.handler.validate_readonly("CONFIG SET maxmemory 100mb")
        assert is_safe is False

    def test_validate_empty_rejected(self):
        """空命令应被拒绝。"""
        is_safe, msg = self.handler.validate_readonly("")
        assert is_safe is False
        assert "空" in msg

    def test_validate_unknown_rejected(self):
        """未知命令应被拒绝。"""
        is_safe, msg = self.handler.validate_readonly("UNKNOWNCMD arg1")
        assert is_safe is False

    # --- 结果格式化 ---

    def test_format_result_nil(self):
        """(nil) 结果应返回单行。"""
        cols, rows = self.handler._format_result("GET", "(nil)")
        assert cols == ["result"]
        assert rows == [{"result": "(nil)"}]

    def test_format_result_simple_string(self):
        """简单字符串结果应返回单行。"""
        cols, rows = self.handler._format_result("GET", "hello")
        assert cols == ["result"]
        assert rows[0]["result"] == "hello"

    def test_format_result_info_keyvalue(self):
        """INFO 结果应解析为 key-value 表格。"""
        info_output = "# Server\nredis_version:7.0.0\ntcp_port:6379\nuptime_in_days:10"
        cols, rows = self.handler._format_result("INFO", info_output)
        assert cols == ["key", "value"]
        assert len(rows) == 3
        assert rows[0]["key"] == "redis_version"
        assert rows[0]["value"] == "7.0.0"

    def test_format_result_hgetall(self):
        """HGETALL 结果应解析为 field-value 表格。"""
        hgetall_output = "name\n" + "张三\n" + "age\n" + "25\n" + "email\n" + "test@example.com"
        cols, rows = self.handler._format_result("HGETALL", hgetall_output)
        assert cols == ["field", "value"]
        assert len(rows) == 3
        assert rows[0]["field"] == "name"
        assert rows[0]["value"] == "张三"

    def test_format_result_keys_list(self):
        """KEYS 结果应解析为 value 列表。"""
        keys_output = "user:1\nuser:2\nuser:3"
        cols, rows = self.handler._format_result("KEYS", keys_output)
        assert cols == ["value"]
        assert len(rows) == 3

    def test_format_result_empty(self):
        """空结果应返回 (nil)。"""
        cols, rows = self.handler._format_result("GET", "")
        assert cols == ["result"]
        assert rows == [{"result": "(nil)"}]

    @pytest.mark.asyncio
    async def test_execute_non_readonly_raises(self):
        """非只读命令应抛出 ValueError。"""
        instance = MagicMock()
        with pytest.raises(ValueError):
            await self.handler.execute(instance, "0", "SET mykey value", 1)


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


class TestMongoQueryHandler:
    """MongoDB 查询处理器单元测试。"""

    def setup_method(self):
        self.handler = MongoQueryHandler()

    # --- JSON 解析 ---

    def test_validate_invalid_json_rejected(self):
        """无效 JSON 应被拒绝。"""
        is_safe, msg = self.handler.validate_readonly("not json at all")
        assert is_safe is False
        assert "JSON" in msg

    def test_validate_missing_operation_rejected(self):
        """缺少 operation 字段应被拒绝。"""
        is_safe, msg = self.handler.validate_readonly(json.dumps({"collection": "users"}))
        assert is_safe is False
        assert "operation" in msg

    # --- 只读操作白名单 ---

    def test_validate_find_pass(self):
        """find 操作应通过验证。"""
        query = json.dumps({"collection": "users", "operation": "find", "filter": {}})
        is_safe, msg = self.handler.validate_readonly(query)
        assert is_safe is True

    def test_validate_findone_pass(self):
        """findOne 操作应通过验证。"""
        query = json.dumps({"collection": "users", "operation": "findOne", "filter": {"_id": 1}})
        is_safe, msg = self.handler.validate_readonly(query)
        assert is_safe is True

    def test_validate_countdocuments_pass(self):
        """countDocuments 操作应通过验证。"""
        query = json.dumps({"collection": "users", "operation": "countDocuments", "filter": {}})
        is_safe, msg = self.handler.validate_readonly(query)
        assert is_safe is True

    def test_validate_estimatedcount_pass(self):
        """estimatedDocumentCount 操作应通过验证。"""
        query = json.dumps({"collection": "users", "operation": "estimatedDocumentCount"})
        is_safe, msg = self.handler.validate_readonly(query)
        assert is_safe is True

    def test_validate_distinct_pass(self):
        """distinct 操作应通过验证。"""
        query = json.dumps({"collection": "users", "operation": "distinct", "field": "status"})
        is_safe, msg = self.handler.validate_readonly(query)
        assert is_safe is True

    def test_validate_aggregate_pass(self):
        """安全的 aggregate 管道应通过验证。"""
        query = json.dumps({
            "collection": "users",
            "operation": "aggregate",
            "pipeline": [{"$match": {"status": "active"}}, {"$group": {"_id": "$city", "count": {"$sum": 1}}}]
        })
        is_safe, msg = self.handler.validate_readonly(query)
        assert is_safe is True

    # --- 写操作拒绝 ---

    def test_validate_insertone_rejected(self):
        """insertOne 操作应被拒绝。"""
        query = json.dumps({"collection": "users", "operation": "insertOne", "document": {"name": "test"}})
        is_safe, msg = self.handler.validate_readonly(query)
        assert is_safe is False
        assert "写操作" in msg or "禁止" in msg

    def test_validate_updatemany_rejected(self):
        """updateMany 操作应被拒绝。"""
        query = json.dumps({"collection": "users", "operation": "updateMany", "filter": {}, "update": {"$set": {"x": 1}}})
        is_safe, msg = self.handler.validate_readonly(query)
        assert is_safe is False

    def test_validate_deleteone_rejected(self):
        """deleteOne 操作应被拒绝。"""
        query = json.dumps({"collection": "users", "operation": "deleteOne", "filter": {"_id": 1}})
        is_safe, msg = self.handler.validate_readonly(query)
        assert is_safe is False

    def test_validate_deletemany_rejected(self):
        """deleteMany 操作应被拒绝。"""
        query = json.dumps({"collection": "users", "operation": "deleteMany", "filter": {}})
        is_safe, msg = self.handler.validate_readonly(query)
        assert is_safe is False

    def test_validate_drop_rejected(self):
        """drop 操作应被拒绝。"""
        query = json.dumps({"collection": "users", "operation": "drop"})
        is_safe, msg = self.handler.validate_readonly(query)
        assert is_safe is False

    # --- aggregate 管道安全检查 ---

    def test_validate_aggregate_with_out_rejected(self):
        """aggregate 包含 $out 阶段应被拒绝。"""
        query = json.dumps({
            "collection": "users",
            "operation": "aggregate",
            "pipeline": [{"$match": {}}, {"$out": "backup_users"}]
        })
        is_safe, msg = self.handler.validate_readonly(query)
        assert is_safe is False
        assert "$out" in msg or "写操作" in msg

    def test_validate_aggregate_with_merge_rejected(self):
        """aggregate 包含 $merge 阶段应被拒绝。"""
        query = json.dumps({
            "collection": "users",
            "operation": "aggregate",
            "pipeline": [{"$match": {}}, {"$merge": {"into": "backup_users"}}]
        })
        is_safe, msg = self.handler.validate_readonly(query)
        assert is_safe is False
        assert "$merge" in msg or "写操作" in msg

    def test_validate_aggregate_pipeline_not_list_rejected(self):
        """aggregate 的 pipeline 不是数组应被拒绝。"""
        query = json.dumps({
            "collection": "users",
            "operation": "aggregate",
            "pipeline": "not a list"
        })
        is_safe, msg = self.handler.validate_readonly(query)
        assert is_safe is False
        assert "数组" in msg

    def test_validate_unknown_operation_rejected(self):
        """未知操作应被拒绝。"""
        query = json.dumps({"collection": "users", "operation": "mapReduce"})
        is_safe, msg = self.handler.validate_readonly(query)
        assert is_safe is False
        assert "允许" in msg or "不在" in msg

    @pytest.mark.asyncio
    async def test_execute_non_readonly_raises(self):
        """写操作应抛出 ValueError。"""
        instance = MagicMock()
        query = json.dumps({"collection": "users", "operation": "insertOne", "document": {"name": "test"}})
        with pytest.raises(ValueError):
            await self.handler.execute(instance, "test_db", query, 1)

    @pytest.mark.asyncio
    async def test_execute_invalid_json_raises(self):
        """无效 JSON 应抛出 ValueError。"""
        instance = MagicMock()
        with pytest.raises(ValueError):
            await self.handler.execute(instance, "test_db", "not json", 1)


# ===== QueryService 路由 =====


class TestQueryServiceRouter:
    """QueryService 统一路由测试。"""

    def test_handlers_include_all_db_types(self):
        """应支持所有 4 种数据库类型。"""
        handlers = QueryService._handlers
        assert "mysql" in handlers
        assert "postgres" in handlers
        assert "postgresql" in handlers
        assert "redis" in handlers
        assert "mongo" in handlers

    def test_postgres_and_postgresql_same_type(self):
        """postgres 和 postgresql 应使用同一类型的处理器。"""
        handlers = QueryService._handlers
        assert isinstance(handlers["postgres"], PostgresQueryHandler)
        assert isinstance(handlers["postgresql"], PostgresQueryHandler)

    @pytest.mark.asyncio
    async def test_unsupported_db_type_raises(self):
        """不支持的数据库类型应抛出 ValueError。"""
        instance = MagicMock()
        instance.db_type = "oracle"
        with pytest.raises(ValueError, match="暂不支持"):
            await QueryService.execute_query(instance, "test_db", "SELECT 1", 1)

    @pytest.mark.asyncio
    async def test_unsupported_db_type_sqlite(self):
        """SQLite 不应被支持（通过统一查询）。"""
        instance = MagicMock()
        instance.db_type = "sqlite"
        with pytest.raises(ValueError):
            await QueryService.execute_query(instance, "test_db", "SELECT 1", 1)


# ===== Redis 白名单完整性检查 =====


class TestRedisWhitelistCompleteness:
    """检查 Redis 只读白名单的完整性。"""

    def test_common_readonly_commands_present(self):
        """常用只读命令应在白名单中。"""
        required = {
            "GET", "MGET", "HGET", "HGETALL", "HKEYS", "HVALS", "HLEN",
            "LRANGE", "LLEN", "LINDEX",
            "SMEMBERS", "SCARD", "SISMEMBER",
            "ZRANGE", "ZCARD", "ZSCORE", "ZRANK",
            "EXISTS", "TTL", "PTTL", "TYPE", "DBSIZE",
            "SCAN", "HSCAN", "SSCAN", "ZSCAN", "KEYS",
            "INFO", "PING", "ECHO", "TIME",
        }
        missing = required - REDIS_READONLY_COMMANDS
        assert not missing, "缺少只读命令: " + str(missing)

    def test_write_commands_not_in_readonly(self):
        """写命令不应出现在只读白名单中。"""
        write_cmds = {
            "SET", "SETEX", "SETNX", "MSET", "DEL",
            "HSET", "HMSET", "HDEL", "HINCRBY",
            "LPUSH", "RPUSH", "LPOP", "RPOP", "LREM", "LSET",
            "SADD", "SREM", "SPOP",
            "ZADD", "ZREM", "ZINCRBY",
            "INCR", "INCRBY", "DECR", "DECRBY",
            "EXPIRE", "EXPIREAT", "PERSIST", "RENAME",
            "FLUSHDB", "FLUSHALL", "SAVE", "BGSAVE",
        }
        leaked = write_cmds & REDIS_READONLY_COMMANDS
        assert not leaked, "写命令不应在只读白名单中: " + str(leaked)

    def test_subcommand_allowlist_has_config_get(self):
        """CONFIG GET 子命令应在白名单中。"""
        assert ("CONFIG", "GET") in REDIS_SUBCOMMAND_ALLOWLIST

    def test_subcommand_allowlist_has_memory_usage(self):
        """MEMORY USAGE 子命令应在白名单中。"""
        assert ("MEMORY", "USAGE") in REDIS_SUBCOMMAND_ALLOWLIST


# ===== MongoDB 白名单完整性检查 =====


class TestMongoWhitelistCompleteness:
    """检查 MongoDB 只读操作白名单的完整性。"""

    def test_readonly_ops_present(self):
        """所有只读操作应在白名单中。"""
        required = {"find", "findOne", "countDocuments", "estimatedDocumentCount", "distinct", "aggregate"}
        assert required == MONGODB_READONLY_OPS

    def test_write_ops_present(self):
        """所有写操作应在 MONGODB_WRITE_OPS 中。"""
        required = {
            "insertOne", "insertMany", "updateOne", "updateMany",
            "deleteOne", "deleteMany", "drop", "createIndex", "dropIndex",
            "replaceOne", "renameCollection", "findOneAndUpdate",
            "findOneAndReplace", "findOneAndDelete", "bulkWrite",
        }
        assert required == MONGODB_WRITE_OPS

    def test_readonly_and_write_no_overlap(self):
        """只读和写操作不应有重叠。"""
        overlap = MONGODB_READONLY_OPS & MONGODB_WRITE_OPS
        assert not overlap, "只读和写操作不应重叠: " + str(overlap)