Detect stale metrics registry entries in the metrics sync prek hook

.pre-commit-config.yaml

@@ -944,9 +944,10 @@ repos:
         name: Check that metrics in the codebase are in sync with the metrics registry YAML file.
         entry: ./scripts/ci/prek/check_metrics_synced_with_the_registry.py
         language: python
-        files: \.py$
+        files: \.py$|^shared/observability/src/airflow_shared/observability/metrics/metrics_template\.yaml$
         exclude: ^(tests/|.*/tests/)
         pass_filenames: true
+        require_serial: true
         additional_dependencies: ["PyYAML>=6.0", "rich>=13.6.0"]
       - id: check-boring-cyborg-configuration
         name: Checks for Boring Cyborg configuration consistency

scripts/ci/prek/check_metrics_synced_with_the_registry.py

@@ -32,6 +32,7 @@
 import argparse
 import ast
 import re
+import subprocess
 import sys
 from dataclasses import dataclass
 from pathlib import Path
@@ -57,6 +58,16 @@
     AIRFLOW_ROOT_PATH / "shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml"
 )
 
+EXCLUDED_TESTS_PATTERN = re.compile(r"(^|/)tests/")
+
+# Registry metrics emitted through a variable that the AST scan cannot resolve back to a
+# string literal, e.g. names built by BaseExecutor._get_metric_name.
+INDIRECTLY_EMITTED_METRICS = {
+    "executor.open_slots",
+    "executor.queued_tasks",
+    "executor.running_tasks",
+}
+
 
 def load_metrics_registry_yaml() -> dict[str, dict[str, Any]]:
     """Load the metrics registry YAML and return a dict keyed by metric name."""
@@ -82,6 +93,12 @@ def normalize_metric_name(registry_metric_name: str) -> str:
 _PREFIX_MATCHED = "__prefix_matched__"
 
 
+def find_prefix_matched_registry_entries(metric_name: str, metrics_registry: dict[str, dict]) -> list[str]:
+    """Return the registry entry names whose name matches the static prefix of a dynamic metric name."""
+    base = metric_name.split("{")[0].rstrip(".")
+    return [name for name in metrics_registry if name == base or name.startswith(base + ".")]
+
+
 def find_registry_match(metric_name: str, metrics_registry: dict[str, dict]) -> str | None:
     """Return the registry entry name that best matches the given metric name, or None."""
     if metric_name in metrics_registry:
@@ -99,16 +116,13 @@ def find_registry_match(metric_name: str, metrics_registry: dict[str, dict]) ->
             return registry_metric_name
 
     # Dynamic metric name.
-    if "{" in metric_name:
-        base = metric_name.split("{")[0].rstrip(".")
-        for registry_metric_name in metrics_registry:
-            if registry_metric_name == base or registry_metric_name.startswith(base + "."):
-                # Metric prefix matches the prefix of a dynamic registry entry.
-                # If the static part before the first variable, matches an exact registry entry name,
-                # or a dotted-prefix of one, then τηε name is considered covered and
-                # _PREFIX_MATCHED is returned. The type check must be skipped because
-                # the resulting metric name with all variables expanded, cannot be determined.
-                return _PREFIX_MATCHED
+    if "{" in metric_name and find_prefix_matched_registry_entries(metric_name, metrics_registry):
+        # Metric prefix matches the prefix of a dynamic registry entry.
+        # If the static part before the first variable, matches an exact registry entry name,
+        # or a dotted-prefix of one, then the name is considered covered and
+        # _PREFIX_MATCHED is returned. The type check must be skipped because
+        # the resulting metric name with all variables expanded, cannot be determined.
+        return _PREFIX_MATCHED
 
     # All checks for matching failed.
     return None
@@ -167,6 +181,22 @@ def extract_metric_name_from_ast_node(name_node: ast.expr) -> str | None:
     return None
 
 
+def extract_metric_names_from_ast_node(name_node: ast.expr) -> list[str]:
+    """Resolve a metric name AST node to all names it can produce.
+
+    A conditional expression like ``"a.success" if ok else "a.failed"`` yields both names.
+    Returns an empty list when the expression is too dynamic to resolve.
+    """
+    if isinstance(name_node, ast.IfExp):
+        return [
+            name
+            for branch in (name_node.body, name_node.orelse)
+            for name in extract_metric_names_from_ast_node(branch)
+        ]
+    metric_name = extract_metric_name_from_ast_node(name_node)
+    return [metric_name] if metric_name is not None else []
+
+
 @dataclass
 class MetricCall:
     file_path: str
@@ -181,8 +211,16 @@ def scan_file_for_metrics(file_path: Path) -> list[MetricCall]:
     """Return all Stats metric calls found in the provided file_path."""
     try:
         source = file_path.read_text(encoding="utf-8")
+    except (OSError, UnicodeDecodeError):
+        return []
+
+    # Cheap text pre-filter to avoid parsing the vast majority of files that emit no metrics.
+    if "Stats." not in source and "stats." not in source:
+        return []
+
+    try:
         tree = ast.parse(source, filename=str(file_path))
-    except (OSError, UnicodeDecodeError, SyntaxError):
+    except SyntaxError:
         return []
 
     metrics_found: list[MetricCall] = []
@@ -200,26 +238,59 @@ def scan_file_for_metrics(file_path: Path) -> list[MetricCall]:
             if name_node is None:
                 continue
 
-            metric_name = extract_metric_name_from_ast_node(name_node)
-            if metric_name is None:
+            metric_names = extract_metric_names_from_ast_node(name_node)
+            if not metric_names:
                 # Metric name is unresolvable. Probably has too many variables.
                 continue
 
             stats_obj = get_stats_obj_name(node.func.value)
-            metrics_found.append(
-                MetricCall(
-                    file_path=str(file_path),
-                    line_num=node.lineno,
-                    metric_name=metric_name,
-                    method=method,
-                    stats_obj=stats_obj or "",
-                    is_dynamic="{" in metric_name,
+            for metric_name in metric_names:
+                metrics_found.append(
+                    MetricCall(
+                        file_path=str(file_path),
+                        line_num=node.lineno,
+                        metric_name=metric_name,
+                        method=method,
+                        stats_obj=stats_obj or "",
+                        is_dynamic="{" in metric_name,
+                    )
                 )
-            )
 
     return metrics_found
 
 
+def list_repository_python_files() -> list[Path]:
+    """Return all git-tracked, non-test Python files in the repository."""
+    result = subprocess.run(
+        ["git", "ls-files", "*.py"],
+        cwd=AIRFLOW_ROOT_PATH,
+        capture_output=True,
+        text=True,
+        check=True,
+    )
+    return [
+        AIRFLOW_ROOT_PATH / line
+        for line in result.stdout.splitlines()
+        if not EXCLUDED_TESTS_PATTERN.search(line)
+    ]
+
+
+def compute_unused_registry_entries(
+    code_metric_names: set[str], metrics_registry: dict[str, dict]
+) -> list[str]:
+    """Return the registry entry names that no code metric name matches."""
+    used_entries = set(INDIRECTLY_EMITTED_METRICS)
+    for metric_name in code_metric_names:
+        registry_metric_name = find_registry_match(metric_name, metrics_registry)
+        if registry_metric_name is None:
+            continue
+        if registry_metric_name is _PREFIX_MATCHED:
+            used_entries.update(find_prefix_matched_registry_entries(metric_name, metrics_registry))
+        else:
+            used_entries.add(registry_metric_name)
+    return sorted(set(metrics_registry) - used_entries)
+
+
 def main() -> None:
     parser = argparse.ArgumentParser(
         description="Check that metrics in the codebase are in sync with the metrics registry YAML file."
@@ -262,11 +333,18 @@ def main() -> None:
         if mismatched:
             metrics_with_type_mismatch[name] = mismatched
 
-    # There is no point in checking whether the metrics exist in the YAML but not in the code,
-    # because the script is comparing the entire YAML against certain files at a time.
-    # For that to work, the script would have to run against all project files EVERY TIME.
+    # Violation 3: the metric exists in the registry but nowhere in the code. The hook only
+    # receives the changed files, so this check scans all git-tracked Python files itself.
+    all_code_metric_names = {
+        call.metric_name
+        for repo_file_path in list_repository_python_files()
+        for call in scan_file_for_metrics(repo_file_path)
+    }
+    unused_registry_entries = compute_unused_registry_entries(all_code_metric_names, metrics_registry)
 
-    total_violations = len(metrics_not_in_registry) + len(metrics_with_type_mismatch)
+    total_violations = (
+        len(metrics_not_in_registry) + len(metrics_with_type_mismatch) + len(unused_registry_entries)
+    )
 
     if total_violations:
         console.print(f"[red]Found {total_violations} violation(s).[/red]")
@@ -301,6 +379,18 @@ def main() -> None:
             console.print("    [yellow]Fix the type mismatch in either the code or the registry.[/yellow]")
             console.print()
 
+        if unused_registry_entries:
+            console.print(
+                f"    [red]-> {len(unused_registry_entries)} metric(s) found in the registry YAML but not emitted anywhere in the code:[/red]"
+            )
+            for metric_name in unused_registry_entries:
+                console.print(f"        [green]{metric_name}[/green]")
+            console.print(
+                "    [yellow]Remove them from the registry, or if they are emitted through a variable "
+                "the scan cannot resolve, add them to INDIRECTLY_EMITTED_METRICS in this script.[/yellow]"
+            )
+            console.print()
+
         sys.exit(1)
 
 

scripts/tests/ci/prek/test_check_metrics_synced_with_the_registry.py

@@ -23,7 +23,10 @@
 import pytest
 from ci.prek.check_metrics_synced_with_the_registry import (
     _PREFIX_MATCHED,
+    compute_unused_registry_entries,
     extract_metric_name_from_ast_node,
+    extract_metric_names_from_ast_node,
+    find_prefix_matched_registry_entries,
     find_registry_match,
     get_stats_obj_name,
     normalize_metric_name,
@@ -180,6 +183,100 @@ def test_extract_metric_name_from_ast_node(code: str, expected_result):
     assert extract_metric_name_from_ast_node(node) == expected_result
 
 
+@pytest.mark.parametrize(
+    "code, expected_result",
+    [
+        pytest.param('"scheduler_heartbeat"', ["scheduler_heartbeat"], id="static_string_single_name"),
+        pytest.param(
+            '"connection_test.success" if success else "connection_test.failed"',
+            ["connection_test.success", "connection_test.failed"],
+            id="conditional_expression_both_branches",
+        ),
+        pytest.param(
+            '"connection_test.success" if success else get_name()',
+            ["connection_test.success"],
+            id="conditional_expression_unresolvable_branch_dropped",
+        ),
+        pytest.param(
+            '"a" if x else ("b" if y else "c")',
+            ["a", "b", "c"],
+            id="nested_conditional_expression",
+        ),
+        pytest.param("some_variable", [], id="unresolvable_name_returns_empty_list"),
+    ],
+)
+def test_extract_metric_names_from_ast_node(code: str, expected_result):
+    node = ast.parse(code, mode="eval").body
+    assert extract_metric_names_from_ast_node(node) == expected_result
+
+
+@pytest.mark.parametrize(
+    "metric_name, expected_result",
+    [
+        pytest.param(
+            "ti.{state}",
+            ["ti.scheduled", "ti.queued", "ti.start.{dag_id}.{task_id}"],
+            id="base_prefix_matches_multiple_entries",
+        ),
+        pytest.param("dagrun.duration.{state}", ["dagrun.duration.success"], id="dotted_base_prefix"),
+        pytest.param("non.existent.{var}", [], id="no_prefix_match_returns_empty_list"),
+    ],
+)
+def test_find_prefix_matched_registry_entries(metric_name, expected_result):
+    assert find_prefix_matched_registry_entries(metric_name, METRICS_REGISTRY) == expected_result
+
+
+# 'executor.open_slots' is in INDIRECTLY_EMITTED_METRICS, so it is never reported as unused.
+@pytest.mark.parametrize(
+    "code_metric_names, expected_unused",
+    [
+        pytest.param(
+            set(),
+            [
+                "dagrun.duration.success",
+                "pool.open_slots",
+                "scheduler.heartbeat",
+                "task.duration",
+                "ti.queued",
+                "ti.scheduled",
+                "ti.start.{dag_id}.{task_id}",
+            ],
+            id="no_code_metrics_reports_all_but_indirectly_emitted",
+        ),
+        pytest.param(
+            {"scheduler.heartbeat", "dag.{x}.{y}.duration", "unknown.metric"},
+            [
+                "dagrun.duration.success",
+                "pool.open_slots",
+                "ti.queued",
+                "ti.scheduled",
+                "ti.start.{dag_id}.{task_id}",
+            ],
+            id="exact_and_legacy_matches_mark_entries_used",
+        ),
+        pytest.param(
+            {"ti.{state}"},
+            ["dagrun.duration.success", "pool.open_slots", "scheduler.heartbeat", "task.duration"],
+            id="prefix_match_marks_all_prefix_entries_used",
+        ),
+        pytest.param(
+            {"pool.open_slots.{my_pool}"},
+            [
+                "dagrun.duration.success",
+                "scheduler.heartbeat",
+                "task.duration",
+                "ti.queued",
+                "ti.scheduled",
+                "ti.start.{dag_id}.{task_id}",
+            ],
+            id="legacy_name_structure_match_marks_entry_used",
+        ),
+    ],
+)
+def test_compute_unused_registry_entries(code_metric_names, expected_unused):
+    assert compute_unused_registry_entries(code_metric_names, METRICS_REGISTRY) == expected_unused
+
+
 @pytest.fixture
 def code_to_py_file(tmp_path):
     """Write python source code to a tmp file and return its path."""
@@ -237,6 +334,14 @@ def _write(code: str) -> Path:
             [{"stats_obj": "stats"}],
             id="self_stats_attribute",
         ),
+        pytest.param(
+            'stats.incr("connection_test.success" if success else "connection_test.failed")',
+            [
+                {"metric_name": "connection_test.success", "method": "incr"},
+                {"metric_name": "connection_test.failed", "method": "incr"},
+            ],
+            id="conditional_expression_yields_call_per_branch",
+        ),
         pytest.param('metrics.incr("triggerer_heartbeat")', [], id="unknown_stats_object_ignored"),
         pytest.param("Stats.incr(get_metric_name())", [], id="unresolvable_metric_name_skipped"),
         pytest.param("def foo(:\n    pass\n", [], id="syntax_error_returns_empty"),

shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml

@@ -39,14 +39,6 @@ metrics:
     legacy_name: "-"
     name_variables: ["job_name"]
 
-  - name: "local_task_job.task_exit"
-    description: "Number of ``LocalTaskJob`` terminations with a ``{return_code}`` while
-    running a task ``{task_id}`` of a Dag  ``{dag_id}``.
-    Metric with job_id, dag_id, task_id and return_code tagging."
-    type: "counter"
-    legacy_name: "local_task_job.task_exit.{job_id}.{dag_id}.{task_id}.{return_code}"
-    name_variables: ["job_id", "dag_id", "task_id", "return_code"]
-
   - name: "operator_failures"
     description: "Operator ``{operator_name}`` failures."
     type: "counter"
@@ -162,12 +154,6 @@ metrics:
     legacy_name: "-"
     name_variables: []
 
-  - name: "dag_file_processor_timeouts"
-    description: "(DEPRECATED) same behavior as ``dag_processing.processor_timeouts``"
-    type: "counter"
-    legacy_name: "-"
-    name_variables: []
-
   - name: "scheduler.tasks.killed_externally"
     description: "Number of tasks killed externally. Metric with dag_id and task_id tagging."
     type: "counter"
@@ -433,12 +419,6 @@ metrics:
     legacy_name: "dag_processing.last_run.seconds_ago.{file_name}"
     name_variables: ["file_name"]
 
-  - name: "dag_processing.last_num_of_db_queries.{dag_file}"
-    description: "Number of queries to Airflow database during parsing per ``{dag_file}``"
-    type: "gauge"
-    legacy_name: "-"
-    name_variables: ["dag_file"]
-
   - name: "scheduler.tasks.starving"
     description: "Number of tasks that cannot be scheduled because of no open slot in pool"
     type: "gauge"