"""多数据库查询服务测试 — 验证 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() def test_auto_add_limit_show_no_limit(self): """SHOW 语句不应添加 LIMIT。""" sql = "SHOW CREATE TABLE users" result = self.handler._auto_add_limit(sql) assert "LIMIT" not in result.upper() assert result == sql def test_auto_add_limit_desc_no_limit(self): """DESC 语句不应添加 LIMIT。""" sql = "DESC users" result = self.handler._auto_add_limit(sql) assert "LIMIT" not 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)