Merge branch 'main' into query-schema-generation-dryrun-caching

Author: scholtzan

.circleci/config.yml

@@ -59,6 +59,7 @@ workflows:
             sql/mozfun/.* validate-routines true
             sql/moz-fx-data-shared-prod/udf/.* validate-routines true
             sql/moz-fx-data-shared-prod/udf_js/.* validate-routines true
+            bqetl_project.yaml trigger-sql-generation true
             bqetl_project.yaml validate-sql true
             bqetl_project.yaml validate-routines true
             dags.yaml validate-sql true

.circleci/workflows.yml

@@ -352,6 +352,15 @@ jobs:
             - *restore_venv_cache
             - *build
             - *attach_generated_sql
+            - &restore_schema_cache
+              run:
+                name: Restore schema cache from generate-sql job
+                command: |
+                  # Restore stable table schema cache from workspace
+                  if [ -d /tmp/workspace/schema-cache/bigquery_etl_schemas ]; then
+                    cp -r /tmp/workspace/schema-cache/bigquery_etl_schemas /tmp/
+                    echo "Restored schema cache from generate-sql job"
+                  fi
             - *authenticate
             - &add_private_bigquery_etl_ssh_keys
               add_ssh_keys:
@@ -407,6 +416,11 @@ jobs:
                 command: |
                   rm ~/telemetry-airflow/dags/* -f || true
                   cp -a /tmp/workspace/generated-sql/dags/. ~/telemetry-airflow/dags/
+            - restore_cache:
+                keys:
+                  - telemetry-airflow-deps-v1-{{ checksum "~/telemetry-airflow/requirements.txt" }}-{{ checksum "~/telemetry-airflow/requirements-dev.txt" }}-{{ checksum "~/telemetry-airflow/requirements-override.txt" }}
+                  - telemetry-airflow-deps-v1-{{ checksum "~/telemetry-airflow/requirements.txt" }}-{{ checksum "~/telemetry-airflow/requirements-dev.txt" }}-
+                  - telemetry-airflow-deps-v1-
             - run:
                 name: Install telemetry-airflow dependencies
                 command: |
@@ -416,6 +430,10 @@ jobs:
                   pip install -r requirements.txt
                   pip install -r requirements-dev.txt
                   pip install -r requirements-override.txt
+            - save_cache:
+                paths:
+                  - ~/telemetry-airflow/.venv
+                key: telemetry-airflow-deps-v1-{{ checksum "~/telemetry-airflow/requirements.txt" }}-{{ checksum "~/telemetry-airflow/requirements-dev.txt" }}-{{ checksum "~/telemetry-airflow/requirements-override.txt" }}
             - run:
                 name: 🧪 Test valid DAGs
                 command: |
@@ -610,6 +628,7 @@ jobs:
                       --ignore derived_view_schemas \
                       --output-dir /tmp/workspace/generated-sql/sql/ \
                       --target-project moz-fx-data-shared-prod
+                    PATH="venv/bin:$PATH" script/bqetl format /tmp/workspace/generated-sql/sql/
                   else
                     echo "Changes made don't affect generated SQL. Use content from generated-sql"
 
@@ -618,18 +637,26 @@ jobs:
                     cp -a sql/. /tmp/workspace/generated-sql/sql
                   fi
 
-                  PATH="venv/bin:$PATH" script/bqetl format /tmp/workspace/generated-sql/sql/
                   PATH="venv/bin:$PATH" script/bqetl dependency record \
                     --skip-existing \
                     "/tmp/workspace/generated-sql/sql/"
                   PATH="venv/bin:$PATH" script/bqetl metadata update \
                     --sql-dir /tmp/workspace/generated-sql/sql/ \
                     /tmp/workspace/generated-sql/sql/
                   PATH="venv/bin:$PATH" script/bqetl monitoring update /tmp/workspace/generated-sql/sql/
+            - run:
+                name: Copy schema cache for reuse in other jobs
+                command: |
+                  # Copy stable table schema cache to workspace for reuse in generate-dags
+                  mkdir -p /tmp/workspace/schema-cache
+                  if [ -d /tmp/bigquery_etl_schemas ]; then
+                    cp -r /tmp/bigquery_etl_schemas /tmp/workspace/schema-cache/
+                  fi
             - persist_to_workspace:
                 root: /tmp/workspace
                 paths:
                   - generated-sql
+                  - schema-cache
       - unless:
           condition: *validate-sql-or-routines
           steps:
@@ -646,6 +673,7 @@ jobs:
             - *restore_venv_cache
             - *build
             - *attach_generated_sql
+            - *restore_schema_cache
             - *copy_generated_sql
             - add_ssh_keys:
                 fingerprints:
@@ -1106,8 +1134,8 @@ jobs:
           name: Manually trigger integration tests for fork
           # yamllint disable rule:line-length
           command: |
-            apt update
-            apt install jq -y
+            sudo apt update
+            sudo apt install jq -y
 
             CIRCLE_PR_BRANCH=`curl -s https://api.github.com/repos/${CIRCLE_PROJECT_USERNAME}/${CIRCLE_PROJECT_REPONAME}/pulls/${CIRCLE_PR_NUMBER} | jq -r '.head.label'`
 

bigquery_etl/glam/templates/clients_daily_histogram_aggregates_v1.sql

@@ -17,6 +17,7 @@ WITH extracted AS (
   WHERE
     DATE(submission_timestamp) = {{ submission_date }}
     AND client_info.client_id IS NOT NULL
+    AND LOWER(IFNULL(metadata.isp.name, "")) <> "browserstack" -- Removes bots data.
 ),
 histograms AS (
   SELECT

bigquery_etl/glam/templates/clients_daily_scalar_aggregates_v1.sql

@@ -17,6 +17,7 @@ WITH extracted AS (
   WHERE
     DATE(submission_timestamp) = {{ submission_date }}
     AND client_info.client_id IS NOT NULL
+    AND LOWER(IFNULL(metadata.isp.name, "")) <> "browserstack" -- Removes bots data.
 ),
 unlabeled_metrics AS (
   SELECT

bigquery_etl/query_scheduling/dag.py

@@ -122,6 +122,7 @@ class Dag:
     repo: str = attr.ib("bigquery-etl")
     tags: List[str] = attr.ib([])
     catchup: bool = attr.ib(False)
+    max_active_runs: Optional[int] = attr.ib(None)
 
     @name.validator
     def validate_dag_name(self, attribute, value):

bigquery_etl/query_scheduling/dag_collection.py

@@ -10,8 +10,31 @@
 import yaml
 from black import FileMode, format_file_contents
 
+from bigquery_etl.dependency import extract_table_references_without_views
 from bigquery_etl.query_scheduling.dag import Dag, InvalidDag, PublicDataJsonDag
 from bigquery_etl.query_scheduling.utils import negate_timedelta_string
+from bigquery_etl.schema.stable_table_schema import get_stable_table_schemas
+
+
+def _precompute_task_references(task):
+    """Pre-compute expensive table references parsing and return result."""
+    if (
+        task.is_python_script
+        or task.is_bigeye_check
+        or task.referenced_tables is not None
+    ):
+        return (task.task_key, task.referenced_tables or [])
+
+    query_files = [Path(task.query_file)]
+    if task.multipart:
+        query_files = list(Path(task.query_file_path).glob("*.sql"))
+
+    table_names = {
+        tuple(table.split("."))
+        for query_file in query_files
+        for table in extract_table_references_without_views(query_file)
+    }
+    return (task.task_key, sorted(table_names))
 
 
 class DagCollection:
@@ -192,10 +215,32 @@ def to_airflow_dags(self, output_dir, dag_to_generate=None):
         except Exception:
             pass
 
+        # Pre-load stable table schemas before spawning workers to avoid loading multiple times
+        # This downloads and caches schemas once in the main process
+        try:
+            get_stable_table_schemas()
+        except Exception:
+            # If schema loading fails, continue anyway (some tasks may not need them)
+            pass
+
+        # Pre-compute referenced tables for all tasks in parallel
+        # This is the expensive I/O-heavy part (SQL parsing via extract_table_references_without_views)
+        all_tasks = [task for dag in self.dags for task in dag.tasks]
+
+        with get_context("spawn").Pool(8) as p:
+            task_references = p.map(_precompute_task_references, all_tasks)
+
+        # Update tasks with precomputed references
+        task_map = {task.task_key: task for task in all_tasks}
+        for task_key, referenced_tables in task_references:
+            task_map[task_key].referenced_tables = referenced_tables
+
+        # Resolve dependencies sequentially
         for dag in self.dags:
             dag.with_upstream_dependencies(self)
             dag.with_downstream_dependencies(self)
 
+        # Finally, parallelize DAG-to-Airflow conversion
         to_airflow_dag = partial(self.dag_to_airflow, output_dir)
         with get_context("spawn").Pool(8) as p:
             p.map(to_airflow_dag, self.dags)

bigquery_etl/query_scheduling/templates/airflow_dag.j2

@@ -82,7 +82,19 @@ default_args = {{
 
 tags = {{ tags }}
 
-with DAG('{{ name }}', default_args=default_args{%+ if schedule_interval != None -%}, schedule_interval={{ schedule_interval | format_timedelta | format_schedule_interval }}{%+ endif -%}, doc_md=docs, tags=tags, catchup={{ catchup }}) as dag:
+with DAG(
+    '{{ name }}',
+    default_args=default_args,
+    {% if schedule_interval != None -%}
+    schedule_interval={{ schedule_interval | format_timedelta | format_schedule_interval }},
+    {%- endif %}
+    {% if max_active_runs -%}
+    max_active_runs={{ max_active_runs }},
+    {%- endif %}
+    doc_md=docs,
+    tags=tags,
+    catchup={{ catchup }},
+) as dag:
 {% for task_group in task_groups | sort %}
     task_group_{{ task_group }} = TaskGroup('{{ task_group }}')
 {% endfor %}

bigquery_etl/query_scheduling/templates/public_data_json_airflow_dag.j2

@@ -38,7 +38,18 @@ default_args = {{
 
 tags = {{ tags }}
 
-with DAG('{{ name }}', default_args=default_args{%+ if schedule_interval != None -%}, schedule_interval={{ schedule_interval | format_timedelta | format_schedule_interval }}{%+ endif -%}, doc_md = docs, tags = tags) as dag:
+with DAG(
+    '{{ name }}',
+    default_args=default_args,
+    {% if schedule_interval != None -%}
+    schedule_interval={{ schedule_interval | format_timedelta | format_schedule_interval }},
+    {%- endif %}
+    {% if max_active_runs -%}
+    max_active_runs={{ max_active_runs }},
+    {%- endif %}
+    doc_md=docs,
+    tags=tags,
+) as dag:
     docker_image = "gcr.io/moz-fx-data-airflow-prod-88e0/bigquery-etl:latest"
 {% for task_group in task_groups | sort %}
     task_group_{{ task_group }} = TaskGroup('{{ task_group }}')

bigquery_etl/util/common.py

@@ -8,6 +8,7 @@
 import string
 import tempfile
 import warnings
+from functools import cache
 from pathlib import Path
 from typing import List, Optional, Set, Tuple
 from uuid import uuid4
@@ -17,7 +18,7 @@
 from google.cloud import bigquery
 from jinja2 import Environment, FileSystemLoader
 
-from bigquery_etl.config import ConfigLoader
+from bigquery_etl.config import BQETL_PROJECT_CONFIG, ConfigLoader
 from bigquery_etl.format_sql.formatter import reformat
 from bigquery_etl.metrics import MetricHub
 
@@ -70,6 +71,18 @@ def random_str(length: int = 12) -> str:
     return "".join(random.choice(string.ascii_lowercase) for i in range(length))
 
 
+@cache
+def get_bqetl_project_root() -> Path | None:
+    """Return the root path of the bqetl project the user is currently in."""
+    cwd = Path.cwd()
+    search_paths = [cwd]
+    search_paths.extend(cwd.parents)
+    for possible_project_root in search_paths:
+        if (possible_project_root / BQETL_PROJECT_CONFIG).exists():
+            return possible_project_root
+    return None
+
+
 def render(
     sql_filename,
     template_folder=".",
@@ -78,7 +91,8 @@ def render(
     **kwargs,
 ) -> str:
     """Render a given template query using Jinja."""
-    path = Path(template_folder) / sql_filename
+    template_folder_path = Path(template_folder)
+    path = template_folder_path / sql_filename
     skip = {
         file
         for skip in ConfigLoader.get("render", "skip", fallback=[])
@@ -122,14 +136,20 @@ def render(
                 checks_template.write_text(
                     macro_imports
                     + "\n"
-                    + (Path(template_folder) / sql_filename).read_text()
+                    + (template_folder_path / sql_filename).read_text()
                 )
 
                 file_loader = FileSystemLoader(f"{str(checks_template.parent)}")
                 env = Environment(loader=file_loader)
                 main_sql = env.get_template(checks_template.name)
         else:
-            file_loader = FileSystemLoader(f"{template_folder}")
+            # Add the bigquery-etl project root to the search path to support Jinja imports/includes.
+            file_loader_search_paths = [template_folder_path, ROOT]
+            # Also dynamically detect the project root so imports/includes in the private-bigquery-etl repo work.
+            if bqetl_project_root := get_bqetl_project_root():
+                if bqetl_project_root not in file_loader_search_paths:
+                    file_loader_search_paths.append(bqetl_project_root)
+            file_loader = FileSystemLoader(file_loader_search_paths)
             env = Environment(loader=file_loader)
             main_sql = env.get_template(sql_filename)
 

bqetl_project.yaml

@@ -374,6 +374,7 @@ generate:
       - firefox_reality_pc
       - lockwise_android
       - lockwise_ios
+      - mozilla_vpn  # VPN events are in `main` pings rather than `events` pings (DENG-8596).
       - tiktokreporter_ios
       - tiktokreporter_android
       metrics_as_struct_apps:
@@ -530,19 +531,21 @@ generate:
       fenix:
         templates:
         - terms_of_use_status_v1
+        - terms_of_use_events_v1
         bigeye:
           collection: Operational Checks
           notification_channel: '#de-bigeye-triage'
       firefox_ios:
         templates:
         - terms_of_use_status_v1
+        - terms_of_use_events_v1
         bigeye:
           collection: Operational Checks
           notification_channel: '#de-bigeye-triage'
       firefox_desktop:
         templates:
         - terms_of_use_status_v1
-        - terms_of_use_events_v1
+        - terms_of_use_messages_v1
         bigeye:
           collection: Operational Checks
           notification_channel: '#de-bigeye-triage'
@@ -562,6 +565,23 @@ generate:
 geo_deprecation:
   include_app_ids:
   - org_mozilla_ios_klar
+  - org_mozilla_klar
+  - org_mozilla_ios_focus
+  - org_mozilla_focus
+  - org_mozilla_focus_beta
+  - org_mozilla_focus_nightly
+  - org_mozilla_ios_firefox
+  - org_mozilla_ios_firefoxbeta
+  - org_mozilla_ios_fennec
+  - org_mozilla_firefox
+  - org_mozilla_firefox_beta
+  - org_mozilla_fenix
+  - org_mozilla_fenix_nightly
+  - org_mozilla_fennec_aurora
+  - firefox_desktop
+  - firefox_desktop_background_update
+  - firefox_desktop_background_defaultagent
+  - firefox_desktop_background_tasks
   skip_tables:
   - newtab
 
@@ -570,6 +590,7 @@ retention_exclusion_list:
 - sql/moz-fx-data-shared-prod/telemetry_derived/clients_first_seen_v3
 - sql/moz-fx-data-shared-prod/telemetry_derived/clients_first_seen_28_days_later_v1
 - sql/moz-fx-data-shared-prod/telemetry_derived/clients_first_seen_28_days_later_v3
+- sql/moz-fx-data-shared-prod/firefox_desktop_derived/clients_first_seen_28_days_later_v1
 - sql/moz-fx-data-shared-prod/org_mozilla_firefox_derived/baseline_clients_first_seen_v1
 - sql/moz-fx-data-shared-prod/org_mozilla_fenix_derived/baseline_clients_first_seen_v1
 - sql/moz-fx-data-shared-prod/fenix_derived/ltv_states_v1
@@ -653,3 +674,4 @@ retention_exclusion_list:
 - sql/moz-fx-data-shared-prod/firefox_desktop_background_update/baseline_clients_city_seen_v1
 - sql/moz-fx-data-shared-prod/firefox_desktop_background_defaultagent/baseline_clients_city_seen_v1
 - sql/moz-fx-data-shared-prod/firefox_desktop_background_tasks/baseline_clients_city_seen_v1
+- sql/moz-fx-data-shared-prod/firefox_desktop_derived/metrics_clients_first_seen_v1