diff --git a/backend/tests/test_query.py b/backend/tests/test_query.py index 15f4044..10c281d 100644 --- a/backend/tests/test_query.py +++ b/backend/tests/test_query.py @@ -1,11 +1,19 @@ import sys +from contextlib import asynccontextmanager from unittest.mock import patch, AsyncMock, MagicMock +import pytest from httpx import AsyncClient +from app.core.crypto import encrypt_password +from app.models.instance_group import InstanceGroup, user_groups from tests.conftest import auth_header +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + def _make_aiomysql_mock(): aiomysql = MagicMock() aiomysql.connect = AsyncMock() @@ -13,58 +21,512 @@ def _make_aiomysql_mock(): return aiomysql +def _build_mock_query_result(): + """Return a mock UnifiedQueryResult for SELECT queries.""" + from app.schemas.query import UnifiedQueryResult + return UnifiedQueryResult( + columns=["id", "name"], + rows=[{"id": 1, "name": "test"}], + total=1, + elapsed_ms=10, + truncated=False, + db_type="mysql", + query_type="SELECT", + ) + + +async def _create_grouped_instance(db, group_id, **overrides): + """Create an instance bound to the given group.""" + from app.models.instance import Instance + defaults = dict( + name="GroupedMySQL", + db_type="mysql", + host="127.0.0.1", + port=3306, + username="root", + password=encrypt_password("test123"), + database="test_db", + env="development", + status="disconnected", + group_id=group_id, + ) + defaults.update(overrides) + inst = Instance(**defaults) + db.add(inst) + await db.commit() + await db.refresh(inst) + return inst + + +# =================================================================== +# 1. POST /query — execute_query +# =================================================================== + class TestExecuteQuery: - async def test_select_query(self, client: AsyncClient, test_instance, admin_token): - from contextlib import asynccontextmanager + """Tests for POST /api/v1/query""" - mock_cursor = AsyncMock() - mock_cursor.description = [("id",), ("name",)] - mock_cursor.fetchall.return_value = [{"id": 1, "name": "test"}] - mock_cursor.execute = AsyncMock() + # ---- happy path ------------------------------------------------ - @asynccontextmanager - async def mock_cursor_ctx(*args, **kwargs): - yield mock_cursor + async def test_select_query_success( + self, client: AsyncClient, test_instance, admin_token + ): + mock_result = _build_mock_query_result() - mock_conn = AsyncMock() - mock_conn.cursor = mock_cursor_ctx - mock_conn.close = MagicMock() + with patch( + "app.api.v1.endpoints.query.QueryService.execute_query", + new_callable=AsyncMock, + return_value=mock_result, + ): + resp = await client.post( + "/api/v1/query", + headers=auth_header(admin_token), + json={ + "instance_id": test_instance.id, + "database": "test_db", + "sql": "SELECT * FROM users LIMIT 10", + }, + ) + assert resp.status_code == 200 + data = resp.json() + assert "columns" in data + assert "rows" in data + assert data["total"] == 1 + assert data["query_type"] == "SELECT" - mock_aiomysql = _make_aiomysql_mock() - mock_aiomysql.connect = AsyncMock(return_value=mock_conn) - mock_aiomysql.DictCursor = MagicMock() + # ---- forbidden SQL (DDL / DML) --------------------------------- - with patch.dict(sys.modules, {"aiomysql": mock_aiomysql}): - resp = await client.post("/api/v1/query", headers=auth_header(admin_token), json={ + async def test_drop_rejected( + self, client: AsyncClient, test_instance, admin_token + ): + """DDL like DROP should be rejected by the query service (ValueError → 400).""" + with patch( + "app.api.v1.endpoints.query.QueryService.execute_query", + new_callable=AsyncMock, + side_effect=ValueError("不允许执行非SELECT语句"), + ): + resp = await client.post( + "/api/v1/query", + headers=auth_header(admin_token), + json={ + "instance_id": test_instance.id, + "database": "test_db", + "sql": "DROP TABLE users", + }, + ) + assert resp.status_code == 400 + + async def test_delete_rejected( + self, client: AsyncClient, test_instance, admin_token + ): + with patch( + "app.api.v1.endpoints.query.QueryService.execute_query", + new_callable=AsyncMock, + side_effect=ValueError("不允许执行非SELECT语句"), + ): + resp = await client.post( + "/api/v1/query", + headers=auth_header(admin_token), + json={ + "instance_id": test_instance.id, + "database": "test_db", + "sql": "DELETE FROM users WHERE 1=1", + }, + ) + assert resp.status_code == 400 + + # ---- guest blocked --------------------------------------------- + + async def test_guest_cannot_query( + self, client: AsyncClient, test_instance, guest_token + ): + resp = await client.post( + "/api/v1/query", + headers=auth_header(guest_token), + json={ "instance_id": test_instance.id, "database": "test_db", - "sql": "SELECT * FROM users LIMIT 10", - }) - assert resp.status_code == 200 - data = resp.json() - assert "columns" in data - assert "rows" in data - - async def test_drop_rejected(self, client: AsyncClient, test_instance, admin_token): - resp = await client.post("/api/v1/query", headers=auth_header(admin_token), json={ - "instance_id": test_instance.id, - "database": "test_db", - "sql": "DROP TABLE users", - }) - assert resp.status_code == 400 - - async def test_delete_rejected(self, client: AsyncClient, test_instance, admin_token): - resp = await client.post("/api/v1/query", headers=auth_header(admin_token), json={ - "instance_id": test_instance.id, - "database": "test_db", - "sql": "DELETE FROM users WHERE 1=1", - }) - assert resp.status_code == 400 - - async def test_guest_cannot_query(self, client: AsyncClient, test_instance, guest_token): - resp = await client.post("/api/v1/query", headers=auth_header(guest_token), json={ - "instance_id": test_instance.id, - "database": "test_db", - "sql": "SELECT 1", - }) + "sql": "SELECT 1", + }, + ) assert resp.status_code == 403 + + # ---- instance not found ---------------------------------------- + + async def test_nonexistent_instance(self, client: AsyncClient, admin_token): + resp = await client.post( + "/api/v1/query", + headers=auth_header(admin_token), + json={ + "instance_id": 99999, + "database": "test_db", + "sql": "SELECT 1", + }, + ) + assert resp.status_code == 404 + + # ---- missing sql parameter (validation) ------------------------ + + async def test_missing_sql_parameter( + self, client: AsyncClient, test_instance, admin_token + ): + resp = await client.post( + "/api/v1/query", + headers=auth_header(admin_token), + json={ + "instance_id": test_instance.id, + "database": "test_db", + # no "sql" field + }, + ) + assert resp.status_code == 422 # Pydantic validation error + + # ---- empty sql string ------------------------------------------ + + async def test_empty_sql_string( + self, client: AsyncClient, test_instance, admin_token + ): + resp = await client.post( + "/api/v1/query", + headers=auth_header(admin_token), + json={ + "instance_id": test_instance.id, + "database": "test_db", + "sql": "", + }, + ) + assert resp.status_code == 422 + + # ---- internal server error from query service ------------------ + + async def test_query_service_exception( + self, client: AsyncClient, test_instance, admin_token + ): + with patch( + "app.api.v1.endpoints.query.QueryService.execute_query", + new_callable=AsyncMock, + side_effect=RuntimeError("connection lost"), + ): + resp = await client.post( + "/api/v1/query", + headers=auth_header(admin_token), + json={ + "instance_id": test_instance.id, + "database": "test_db", + "sql": "SELECT 1", + }, + ) + assert resp.status_code == 500 + + # ---- mask=false query param ------------------------------------ + + async def test_select_query_no_mask( + self, client: AsyncClient, test_instance, admin_token + ): + mock_result = _build_mock_query_result() + + with patch( + "app.api.v1.endpoints.query.QueryService.execute_query", + new_callable=AsyncMock, + return_value=mock_result, + ): + resp = await client.post( + "/api/v1/query?mask=false", + headers=auth_header(admin_token), + json={ + "instance_id": test_instance.id, + "database": "test_db", + "sql": "SELECT * FROM users LIMIT 10", + }, + ) + assert resp.status_code == 200 + + # ---- save_history=false query param ---------------------------- + + async def test_select_query_no_history( + self, client: AsyncClient, test_instance, admin_token + ): + mock_result = _build_mock_query_result() + + with patch( + "app.api.v1.endpoints.query.QueryService.execute_query", + new_callable=AsyncMock, + return_value=mock_result, + ), patch( + "app.api.v1.endpoints.query.QueryHistoryService.record_query", + new_callable=AsyncMock, + ) as mock_hist: + resp = await client.post( + "/api/v1/query?save_history=false", + headers=auth_header(admin_token), + json={ + "instance_id": test_instance.id, + "database": "test_db", + "sql": "SELECT 1", + }, + ) + assert resp.status_code == 200 + # history should NOT be recorded + mock_hist.assert_not_called() + + # ---- unauthenticated request ----------------------------------- + + async def test_unauthenticated_request(self, client: AsyncClient, test_instance): + resp = await client.post( + "/api/v1/query", + json={ + "instance_id": test_instance.id, + "database": "test_db", + "sql": "SELECT 1", + }, + ) + assert resp.status_code == 401 + + # ---- developer can query public instance ----------------------- + + async def test_developer_can_query_public_instance( + self, client: AsyncClient, test_instance, developer_token + ): + """Public instance (no group_id) is accessible to developer.""" + mock_result = _build_mock_query_result() + + with patch( + "app.api.v1.endpoints.query.QueryService.execute_query", + new_callable=AsyncMock, + return_value=mock_result, + ): + resp = await client.post( + "/api/v1/query", + headers=auth_header(developer_token), + json={ + "instance_id": test_instance.id, + "database": "test_db", + "sql": "SELECT 1", + }, + ) + assert resp.status_code == 200 + + # ---- developer blocked from grouped instance without membership -- + + async def test_developer_blocked_from_grouped_instance( + self, client: AsyncClient, db, developer_token + ): + """Instance in a group where developer is not a member → 403.""" + group = InstanceGroup(name="private-group", db_type="mysql") + db.add(group) + await db.commit() + await db.refresh(group) + + inst = await _create_grouped_instance(db, group.id) + + resp = await client.post( + "/api/v1/query", + headers=auth_header(developer_token), + json={ + "instance_id": inst.id, + "database": "test_db", + "sql": "SELECT 1", + }, + ) + assert resp.status_code == 403 + + # ---- developer with group membership can query ----------------- + + async def test_developer_with_group_membership_can_query( + self, client: AsyncClient, db, developer_user, developer_token + ): + """Developer belongs to the instance group → 200.""" + group = InstanceGroup( + name="dev-group", + db_type="mysql", + permissions='["query:select"]', + ) + db.add(group) + await db.commit() + await db.refresh(group) + + await db.execute( + user_groups.insert().values( + user_id=developer_user.id, + group_id=group.id, + permissions='["query:select"]', + ) + ) + await db.commit() + + inst = await _create_grouped_instance(db, group.id) + + mock_result = _build_mock_query_result() + with patch( + "app.api.v1.endpoints.query.QueryService.execute_query", + new_callable=AsyncMock, + return_value=mock_result, + ): + resp = await client.post( + "/api/v1/query", + headers=auth_header(developer_token), + json={ + "instance_id": inst.id, + "database": "test_db", + "sql": "SELECT 1", + }, + ) + assert resp.status_code == 200 + + # ---- permission check: missing query:select → 403 -------------- + + async def test_developer_without_query_select_permission( + self, client: AsyncClient, db, developer_user, developer_token + ): + """Developer in group but without query:select permission → 403.""" + group = InstanceGroup( + name="no-query-group", + db_type="mysql", + permissions='["slowlog:view"]', # no query:select + ) + db.add(group) + await db.commit() + await db.refresh(group) + + await db.execute( + user_groups.insert().values( + user_id=developer_user.id, + group_id=group.id, + permissions='["slowlog:view"]', + ) + ) + await db.commit() + + inst = await _create_grouped_instance(db, group.id) + + resp = await client.post( + "/api/v1/query", + headers=auth_header(developer_token), + json={ + "instance_id": inst.id, + "database": "test_db", + "sql": "SELECT 1", + }, + ) + assert resp.status_code == 403 + + # ---- DBA can always query (admin bypass) ----------------------- + + async def test_dba_can_query_any_instance( + self, client: AsyncClient, db, dba_token + ): + """DBA is not admin, but can query a public instance.""" + from app.models.instance import Instance + + inst = Instance( + name="dba-test", + db_type="mysql", + host="127.0.0.1", + port=3306, + username="root", + password=encrypt_password("pw"), + database="test_db", + env="development", + status="disconnected", + ) + db.add(inst) + await db.commit() + await db.refresh(inst) + + mock_result = _build_mock_query_result() + with patch( + "app.api.v1.endpoints.query.QueryService.execute_query", + new_callable=AsyncMock, + return_value=mock_result, + ): + resp = await client.post( + "/api/v1/query", + headers=auth_header(dba_token), + json={ + "instance_id": inst.id, + "database": "test_db", + "sql": "SELECT 1", + }, + ) + assert resp.status_code == 200 + + +# =================================================================== +# 2. GET /query/supported-types +# =================================================================== + +class TestSupportedTypes: + """Tests for GET /api/v1/query/supported-types""" + + async def test_supported_types_returns_200( + self, client: AsyncClient, admin_token + ): + resp = await client.get( + "/api/v1/query/supported-types", + headers=auth_header(admin_token), + ) + assert resp.status_code == 200 + data = resp.json() + assert "types" in data + types = data["types"] + assert len(types) >= 4 + type_names = [t["type"] for t in types] + assert "mysql" in type_names + assert "postgres" in type_names + assert "redis" in type_names + assert "mongo" in type_names + + async def test_supported_types_each_entry_has_required_fields( + self, client: AsyncClient, admin_token + ): + resp = await client.get( + "/api/v1/query/supported-types", + headers=auth_header(admin_token), + ) + data = resp.json() + for entry in data["types"]: + assert "type" in entry + assert "name" in entry + assert "query_format" in entry + assert "example" in entry + assert "security" in entry + assert "features" in entry + + async def test_supported_types_unauthenticated( + self, client: AsyncClient + ): + resp = await client.get("/api/v1/query/supported-types") + assert resp.status_code == 401 + + async def test_supported_types_developer( + self, client: AsyncClient, developer_token + ): + """Non-admin authenticated user can also view supported types.""" + resp = await client.get( + "/api/v1/query/supported-types", + headers=auth_header(developer_token), + ) + assert resp.status_code == 200 + + async def test_supported_types_redis_has_readonly_commands( + self, client: AsyncClient, admin_token + ): + resp = await client.get( + "/api/v1/query/supported-types", + headers=auth_header(admin_token), + ) + data = resp.json() + redis_entry = next(t for t in data["types"] if t["type"] == "redis") + assert "readonly_commands" in redis_entry + assert "GET" in redis_entry["readonly_commands"] + + async def test_supported_types_mongo_has_readonly_operations( + self, client: AsyncClient, admin_token + ): + resp = await client.get( + "/api/v1/query/supported-types", + headers=auth_header(admin_token), + ) + data = resp.json() + mongo_entry = next(t for t in data["types"] if t["type"] == "mongo") + assert "readonly_operations" in mongo_entry + assert "find" in mongo_entry["readonly_operations"]