fix: use actual analysis unit for events stream queries

src/mozanalysis/experiment.py

@@ -738,7 +738,7 @@ def _build_enrollments_query(
         if enrollments_query_type == EnrollmentsQueryType.NORMANDY:
             if use_glean_ids:
                 return self._build_enrollments_query_glean_events_stream(
-                    time_limits, self.app_id, sample_size
+                    time_limits, self.app_id, sample_size, self.analysis_unit.value
                 )
             else:
                 return self._build_enrollments_query_normandy(
@@ -798,6 +798,7 @@ def _build_exposure_query(
             if use_glean_ids:
                 return self._build_exposure_query_glean_events_stream(
                     time_limits,
+                    client_id_field=self.analysis_unit.value,
                 )
             else:
                 return self._build_exposure_query_normandy(time_limits)
@@ -927,7 +928,7 @@ def _build_enrollments_query_glean_events_stream(
                 COUNT(submission_timestamp) AS num_enrollment_events
             FROM `moz-fx-data-shared-prod.{self.app_id or dataset}.events_stream`
             WHERE
-                client_id IS NOT NULL
+                {analysis_id} IS NOT NULL
                 AND DATE(submission_timestamp)
                     BETWEEN '{time_limits.first_enrollment_date}' AND '{time_limits.last_enrollment_date}'
                 AND event_category = "nimbus_events"

tests/test_experiment.py

@@ -1137,6 +1137,92 @@ def test_enrollments_query_analysis_unit(analysis_unit):
     assert dedent(enrollments_sql) == expected
 
 
+@pytest.mark.parametrize(
+    "analysis_unit", [AnalysisUnit.CLIENT, AnalysisUnit.PROFILE_GROUP]
+)
+def test_enrollments_query_glean_ids_analysis_unit(analysis_unit):
+    exp = Experiment(
+        "slug", "2019-01-01", 8, analysis_unit=analysis_unit, app_id="my_cool_app"
+    )
+
+    tl = TimeLimits.for_ts(
+        first_enrollment_date="2019-01-01",
+        last_date_full_data="2019-03-01",
+        time_series_period="weekly",
+        num_dates_enrollment=8,
+    )
+
+    enrollments_sql = exp.build_enrollments_query(
+        time_limits=tl,
+        enrollments_query_type=EnrollmentsQueryType.NORMANDY,
+        use_glean_ids=True,
+    )
+
+    sql_lint(enrollments_sql)
+
+    expected = f"""
+    WITH raw_enrollments AS (
+    SELECT
+        {analysis_unit.value} AS analysis_id,
+        JSON_VALUE(event_extra, '$.branch') AS branch,
+        DATE(MIN(submission_timestamp)) AS enrollment_date,
+        COUNT(submission_timestamp) AS num_enrollment_events
+    FROM `moz-fx-data-shared-prod.my_cool_app.events_stream`
+    WHERE
+        {analysis_unit.value} IS NOT NULL
+        AND DATE(submission_timestamp)
+            BETWEEN '2019-01-01' AND '2019-01-08'
+        AND event_category = "nimbus_events"
+        AND JSON_VALUE(event_extra, "$.experiment") = "slug"
+        AND event_name = "enrollment"
+        AND sample_id < 100
+    GROUP BY ALL
+    ),
+    segmented_enrollments AS (
+SELECT
+    raw_enrollments.*,
+
+FROM raw_enrollments
+
+),
+    exposures AS (
+    SELECT
+        exposures.analysis_id AS analysis_id,
+        exposures.branch,
+        DATE(MIN(exposures.submission_date)) AS exposure_date,
+        COUNT(exposures.submission_date) AS num_exposure_events
+    FROM raw_enrollments re
+    LEFT JOIN (
+        SELECT
+            {analysis_unit.value} AS analysis_id,
+            JSON_VALUE(event_extra, '$.branch') AS branch,
+            DATE(submission_timestamp) AS submission_date
+        FROM
+            `moz-fx-data-shared-prod.my_cool_app.events_stream`
+        WHERE
+            DATE(submission_timestamp)
+            BETWEEN '2019-01-01' AND '2019-01-08'
+            AND event_category = 'nimbus_events'
+            AND JSON_VALUE(event_extra, "$.experiment") = 'slug'
+            AND (event_name = 'expose' OR event_name = 'exposure')
+    ) exposures
+    ON re.analysis_id = exposures.analysis_id AND
+        re.branch = exposures.branch AND
+        exposures.submission_date >= re.enrollment_date
+    GROUP BY ALL
+    )
+
+    SELECT
+        se.*,
+        e.* EXCEPT (analysis_id, branch)
+    FROM segmented_enrollments se
+    LEFT JOIN exposures e
+    USING (analysis_id, branch)
+"""
+
+    assert dedent(enrollments_sql) == expected
+
+
 @pytest.mark.parametrize(
     "analysis_unit", [AnalysisUnit.CLIENT, AnalysisUnit.PROFILE_GROUP]
 )