Add tests

snuba/admin/production_queries/prod_queries.py

@@ -29,13 +29,13 @@ def run_query_with_audit(query: str, user: str) -> Response:
             return response
 
         body["dry_run"] = False
-        validate_projects_in_query(body, dataset)
+        _validate_projects_in_query(body, dataset)
         return dataset_query(dataset, body, Timer("admin"))
 
     return run_query_with_audit(body["query"], user)
 
 
-def validate_projects_in_query(body: Dict[str, Any], dataset: Dataset) -> None:
+def _validate_projects_in_query(body: Dict[str, Any], dataset: Dataset) -> None:
     request_parts = RequestSchema.build(HTTPQuerySettings).validate(body)
     query = parse_snql_query(request_parts.query["query"], dataset)[0]
     project_ids = get_object_ids_in_query_ast(query, "project_id")

snuba/settings/settings_test.py

@@ -37,6 +37,8 @@
 
 OPTIMIZE_PARALLEL_MAX_JITTER_MINUTES = 0
 
+ADMIN_ALLOWED_PROD_PROJECTS = [1]
+
 REDIS_CLUSTERS = {
     key: {
         "use_redis_cluster": os.environ.get("USE_REDIS_CLUSTER", "0") != "0",

tests/admin/test_api.py

@@ -490,3 +490,43 @@ def mock_record(user: Any, action: Any, data: Any, notify: Any) -> None:
         } not in response.json[0]["configs"]
         # make sure an auditlog entry was recorded
         assert auditlog_records.pop()
+
+
+def test_prod_snql_query_invalid_dataset(admin_api: FlaskClient) -> None:
+    response = admin_api.post(
+        "/production_snql_query", data=json.dumps({"dataset": "", "query": ""})
+    )
+    assert response.status_code == 400
+    data = json.loads(response.data)
+    assert data["error"]["message"] == "dataset '' does not exist"
+
+
+@pytest.mark.redis_db
+def test_prod_snql_query_invalid_invalid_query(admin_api: FlaskClient) -> None:
+    response = admin_api.post(
+        "/production_snql_query", data=json.dumps({"dataset": "sessions", "query": ""})
+    )
+    assert response.status_code == 400
+    data = json.loads(response.data)
+    assert (
+        data["error"]["message"]
+        == "Rule 'query_exp' didn't match at '' (line 1, column 1)."
+    )
+
+
+@pytest.mark.redis_db
+@pytest.mark.clickhouse_db
+def test_prod_snql_query_valid_query(admin_api: FlaskClient) -> None:
+    snql_query = """
+    MATCH (events)
+    SELECT title
+    WHERE project_id = 1
+    AND timestamp >= toDateTime('2023-01-01 00:00:00')
+    AND timestamp < toDateTime('2023-02-01 00:00:00')
+    """
+    response = admin_api.post(
+        "/snql_to_sql", data=json.dumps({"dataset": "events", "query": snql_query})
+    )
+    assert response.status_code == 200
+    data = json.loads(response.data)
+    assert "data" in data