22import sys
33from contextlib import asynccontextmanager
44from pathlib import Path
5- from typing import Any , List , NamedTuple
5+ from typing import Any , List , NamedTuple , TypedDict
66
77import sentry_sdk
88from apscheduler .schedulers .asyncio import AsyncIOScheduler # type: ignore
@@ -161,53 +161,59 @@ def initialize_glean():
161161
162162
163163class EnrollmentMetricData (NamedTuple ):
164+ nimbus_user_id : str
165+ app_id : str
164166 experiment_slug : str
165167 branch_slug : str
166168 experiment_type : str
169+ is_preview : bool
170+
171+
172+ class ComputeFeaturesEnrollmentResult (TypedDict ):
173+ features : dict [str , dict [str , Any ]]
174+ enrollments : list [EnrollmentMetricData ]
167175
168176
169177def collate_enrollment_metric_data (
170- enrolled_partial_configuration : dict [str , Any ], nimbus_preview_flag : bool
178+ enrolled_partial_configuration : dict [str , Any ],
179+ client_id : str ,
180+ nimbus_preview_flag : bool ,
171181) -> list [EnrollmentMetricData ]:
172182 events : list [dict [str , Any ]] = enrolled_partial_configuration .get ("events" , [])
183+ remote_settings = (
184+ app .state .remote_setting_preview
185+ if nimbus_preview_flag
186+ else app .state .remote_setting_live
187+ )
173188 data : list [EnrollmentMetricData ] = []
174189 for event in events :
175190 if event .get ("change" ) == "Enrollment" :
176191 experiment_slug = event .get ("experiment_slug" , "" )
177192 branch_slug = event .get ("branch_slug" , "" )
178- experiment_type = None
179- remote_settings = app .state .remote_setting_live
180- if nimbus_preview_flag :
181- remote_settings = app .state .remote_setting_preview
182193 experiment_type = remote_settings .get_recipe_type (experiment_slug )
183194 data .append (
184195 EnrollmentMetricData (
196+ nimbus_user_id = client_id ,
197+ app_id = app_id ,
185198 experiment_slug = experiment_slug ,
186199 branch_slug = branch_slug ,
187200 experiment_type = experiment_type ,
201+ is_preview = nimbus_preview_flag ,
188202 )
189203 )
190204 return data
191205
192206
193- async def record_metrics (
194- enrolled_partial_configuration : dict [str , Any ],
195- client_id : str ,
196- nimbus_preview_flag : bool ,
197- ):
198- metrics = collate_enrollment_metric_data (
199- enrolled_partial_configuration = enrolled_partial_configuration ,
200- nimbus_preview_flag = nimbus_preview_flag ,
201- )
202- for experiment_slug , branch_slug , experiment_type in metrics :
207+ async def record_metrics (enrollment_data : list [EnrollmentMetricData ]):
208+ for enrollment in enrollment_data :
203209 app .state .metrics .cirrus_events .enrollment .record (
204210 app .state .metrics .cirrus_events .EnrollmentExtra (
205- user_id = client_id ,
206- app_id = app_id ,
207- experiment = experiment_slug ,
208- branch = branch_slug ,
209- experiment_type = experiment_type ,
210- is_preview = nimbus_preview_flag ,
211+ user_id = enrollment . nimbus_user_id ,
212+ app_id = enrollment . app_id ,
213+ experiment = enrollment . experiment_slug ,
214+ branch = enrollment . branch_slug ,
215+ experiment_type = enrollment . experiment_type ,
216+ is_preview = enrollment . is_preview ,
211217 )
212218 )
213219 app .state .pings .enrollment .submit ()
@@ -221,11 +227,10 @@ def read_root():
221227 return {"Hello" : "World" }
222228
223229
224- @app .post ("/v1/features/" , status_code = status .HTTP_200_OK )
225- async def compute_features (
230+ async def compute_features_enrollments (
226231 request_data : FeatureRequest ,
227232 nimbus_preview : bool = Query (default = False , alias = "nimbus_preview" ),
228- ):
233+ ) -> ComputeFeaturesEnrollmentResult :
229234 if not request_data .client_id :
230235 raise HTTPException (
231236 status_code = status .HTTP_400_BAD_REQUEST ,
@@ -246,9 +251,8 @@ async def compute_features(
246251 "clientId" : request_data .client_id ,
247252 "requestContext" : request_data .context ,
248253 }
249- sdk = app .state .sdk_live
250- if nimbus_preview :
251- sdk = app .state .sdk_preview
254+
255+ sdk = app .state .sdk_preview if nimbus_preview else app .state .sdk_live
252256 enrolled_partial_configuration : dict [str , Any ] = sdk .compute_enrollments (
253257 targeting_context
254258 )
@@ -257,13 +261,51 @@ async def compute_features(
257261 app .state .fml .compute_feature_configurations (enrolled_partial_configuration )
258262 )
259263
260- await record_metrics (
261- enrolled_partial_configuration = enrolled_partial_configuration ,
264+ # Enrollments data
265+ enrollment_data = collate_enrollment_metric_data (
266+ enrolled_partial_configuration ,
262267 client_id = request_data .client_id ,
263- nimbus_preview_flag = nimbus_preview or False ,
268+ nimbus_preview_flag = nimbus_preview ,
264269 )
265270
266- return client_feature_configuration
271+ # Record metrics
272+ await record_metrics (enrollment_data )
273+
274+ return {
275+ "features" : client_feature_configuration ,
276+ "enrollments" : enrollment_data ,
277+ }
278+
279+
280+ @app .post ("/v1/features/" , status_code = status .HTTP_200_OK )
281+ async def compute_features_v1 (
282+ request_data : FeatureRequest ,
283+ nimbus_preview : bool = Query (default = False , alias = "nimbus_preview" ),
284+ ):
285+ result = await compute_features_enrollments (request_data , nimbus_preview )
286+ return result ["features" ]
287+
288+
289+ @app .post ("/v2/features/" , status_code = status .HTTP_200_OK )
290+ async def compute_features_enrollments_v2 (
291+ request_data : FeatureRequest ,
292+ nimbus_preview : bool = Query (default = False , alias = "nimbus_preview" ),
293+ ):
294+ result = await compute_features_enrollments (request_data , nimbus_preview )
295+ return {
296+ "Features" : result ["features" ],
297+ "Enrollments" : [
298+ {
299+ "nimbus_user_id" : enrollment .nimbus_user_id ,
300+ "app_id" : enrollment .app_id ,
301+ "experiment" : enrollment .experiment_slug ,
302+ "branch" : enrollment .branch_slug ,
303+ "experiment_type" : enrollment .experiment_type ,
304+ "is_preview" : enrollment .is_preview ,
305+ }
306+ for enrollment in result ["enrollments" ]
307+ ],
308+ }
267309
268310
269311async def fetch_schedule_recipes () -> None :
0 commit comments