diff --git a/snuba/admin/views.py b/snuba/admin/views.py index b891c04baa..c469b535ea 100644 --- a/snuba/admin/views.py +++ b/snuba/admin/views.py @@ -80,6 +80,8 @@ runner = Runner() audit_log = AuditLog() +ORG_ID = 1 + @application.errorhandler(UnauthorizedException) def handle_invalid_json(exception: UnauthorizedException) -> Response: @@ -237,9 +239,11 @@ def do_action() -> None: if not dry_run: audit_log.record( user or "", - AuditLogAction.RAN_MIGRATION_STARTED - if action == "run" - else AuditLogAction.REVERSED_MIGRATION_STARTED, + ( + AuditLogAction.RAN_MIGRATION_STARTED + if action == "run" + else AuditLogAction.REVERSED_MIGRATION_STARTED + ), {"migration": str(migration_key), "force": force, "fake": fake}, ) check_for_inactive_replicas( @@ -256,9 +260,11 @@ def do_action() -> None: if not dry_run: audit_log.record( user or "", - AuditLogAction.RAN_MIGRATION_COMPLETED - if action == "run" - else AuditLogAction.REVERSED_MIGRATION_COMPLETED, + ( + AuditLogAction.RAN_MIGRATION_COMPLETED + if action == "run" + else AuditLogAction.REVERSED_MIGRATION_COMPLETED + ), {"migration": str(migration_key), "force": force, "fake": fake}, notify=True, ) @@ -266,9 +272,11 @@ def do_action() -> None: def notify_error() -> None: audit_log.record( user or "", - AuditLogAction.RAN_MIGRATION_FAILED - if action == "run" - else AuditLogAction.REVERSED_MIGRATION_FAILED, + ( + AuditLogAction.RAN_MIGRATION_FAILED + if action == "run" + else AuditLogAction.REVERSED_MIGRATION_FAILED + ), {"migration": str(migration_key), "force": force, "fake": fake}, notify=True, ) @@ -667,7 +675,7 @@ def config(config_key: str) -> Response: 400, {"Content-Type": "application/json"}, ) - except (state.MismatchedTypeException): + except state.MismatchedTypeException: return Response( json.dumps({"error": "Mismatched type"}), 400, @@ -997,7 +1005,7 @@ def dlq_replay() -> Response: @check_tool_perms(tools=[AdminTools.PRODUCTION_QUERIES]) def production_snql_query() -> Response: body = json.loads(request.data) - body["tenant_ids"] = {"referrer": request.referrer} + body["tenant_ids"] = {"referrer": request.referrer, "organization_id": ORG_ID} try: ret = run_snql_query(body, g.user.email) return ret diff --git a/snuba/settings/settings_test.py b/snuba/settings/settings_test.py index 01673da983..c59bb3b37e 100644 --- a/snuba/settings/settings_test.py +++ b/snuba/settings/settings_test.py @@ -44,7 +44,7 @@ OPTIMIZE_PARALLEL_MAX_JITTER_MINUTES = 0 -ADMIN_ALLOWED_PROD_PROJECTS = [1] +ADMIN_ALLOWED_PROD_PROJECTS = [1, 11276] REDIS_CLUSTERS = { key: { diff --git a/tests/admin/test_api.py b/tests/admin/test_api.py index aac9582f5b..d465e9d7e2 100644 --- a/tests/admin/test_api.py +++ b/tests/admin/test_api.py @@ -599,6 +599,26 @@ def test_prod_snql_query_valid_query(admin_api: FlaskClient) -> None: assert "data" in data +@pytest.mark.redis_db +@pytest.mark.clickhouse_db +def test_prod_snql_query_multiple_allowed_projects(admin_api: FlaskClient) -> None: + snql_query = """ + MATCH (transactions) + SELECT title + WHERE project_id IN array(1, 11276) + AND finish_ts >= toDateTime('2023-01-01 00:00:00') + AND finish_ts < toDateTime('2023-02-01 00:00:00') + """ + response = admin_api.post( + "/production_snql_query", + data=json.dumps({"dataset": "transactions", "query": snql_query}), + headers={"Referer": "https://snuba-admin.getsentry.net/"}, + ) + assert response.status_code == 200 + data = json.loads(response.data) + assert "data" in data + + @pytest.mark.redis_db @pytest.mark.clickhouse_db def test_prod_snql_query_invalid_project_query(admin_api: FlaskClient) -> None: