This is an automated email from the ASF dual-hosted git repository.
amoghrajesh pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new c68e9f16325 Update the metrics registry hook to ban direct imports of
stats emitting module-level functions (#68675)
c68e9f16325 is described below
commit c68e9f163250efd2970f11d6b506a51350bcd818
Author: Christos Bisias <[email protected]>
AuthorDate: Wed Jul 15 15:03:45 2026 +0300
Update the metrics registry hook to ban direct imports of stats emitting
module-level functions (#68675)
---
.pre-commit-config.yaml | 2 +-
.../prek/check_metrics_synced_with_the_registry.py | 139 ++++++++++++++-
.../test_check_metrics_synced_with_the_registry.py | 190 +++++++++++++++++++++
3 files changed, 328 insertions(+), 3 deletions(-)
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index c5f8e587ddb..c3722a2347f 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -945,7 +945,7 @@ repos:
entry: ./scripts/ci/prek/check_metrics_synced_with_the_registry.py
language: python
files: \.py$
- exclude: ^(tests/|.*/tests/)
+ exclude:
^(tests/|.*/tests/|airflow-core/src/airflow/observability/stats\.py$|task-sdk/src/airflow/sdk/observability/stats\.py$)
pass_filenames: true
additional_dependencies: ["PyYAML>=6.0", "rich>=13.6.0"]
- id: check-boring-cyborg-configuration
diff --git a/scripts/ci/prek/check_metrics_synced_with_the_registry.py
b/scripts/ci/prek/check_metrics_synced_with_the_registry.py
index 94b5bc254de..3d42939b6ce 100644
--- a/scripts/ci/prek/check_metrics_synced_with_the_registry.py
+++ b/scripts/ci/prek/check_metrics_synced_with_the_registry.py
@@ -53,6 +53,16 @@ STATS_METHOD_TO_TYPE: dict[str, str] = {
STATS_OBJECTS = {"Stats", "stats"}
+# Stats module path suffix. Currently, there are the following possible module
paths:
+# ``airflow._shared.observability.metrics.stats``,
+# ``airflow.sdk._shared.observability.metrics.stats``, and
+# ``airflow_shared.observability.metrics.stats``.
+STATS_MODULE_SUFFIX = "observability.metrics.stats"
+
+# Names of exception classes that when present in a try-except, then it means
that
+# the import is used for a back-compat version check, and it should be ignored.
+IMPORT_ERROR_NAMES = {"ImportError", "ModuleNotFoundError"}
+
METRICS_REGISTRY_PATH = (
AIRFLOW_ROOT_PATH /
"shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml"
)
@@ -177,6 +187,47 @@ class MetricCall:
is_dynamic: bool
+@dataclass
+class DirectStatsImport:
+ file_path: str
+ line_num: int
+ module: str
+ imported_names: list[str]
+
+
+def _is_stats_module_path(module: str | None) -> bool:
+ return module is not None and module.endswith(STATS_MODULE_SUFFIX)
+
+
+def _except_handler_catches_expected_error(handler: ast.ExceptHandler) -> bool:
+ """Return True if this except handler catches only names in
``IMPORT_ERROR_NAMES``."""
+ exc_type = handler.type
+ if isinstance(exc_type, ast.Name):
+ return exc_type.id in IMPORT_ERROR_NAMES
+ if isinstance(exc_type, ast.Tuple):
+ return bool(exc_type.elts) and all(
+ isinstance(elt, ast.Name) and elt.id in IMPORT_ERROR_NAMES for elt
in exc_type.elts
+ )
+ return False
+
+
+def find_back_compat_check_import_ids(tree: ast.AST) -> set[int]:
+ """Return ``id()`` of every ImportFrom inside a try-except ImportError
block.
+
+ Imports inside such blocks are used under providers for back-compat checks.
+ """
+ back_compat_check_import_ids: set[int] = set()
+ for node in ast.walk(tree):
+ if isinstance(node, ast.Try) and any(
+ _except_handler_catches_expected_error(h) for h in node.handlers
+ ):
+ for stmt in node.body:
+ for child in ast.walk(stmt):
+ if isinstance(child, ast.ImportFrom):
+ back_compat_check_import_ids.add(id(child))
+ return back_compat_check_import_ids
+
+
def scan_file_for_metrics(file_path: Path) -> list[MetricCall]:
"""Return all Stats metric calls found in the provided file_path."""
try:
@@ -220,6 +271,53 @@ def scan_file_for_metrics(file_path: Path) ->
list[MetricCall]:
return metrics_found
+def scan_file_for_direct_stats_imports(file_path: Path) ->
list[DirectStatsImport]:
+ """Return direct imports of stats module-level functions.
+
+ Only the metric-emitting methods in ``STATS_METHOD_TO_TYPE`` (``incr``,
+ ``decr``, ``gauge``, ``timing``, ``timer``) are restricted; other names
+ exported from ``stats.py`` (helpers, the ``Stats`` shim, ``initialize``,
+ etc.) may still be imported directly.
+
+ Imports inside a ``try`` block that catches ``ImportError`` (or
+ ``ModuleNotFoundError``) are exempt: those are intentional
+ back-compat version checks whose success is the signal the code is checking
+ for, and rewriting them to namespace form would defeat the check.
+ """
+ try:
+ source = file_path.read_text(encoding="utf-8")
+ tree = ast.parse(source, filename=str(file_path))
+ except (OSError, UnicodeDecodeError, SyntaxError):
+ return []
+
+ # Imports wrapped with try-except. These should be ignored.
+ back_compat_check_ids = find_back_compat_check_import_ids(tree)
+
+ violations: list[DirectStatsImport] = []
+ for node in ast.walk(tree):
+ if not isinstance(node, ast.ImportFrom):
+ continue
+ if not _is_stats_module_path(node.module):
+ continue
+ if id(node) in back_compat_check_ids:
+ continue
+ offending: list[str] = []
+ for alias in node.names:
+ if alias.name in STATS_METHOD_TO_TYPE:
+ offending.append(f"{alias.name} as {alias.asname}" if
alias.asname else alias.name)
+ if not offending:
+ continue
+ violations.append(
+ DirectStatsImport(
+ file_path=str(file_path),
+ line_num=node.lineno,
+ module=node.module or "",
+ imported_names=offending,
+ )
+ )
+ return violations
+
+
def main() -> None:
parser = argparse.ArgumentParser(
description="Check that metrics in the codebase are in sync with the
metrics registry YAML file."
@@ -232,9 +330,26 @@ def main() -> None:
metrics_registry = load_metrics_registry_yaml()
- # Collect all metric calls across all provided files.
+ # The scanner only sees `<obj>.<method>`, but since the class functions
were
+ # converted to module-level functions, if they are imported directly, the
scanner
+ # won't recognize them as a metric calls that need validation.
+ # E.g.
+ # This is recognizable for the scanner
+ # from airflow._shared.observability.metrics import stats
+ # stats.incr()
+ # This isn't
+ # from airflow._shared.observability.metrics.stats import incr
+ # incr()
+ #
+ # The simpler solution is to ban direct imports.
+ direct_stats_imports: list[DirectStatsImport] = []
code_metrics: dict[str, list[MetricCall]] = {}
for file_path in [Path(f) for f in args.files]:
+ file_direct_imports = scan_file_for_direct_stats_imports(file_path)
+ if file_direct_imports:
+ direct_stats_imports.extend(file_direct_imports)
+ # If there are direct imports, then skip checking for metric calls
in this file.
+ continue
for call in scan_file_for_metrics(file_path):
code_metrics.setdefault(call.metric_name, []).append(call)
@@ -266,7 +381,9 @@ def main() -> None:
# 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.
- 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(direct_stats_imports)
+ )
if total_violations:
console.print(f"[red]Found {total_violations} violation(s).[/red]")
@@ -301,6 +418,24 @@ def main() -> None:
console.print(" [yellow]Fix the type mismatch in either the
code or the registry.[/yellow]")
console.print()
+ if direct_stats_imports:
+ console.print(
+ f" [red]-> {len(direct_stats_imports)} direct import(s) of
stats functions (use the `stats` namespace instead):[/red]"
+ )
+ for violation in direct_stats_imports:
+ console.print(
+ f" [yellow]{violation.file_path}[/yellow] line
[yellow]{violation.line_num}[/yellow]: "
+ f"[magenta]from {violation.module} import {',
'.join(violation.imported_names)}[/magenta]"
+ )
+ console.print(
+ " [yellow]Replace direct imports with namespace access: "
+ "`from <parent>.observability.metrics import stats` and call
`stats.<method>(...)`.[/yellow]"
+ )
+ console.print(
+ " [yellow]Imports inside a `try` block that catches
`ImportError` are exempt (back-compat checks).[/yellow]"
+ )
+ console.print()
+
sys.exit(1)
diff --git
a/scripts/tests/ci/prek/test_check_metrics_synced_with_the_registry.py
b/scripts/tests/ci/prek/test_check_metrics_synced_with_the_registry.py
index 98f7f1ee256..fdd91fb4dce 100644
--- a/scripts/tests/ci/prek/test_check_metrics_synced_with_the_registry.py
+++ b/scripts/tests/ci/prek/test_check_metrics_synced_with_the_registry.py
@@ -23,10 +23,13 @@ from pathlib import Path
import pytest
from ci.prek.check_metrics_synced_with_the_registry import (
_PREFIX_MATCHED,
+ _except_handler_catches_expected_error,
+ _is_stats_module_path,
extract_metric_name_from_ast_node,
find_registry_match,
get_stats_obj_name,
normalize_metric_name,
+ scan_file_for_direct_stats_imports,
scan_file_for_metrics,
)
@@ -275,3 +278,190 @@ def test_scan_file_with_multiple_calls(code_to_py_file):
def test_scan_file_nonexistent_file_returns_empty(tmp_path):
assert scan_file_for_metrics(tmp_path / "non_existent.py") == []
+
+
[email protected](
+ "code, expected_imports",
+ [
+ pytest.param(
+ "from airflow._shared.observability.metrics.stats import
incr\nincr('foo')",
+ [{"imported_names": ["incr"], "line_num": 1}],
+ id="direct_import_of_metric_function",
+ ),
+ pytest.param(
+ "from airflow.sdk._shared.observability.metrics.stats import
gauge\n",
+ [{"imported_names": ["gauge"]}],
+ id="direct_import_via_sdk_path",
+ ),
+ pytest.param(
+ "from airflow_shared.observability.metrics.stats import timer\n",
+ [{"imported_names": ["timer"]}],
+ id="direct_import_via_shared_path",
+ ),
+ pytest.param(
+ "from airflow._shared.observability.metrics.stats import incr as
foo\n",
+ [{"imported_names": ["incr as foo"]}],
+ id="aliased_direct_import",
+ ),
+ pytest.param(
+ "from airflow._shared.observability.metrics.stats import incr,
gauge, timing\n",
+ [{"imported_names": ["incr", "gauge", "timing"]}],
+ id="multiple_methods_one_statement",
+ ),
+ pytest.param(
+ "from airflow._shared.observability.metrics import stats\n",
+ [],
+ id="namespace_import_allowed",
+ ),
+ pytest.param(
+ "from airflow._shared.observability.metrics.stats import
normalize_name_for_stats\n",
+ [],
+ id="non_metric_function_allowed",
+ ),
+ pytest.param(
+ "from airflow._shared.observability.metrics.stats import Stats\n",
+ [],
+ id="class_shim_allowed",
+ ),
+ pytest.param(
+ "from airflow._shared.observability.metrics.stats import incr,
normalize_name_for_stats\n",
+ [{"imported_names": ["incr"]}],
+ id="only_metric_function_reported_from_mixed_import",
+ ),
+ pytest.param(
+ "try:\n"
+ " from airflow._shared.observability.metrics.stats import gauge
# noqa: F401\n"
+ "except ImportError:\n"
+ " gauge = None\n",
+ [],
+ id="exempt_inside_try_except_import_error",
+ ),
+ pytest.param(
+ "try:\n"
+ " from airflow._shared.observability.metrics.stats import
gauge\n"
+ "except ModuleNotFoundError:\n"
+ " gauge = None\n",
+ [],
+ id="exempt_inside_try_except_module_not_found_error",
+ ),
+ pytest.param(
+ "try:\n"
+ " from airflow._shared.observability.metrics.stats import
gauge\n"
+ "except (ImportError, ModuleNotFoundError):\n"
+ " gauge = None\n",
+ [],
+ id="exempt_inside_try_except_tuple_of_only_expected",
+ ),
+ pytest.param(
+ "try:\n"
+ " from airflow._shared.observability.metrics.stats import
gauge\n"
+ "except (ImportError, OSError):\n"
+ " gauge = None\n",
+ [{"imported_names": ["gauge"]}],
+ id="not_exempt_for_tuple_mixing_expected_and_unrelated",
+ ),
+ pytest.param(
+ "try:\n"
+ " from airflow._shared.observability.metrics.stats import
gauge\n"
+ "except ValueError:\n"
+ " gauge = None\n",
+ [{"imported_names": ["gauge"]}],
+ id="not_exempt_for_unrelated_exception",
+ ),
+ pytest.param(
+ "try:\n"
+ " from airflow._shared.observability.metrics.stats import
gauge\n"
+ "except:\n"
+ " gauge = None\n",
+ [{"imported_names": ["gauge"]}],
+ id="not_exempt_for_bare_except",
+ ),
+ pytest.param(
+ "from some.unrelated.module import incr\n",
+ [],
+ id="unrelated_module_import_ignored",
+ ),
+ ],
+)
+def test_scan_file_for_direct_stats_imports(code_to_py_file, code,
expected_imports):
+ violations = scan_file_for_direct_stats_imports(code_to_py_file(code))
+ assert len(violations) == len(expected_imports)
+ for violation, expected in zip(violations, expected_imports):
+ for field, value in expected.items():
+ assert getattr(violation, field) == value
+
+
+def test_scan_file_for_direct_stats_imports_records_module(code_to_py_file):
+ path = code_to_py_file("from airflow._shared.observability.metrics.stats
import incr\n")
+ violations = scan_file_for_direct_stats_imports(path)
+ assert len(violations) == 1
+ assert violations[0].module ==
"airflow._shared.observability.metrics.stats"
+ assert violations[0].file_path == str(path)
+
+
+def
test_scan_file_for_direct_stats_imports_nonexistent_file_returns_empty(tmp_path):
+ assert scan_file_for_direct_stats_imports(tmp_path / "non_existent.py") ==
[]
+
+
[email protected](
+ "module, expected_bool_result",
+ [
+ pytest.param(None, False, id="none_returns_false"),
+ pytest.param("", False, id="empty_string_returns_false"),
+ pytest.param("observability.metrics.stats", True,
id="bare_suffix_matches"),
+ pytest.param(
+ "airflow._shared.observability.metrics.stats", True,
id="airflow_core_shared_path_matches"
+ ),
+ pytest.param(
+ "airflow.sdk._shared.observability.metrics.stats", True,
id="task_sdk_shared_path_matches"
+ ),
+ pytest.param("airflow_shared.observability.metrics.stats", True,
id="shared_path_matches"),
+ pytest.param(
+ "airflow._shared.observability.metrics.statsd_logger",
+ False,
+ id="module_with_stats_prefix_rejected",
+ ),
+ pytest.param(
+ "airflow._shared.observability.metrics.base_stats_logger",
+ False,
+ id="module_without_stats_at_tail_rejected",
+ ),
+ pytest.param(
+ "airflow._shared.observability.metrics", False,
id="parent_module_without_stats_rejected"
+ ),
+ pytest.param("some.random.module", False, id="random_module_rejected"),
+ ],
+)
+def test_is_stats_module_path(module, expected_bool_result):
+ assert _is_stats_module_path(module) is expected_bool_result
+
+
+def _create_ast_except_handler(exception_clause: str) -> ast.ExceptHandler:
+ """Parse a ``try`` block with the given except clause and return its
handler."""
+ suffix = f" {exception_clause}" if exception_clause else ""
+ tree = ast.parse(f"try:\n pass\nexcept{suffix}:\n pass\n")
+ return tree.body[0].handlers[0] # type: ignore[attr-defined]
+
+
[email protected](
+ "exception_clause, expected_bool_result",
+ [
+ pytest.param("ImportError", True, id="import_error_matches"),
+ pytest.param("ModuleNotFoundError", True,
id="module_not_found_error_matches"),
+ pytest.param("(ImportError, ModuleNotFoundError)", True,
id="tuple_of_only_expected_matches"),
+ pytest.param("(ImportError, OSError)", False,
id="tuple_with_unrelated_error_rejected1"),
+ pytest.param(
+ "(OSError, ModuleNotFoundError)",
+ False,
+ id="tuple_with_unrelated_error_rejected2",
+ ),
+ pytest.param("OSError", False,
id="unrelated_single_exception_rejected"),
+ pytest.param("(OSError, ValueError)", False,
id="tuple_of_unrelated_rejected"),
+ pytest.param("Exception", False, id="broad_exception_rejected"),
+ pytest.param("BaseException", False, id="base_exception_rejected"),
+ pytest.param("", False, id="empty_except_rejected"),
+ ],
+)
+def test_except_handler_catches_expected_error(exception_clause: str,
expected_bool_result):
+ handler = _create_ast_except_handler(exception_clause)
+ assert _except_handler_catches_expected_error(handler) is
expected_bool_result
