chore(on-call): Default org_id in admin query tool (#5674)

* default prod query tool to use org id 1

* add test

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,19 +260,23 @@ 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,
             )
 
     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

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: {

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: