Skip to content

Commit e2015f2

Browse files
Rachel ChenRachel Chen
Rachel Chen
authored and
Rachel Chen
committed
tenant_ids
1 parent 0282a73 commit e2015f2

File tree

6 files changed

+14
-85
lines changed

6 files changed

+14
-85
lines changed

‎snuba/web/rpc/__init__.py‎

Lines changed: 0 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,6 @@
3131
RPCRequestException,
3232
convert_rpc_exception_to_proto,
3333
)
34-
from snuba.web.rpc.storage_routing.common import extract_message_meta
3534
from snuba.web.rpc.storage_routing.routing_strategies.storage_routing import (
3635
RoutingContext,
3736
RoutingDecision,
@@ -192,20 +191,10 @@ def execute(self, in_msg: Tin) -> Tout:
192191
span = scope.span
193192
if span is not None:
194193
span.description = self.config_key()
195-
request_meta = extract_message_meta(in_msg)
196194
self.routing_context = RoutingContext(
197195
timer=self._timer,
198196
in_msg=in_msg,
199197
query_id=uuid.uuid4().hex,
200-
tenant_ids={
201-
"organization_id": request_meta.organization_id,
202-
"referrer": request_meta.referrer,
203-
**(
204-
{"project_id": request_meta.project_ids[0]}
205-
if hasattr(request_meta, "project_ids") and len(request_meta.project_ids) == 1
206-
else {}
207-
),
208-
},
209198
)
210199

211200
self.__before_execute(in_msg)

‎snuba/web/rpc/storage_routing/routing_strategies/storage_routing.py‎

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -75,12 +75,24 @@ class RoutingContext:
7575
timer: Timer
7676
in_msg: ProtobufMessage
7777
query_id: str
78-
tenant_ids: dict[str, str | int]
7978
query_result: Optional[QueryResult] = field(default=None)
8079
extra_info: dict[str, Any] = field(default_factory=dict)
8180
allocation_policies_recommendations: dict[str, QuotaAllowance] = field(default_factory=dict)
8281
cluster_load_info: LoadInfo | None = field(default=None)
8382

83+
@property
84+
def tenant_ids(self) -> dict[str, str | int]:
85+
request_meta = extract_message_meta(self.in_msg)
86+
return {
87+
"organization_id": request_meta.organization_id,
88+
"referrer": request_meta.referrer,
89+
**(
90+
{"project_id": request_meta.project_ids[0]}
91+
if hasattr(request_meta, "project_ids") and len(request_meta.project_ids) == 1
92+
else {}
93+
),
94+
}
95+
8496

8597
@dataclass
8698
class TimeWindow:
@@ -455,7 +467,7 @@ def after_execute(self, routing_decision: RoutingDecision, error: Exception | No
455467
try:
456468
assert routing_decision.routing_context is not None
457469

458-
routing_decision.strategy.update_allocation_policies_balances(routing_decision, error)
470+
self.update_allocation_policies_balances(routing_decision, error)
459471

460472
# these metrics are meant to track reject/throttle/success decisions, so they get emitted even if the query did not run successfully after routing
461473
if not routing_decision.can_run:

‎tests/web/rpc/v1/routing_strategies/test_outcomes_based.py‎

Lines changed: 0 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -71,10 +71,6 @@ def test_outcomes_based_routing_queries_daily_table() -> None:
7171
in_msg=request,
7272
timer=Timer("test"),
7373
query_id=uuid.uuid4().hex,
74-
tenant_ids={
75-
"organization_id": _ORG_ID,
76-
"referrer": "something",
77-
},
7874
)
7975
in_msg_meta = extract_message_meta(context.in_msg)
8076
assert strategy._use_daily(in_msg_meta=in_msg_meta)
@@ -102,10 +98,6 @@ def test_outcomes_based_routing_sampled_data_past_thirty_days() -> None:
10298
in_msg=request,
10399
timer=Timer("test"),
104100
query_id=uuid.uuid4().hex,
105-
tenant_ids={
106-
"organization_id": _ORG_ID,
107-
"referrer": "something",
108-
},
109101
)
110102

111103
routing_decision = strategy.get_routing_decision(context)
@@ -126,10 +118,6 @@ def test_outcomes_based_routing_sampled_data_past_thirty_days() -> None:
126118
in_msg=request,
127119
timer=Timer("test"),
128120
query_id=uuid.uuid4().hex,
129-
tenant_ids={
130-
"organization_id": _ORG_ID,
131-
"referrer": "something",
132-
},
133121
)
134122

135123
routing_decision = strategy.get_routing_decision(context)
@@ -144,10 +132,6 @@ def test_outcomes_based_routing_sampled_data_past_thirty_days() -> None:
144132
in_msg=request,
145133
timer=Timer("test"),
146134
query_id=uuid.uuid4().hex,
147-
tenant_ids={
148-
"organization_id": _ORG_ID,
149-
"referrer": "something",
150-
},
151135
)
152136

153137
routing_decision = strategy.get_routing_decision(context)
@@ -169,10 +153,6 @@ def test_outcomes_based_routing_normal_mode(store_outcomes_fixture: Any) -> None
169153
in_msg=request,
170154
timer=Timer("test"),
171155
query_id=uuid.uuid4().hex,
172-
tenant_ids={
173-
"organization_id": _ORG_ID,
174-
"referrer": "something",
175-
},
176156
)
177157
)
178158
assert routing_decision.tier == Tier.TIER_1
@@ -194,10 +174,6 @@ def test_outcomes_based_routing_downsample(store_outcomes_fixture: Any) -> None:
194174
in_msg=request,
195175
timer=Timer("test"),
196176
query_id=uuid.uuid4().hex,
197-
tenant_ids={
198-
"organization_id": _ORG_ID,
199-
"referrer": "something",
200-
},
201177
)
202178
)
203179
assert routing_decision.tier == Tier.TIER_8
@@ -209,10 +185,6 @@ def test_outcomes_based_routing_downsample(store_outcomes_fixture: Any) -> None:
209185
in_msg=request,
210186
timer=Timer("test"),
211187
query_id=uuid.uuid4().hex,
212-
tenant_ids={
213-
"organization_id": _ORG_ID,
214-
"referrer": "something",
215-
},
216188
)
217189
)
218190
assert routing_decision.tier == Tier.TIER_64
@@ -225,10 +197,6 @@ def test_outcomes_based_routing_downsample(store_outcomes_fixture: Any) -> None:
225197
in_msg=request,
226198
timer=Timer("test"),
227199
query_id=uuid.uuid4().hex,
228-
tenant_ids={
229-
"organization_id": _ORG_ID,
230-
"referrer": "something",
231-
},
232200
)
233201
)
234202
assert routing_decision.tier == Tier.TIER_512
@@ -248,10 +216,6 @@ def test_outcomes_based_routing_highest_accuracy_mode(store_outcomes_fixture: An
248216
in_msg=request,
249217
timer=Timer("test"),
250218
query_id=uuid.uuid4().hex,
251-
tenant_ids={
252-
"organization_id": _ORG_ID,
253-
"referrer": "something",
254-
},
255219
)
256220
)
257221

@@ -275,10 +239,6 @@ def test_outcomes_based_routing_defaults_to_spans_for_unspecified_item_type(
275239
in_msg=request,
276240
timer=Timer("test"),
277241
query_id=uuid.uuid4().hex,
278-
tenant_ids={
279-
"organization_id": _ORG_ID,
280-
"referrer": "something",
281-
},
282242
)
283243
)
284244
assert routing_decision.tier == Tier.TIER_512

‎tests/web/rpc/v1/routing_strategies/test_outcomes_flex_time.py‎

Lines changed: 0 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -76,10 +76,6 @@ def test_outcomes_flex_time_routing_strategy_with_data() -> None:
7676
in_msg=request,
7777
timer=Timer("test"),
7878
query_id=uuid.uuid4().hex,
79-
tenant_ids={
80-
"organization_id": _ORG_ID,
81-
"referrer": "something",
82-
},
8379
)
8480
)
8581
assert routing_decision.time_window is not None
@@ -126,10 +122,6 @@ def test_outcomes_flex_time_routing_strategy_with_data_and_page_token() -> None:
126122
in_msg=request,
127123
timer=Timer("test"),
128124
query_id=uuid.uuid4().hex,
129-
tenant_ids={
130-
"organization_id": _ORG_ID,
131-
"referrer": "something",
132-
},
133125
)
134126
)
135127
assert routing_decision.time_window is not None

‎tests/web/rpc/v1/routing_strategies/test_strategy_selector.py‎

Lines changed: 0 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -110,10 +110,6 @@ def test_selects_same_strategy_for_same_org_and_project_ids() -> None:
110110
),
111111
timer=Timer(name="doesntmatter"),
112112
query_id=uuid.uuid4().hex,
113-
tenant_ids={
114-
"organization_id": 1,
115-
"referrer": "something",
116-
},
117113
)
118114

119115
for _ in range(50):
@@ -147,10 +143,6 @@ def test_selects_strategy_based_on_non_uniform_distribution() -> None:
147143
),
148144
timer=Timer(name="doesntmatter"),
149145
query_id=uuid.uuid4().hex,
150-
tenant_ids={
151-
"organization_id": 1,
152-
"referrer": "something",
153-
},
154146
)
155147
strategy = selector.select_routing_strategy(routing_context)
156148
strategy_counts[type(strategy)] += 1
@@ -177,10 +169,6 @@ def test_config_ordering_does_not_affect_routing_consistency() -> None:
177169
),
178170
timer=Timer(name="doesntmatter"),
179171
query_id=uuid.uuid4().hex,
180-
tenant_ids={
181-
"organization_id": 1,
182-
"referrer": "something",
183-
},
184172
)
185173

186174
assert isinstance(
@@ -220,10 +208,6 @@ def test_selects_override_if_it_exists() -> None:
220208
),
221209
timer=Timer(name="doesntmatter"),
222210
query_id=uuid.uuid4().hex,
223-
tenant_ids={
224-
"organization_id": 1,
225-
"referrer": "something",
226-
},
227211
)
228212

229213
assert RoutingStrategySelector().get_storage_routing_config(
@@ -255,10 +239,6 @@ def test_does_not_override_if_organization_id_is_different() -> None:
255239
),
256240
timer=Timer(name="doesntmatter"),
257241
query_id=uuid.uuid4().hex,
258-
tenant_ids={
259-
"organization_id": 1,
260-
"referrer": "something",
261-
},
262242
)
263243

264244
assert RoutingStrategySelector().get_storage_routing_config(

‎tests/web/rpc/v1/test_storage_routing.py‎

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -164,10 +164,6 @@ def get_allocation_policies(self) -> list[AllocationPolicy]:
164164
query_result=MagicMock(spec=QueryResult),
165165
extra_info={},
166166
query_id=uuid.uuid4().hex,
167-
tenant_ids={
168-
"organization_id": 1,
169-
"referrer": "something",
170-
},
171167
)
172168

173169

0 commit comments

Comments
 (0)