Skip to content

Commit

Permalink
feat(capman): add project_id as implicit tenant (#4474)
Browse files Browse the repository at this point in the history
In order to port the project based rate limiters to the capman framework, we have to have project ids as tenants, we currently do not. Add logic in the build_request function to add project_id to AttributionInfo.tenant_ids

This logic makes an explicit decision that when there are multiple projects being searched in one request, there is no project tenant, instead the organization_id becomes the tenant.
  • Loading branch information
volokluev authored Jul 11, 2023
1 parent 6c43e8e commit 9c94e09
Show file tree
Hide file tree
Showing 2 changed files with 68 additions and 1 deletion.
9 changes: 8 additions & 1 deletion snuba/request/validation.py
Original file line number Diff line number Diff line change
Expand Up @@ -98,8 +98,10 @@ def build_request(
)

project_ids = get_object_ids_in_query_ast(query, "project_id")
query_project_id = None
if project_ids is not None and len(project_ids) == 1:
sentry_sdk.set_tag("snuba_project_id", project_ids.pop())
query_project_id = project_ids.pop()
sentry_sdk.set_tag("snuba_project_id", query_project_id)

org_ids = get_object_ids_in_query_ast(query, "org_id")
if org_ids is not None and len(org_ids) == 1:
Expand All @@ -113,6 +115,11 @@ def build_request(
attribution_info["tenant_ids"] = request_parts.attribution_info[
"tenant_ids"
]
if (
"project_id" not in attribution_info["tenant_ids"]
and query_project_id is not None
):
attribution_info["tenant_ids"]["project_id"] = query_project_id

request_id = uuid.uuid4().hex
request = Request(
Expand Down
60 changes: 60 additions & 0 deletions tests/request/test_build_request.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
from __future__ import annotations

from datetime import datetime
from typing import Any, Dict

Expand Down Expand Up @@ -35,6 +37,7 @@
"GRANULARITY 60"
),
"parent_api": "<unknown>",
"tenant_ids": {"organization_id": 1, "referrer": "test"},
},
binary_condition(
BooleanFunctions.AND,
Expand Down Expand Up @@ -97,3 +100,60 @@ def test_build_request(body: Dict[str, Any], condition: Expression) -> None:
assert dict(request.original_body) == body
status, differences = request.query.equals(expected_query)
assert status == True, f"Query mismatch: {differences}"


TENANT_ID_TESTS = [
pytest.param(
{
"query": (
"MATCH (events) "
"SELECT count() AS count BY time "
"WHERE "
"project_id IN tuple(1) AND "
"timestamp >= toDateTime('2011-07-01T19:54:15') AND"
"timestamp < toDateTime('2018-07-06T19:54:15') "
"LIMIT 1000 "
"GRANULARITY 60"
),
"tenant_ids": {"organization_id": 1, "referrer": "test"},
},
{"organization_id": 1, "referrer": "test", "project_id": 1},
id="one project id in query",
),
pytest.param(
{
"query": (
"MATCH (events) "
"SELECT count() AS count BY time "
"WHERE "
"project_id IN tuple(1, 2, 3, 4) AND"
"timestamp >= toDateTime('2011-07-01t19:54:15') AND"
"timestamp < toDateTime('2018-07-06t19:54:15') "
"LIMIT 1000 "
"GRANULARITY 60"
),
"tenant_ids": {"organization_id": 1, "referrer": "test"},
},
{"organization_id": 1, "referrer": "test"},
id="multiple projects, no project tenant",
),
]


@pytest.mark.parametrize("request_payload, expected_tenant_ids", TENANT_ID_TESTS)
def test_tenant_ids(
request_payload: dict[str, Any], expected_tenant_ids: dict[str, Any]
) -> None:
dataset = get_dataset("events")
schema = RequestSchema.build(HTTPQuerySettings)

request = build_request(
request_payload,
parse_snql_query,
HTTPQuerySettings,
schema,
dataset,
Timer("test"),
"my_request",
)
assert request.attribution_info.tenant_ids == expected_tenant_ids

0 comments on commit 9c94e09

Please sign in to comment.