Skip to content

fix: automatic groupings for enrollment/exposure queries #523

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
mikewilli merged 2 commits into from
Oct 22, 2025
Merged

fix: automatic groupings for enrollment/exposure queries #523

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
b8cc253
fix: automatic groupings for enrollment/exposure queries
mikewilli Oct 22, 2025
03eb325
fix tests & ci run condition
mikewilli Oct 22, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions .github/workflows/build.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ name: Build, Test, and Deploy
on:
push:
branches:
- main
pull_request:

jobs:
Expand Down
49 changes: 11 additions & 38 deletions src/mozanalysis/experiment.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -858,7 +858,7 @@ def _build_enrollments_query_normandy(
BETWEEN '{time_limits.first_enrollment_date}' AND '{time_limits.last_enrollment_date}'
AND e.event_string_value = '{self.experiment_slug}'
AND e.sample_id < {sample_size}
GROUP BY e.{self.analysis_unit.value}, branch
GROUP BY ALL
""" # noqa:E501

def _build_enrollments_query_fenix_baseline(
Expand Down Expand Up @@ -895,7 +895,7 @@ def _build_enrollments_query_fenix_baseline(
'{experiment_slug}'
).branch IS NOT NULL
AND b.sample_id < {sample_size}
GROUP BY b.client_info.client_id, branch
GROUP BY ALL
HAVING enrollment_date >= '{first_enrollment_date}'
""".format(
experiment_slug=self.experiment_slug,
Expand All @@ -905,36 +905,6 @@ def _build_enrollments_query_fenix_baseline(
sample_size=sample_size,
)

def _build_enrollments_query_glean_event(
self, time_limits: TimeLimits, dataset: str, sample_size: int = 100
) -> str:
"""Deprecated; see _build_enrollments_query_glean_events_stream below

Return SQL to query enrollments for a Glean experiment from the
events table for the application or dataset.
"""

return f"""
SELECT events.client_info.client_id AS analysis_id,
mozfun.map.get_key(
e.extra,
'branch'
) AS branch,
DATE(MIN(events.submission_timestamp)) AS enrollment_date,
COUNT(events.submission_timestamp) AS num_enrollment_events
FROM `moz-fx-data-shared-prod.{self.app_id or dataset}.events` events,
UNNEST(events.events) AS e
WHERE
events.client_info.client_id IS NOT NULL AND
DATE(events.submission_timestamp)
BETWEEN '{time_limits.first_enrollment_date}' AND '{time_limits.last_enrollment_date}'
AND e.category = "nimbus_events"
AND mozfun.map.get_key(e.extra, "experiment") = '{self.experiment_slug}'
AND e.name = 'enrollment'
AND sample_id < {sample_size}
GROUP BY events.client_info.client_id, branch
""" # noqa:E501

def _build_enrollments_query_glean_events_stream(
self,
time_limits: TimeLimits,
Expand All @@ -961,7 +931,7 @@ def _build_enrollments_query_glean_events_stream(
AND JSON_VALUE(event_extra, "$.experiment") = "{self.experiment_slug}"
AND event_name = "enrollment"
AND sample_id < {sample_size}
GROUP BY client_id, branch
GROUP BY ALL
""" # noqa:E501

def _build_enrollments_query_glean_events_stream_enrollment_status(
Expand All @@ -972,6 +942,9 @@ def _build_enrollments_query_glean_events_stream_enrollment_status(
enrollments from the enrollment_status event.
"""

logger.warning("enrollment_status is not ready for production use.")
# TODO: dynamic analysis_id

return f"""
SELECT
client_id AS analysis_id,
Expand All @@ -989,7 +962,7 @@ def _build_enrollments_query_glean_events_stream_enrollment_status(
AND JSON_VALUE(event_extra, "$.status") = "Enrolled"
AND JSON_VALUE(event_extra, "$.reason") = "Qualified"
AND sample_id < {sample_size}
GROUP BY client_id, branch
GROUP BY ALL
""" # noqa:E501

def _build_enrollments_query_cirrus(
Expand Down Expand Up @@ -1024,7 +997,7 @@ def _build_enrollments_query_cirrus(
AND mozfun.map.get_key(e.extra, "experiment") = '{self.experiment_slug}'
AND e.name = 'enrollment'
AND client_info.app_channel = 'production'
GROUP BY mozfun.map.get_key(e.extra, "nimbus_user_id"), branch
GROUP BY ALL
""" # noqa:E501

def _build_exposure_query_normandy(self, time_limits: TimeLimits) -> str:
Expand Down Expand Up @@ -1053,7 +1026,7 @@ def _build_exposure_query_normandy(self, time_limits: TimeLimits) -> str:
ON re.analysis_id = e.analysis_id AND
re.branch = e.branch AND
e.submission_date >= re.enrollment_date
GROUP BY e.analysis_id, e.branch
GROUP BY ALL
""" # noqa: E501

def _build_exposure_query_glean_event(
Expand Down Expand Up @@ -1091,7 +1064,7 @@ def _build_exposure_query_glean_event(
ON re.analysis_id = exposures.analysis_id AND
re.branch = exposures.branch AND
exposures.submission_date >= re.enrollment_date
GROUP BY analysis_id, branch
GROUP BY ALL
""" # noqa: E501

def _build_exposure_query_glean_events_stream(
Expand Down Expand Up @@ -1126,7 +1099,7 @@ def _build_exposure_query_glean_events_stream(
ON re.analysis_id = exposures.analysis_id AND
re.branch = exposures.branch AND
exposures.submission_date >= re.enrollment_date
GROUP BY analysis_id, branch
GROUP BY ALL
""" # noqa: E501

def _build_metrics_query_bits(
Expand Down
8 changes: 4 additions & 4 deletions tests/test_experiment.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -304,7 +304,7 @@ def test_query_not_detectably_malformed(analysis_unit: AnalysisUnit):
sql_lint(enrollments_sql)
assert "sample_id < None" not in enrollments_sql

assert enrollments_sql.count(analysis_unit.value) == 3
assert enrollments_sql.count(analysis_unit.value) == 2

metrics_sql = exp.build_metrics_query(
metric_list=[],
Expand Down Expand Up @@ -345,7 +345,7 @@ def test_megaquery_not_detectably_malformed(

sql_lint(enrollments_sql)

assert enrollments_sql.count(analysis_unit.value) == 3
assert enrollments_sql.count(analysis_unit.value) == 2

metrics_sql = exp.build_metrics_query(
metric_list=desktop_metrics,
Expand Down Expand Up @@ -1069,7 +1069,7 @@ def test_enrollments_query_analysis_unit(analysis_unit):
BETWEEN '2019-01-01' AND '2019-01-08'
AND e.event_string_value = 'slug'
AND e.sample_id < 100
GROUP BY e.{analysis_unit.value}, branch
GROUP BY ALL
),
segmented_enrollments AS (
SELECT
Expand Down Expand Up @@ -1102,7 +1102,7 @@ def test_enrollments_query_analysis_unit(analysis_unit):
ON re.analysis_id = e.analysis_id AND
re.branch = e.branch AND
e.submission_date >= re.enrollment_date
GROUP BY e.analysis_id, e.branch
GROUP BY ALL
)

SELECT
Expand Down