Implemented get_name in StatsLogger, updated Otel and StatsD

airflow/metrics/base_stats_logger.py

@@ -23,6 +23,8 @@
 from airflow.typing_compat import Protocol
 
 if TYPE_CHECKING:
+    from opentelemetry.util.types import Attributes
+
     from airflow.metrics.protocols import DeltaType, TimerProtocol
 
 
@@ -31,55 +33,53 @@ class StatsLogger(Protocol):
 
     instance: StatsLogger | NoStatsLogger | None = None
 
-    @classmethod
     def incr(
-        cls,
-        stat: str,
+        self,
+        metric_name: str,
         count: int = 1,
         rate: int | float = 1,
         *,
         tags: dict[str, Any] | None = None,
     ) -> None:
-        """Increment stat."""
+        """Increment metric_name."""
 
-    @classmethod
     def decr(
-        cls,
-        stat: str,
+        self,
+        metric_name: str,
         count: int = 1,
         rate: int | float = 1,
         *,
         tags: dict[str, Any] | None = None,
     ) -> None:
-        """Decrement stat."""
+        """Decrement metric_name."""
 
-    @classmethod
     def gauge(
-        cls,
-        stat: str,
+        self,
+        metric_name: str,
         value: float,
         rate: int | float = 1,
         delta: bool = False,
         *,
         tags: dict[str, Any] | None = None,
     ) -> None:
-        """Gauge stat."""
+        """Gauge metric_name."""
 
-    @classmethod
     def timing(
-        cls,
-        stat: str,
+        self,
+        metric_name: str,
         dt: DeltaType | None,
         *,
         tags: dict[str, Any] | None = None,
     ) -> None:
         """Stats timing."""
 
-    @classmethod
-    def timer(cls, *args, **kwargs) -> TimerProtocol:
+    def timer(self, *args, **kwargs) -> TimerProtocol:
         """Timer metric that can be cancelled."""
         raise NotImplementedError()
 
+    def get_name(self, metric_name: str, tags: Attributes | None = None) -> str:
+        raise NotImplementedError()
+
 
 class NoStatsLogger:
     """If no StatsLogger is configured, NoStatsLogger is used as a fallback."""

airflow/metrics/otel_logger.py

@@ -31,7 +31,8 @@
 from opentelemetry.sdk.resources import HOST_NAME, SERVICE_NAME, Resource
 
 from airflow.configuration import conf
-from airflow.exceptions import AirflowProviderDeprecationWarning
+from airflow.exceptions import AirflowProviderDeprecationWarning, InvalidStatsNameException
+from airflow.metrics.base_stats_logger import StatsLogger
 from airflow.metrics.protocols import Timer
 from airflow.metrics.validators import (
     OTEL_NAME_MAX_LENGTH,
@@ -173,7 +174,7 @@ def stop(self, send: bool = True) -> None:
             )
 
 
-class SafeOtelLogger:
+class SafeOtelLogger(StatsLogger):
     """Otel Logger."""
 
     def __init__(
@@ -190,15 +191,15 @@ def __init__(
 
     def incr(
         self,
-        stat: str,
+        metric_name: str,
         count: int = 1,
         rate: float = 1,
         tags: Attributes = None,
     ):
         """
         Increment stat by count.
 
-        :param stat: The name of the stat to increment.
+        :param metric_name: The name of the metric to increment.
         :param count: A positive integer to add to the current value of stat.
         :param rate: value between 0 and 1 that represents the sample rate at
             which the metric is going to be emitted.
@@ -209,22 +210,23 @@ def incr(
         if count < 0:
             raise ValueError("count must be a positive value.")
 
-        if self.metrics_validator.test(stat) and name_is_otel_safe(self.prefix, stat):
-            counter = self.metrics_map.get_counter(full_name(prefix=self.prefix, name=stat), attributes=tags)
+        if self.metrics_validator.test(metric_name):
+            full_metric_name = self.get_name(metric_name, tags)
+            counter = self.metrics_map.get_counter(full_metric_name, attributes=tags)
             counter.add(count, attributes=tags)
             return counter
 
     def decr(
         self,
-        stat: str,
+        metric_name: str,
         count: int = 1,
         rate: float = 1,
         tags: Attributes = None,
     ):
         """
         Decrement stat by count.
 
-        :param stat: The name of the stat to decrement.
+        :param metric_name: The name of the metric to decrement.
         :param count: A positive integer to subtract from current value of stat.
         :param rate: value between 0 and 1 that represents the sample rate at
             which the metric is going to be emitted.
@@ -235,14 +237,15 @@ def decr(
         if count < 0:
             raise ValueError("count must be a positive value.")
 
-        if self.metrics_validator.test(stat) and name_is_otel_safe(self.prefix, stat):
-            counter = self.metrics_map.get_counter(full_name(prefix=self.prefix, name=stat))
+        if self.metrics_validator.test(metric_name):
+            full_metric_name = self.get_name(metric_name, tags)
+            counter = self.metrics_map.get_counter(full_metric_name)
             counter.add(-count, attributes=tags)
             return counter
 
     def gauge(
         self,
-        stat: str,
+        metric_name: str,
         value: int | float,
         rate: float = 1,
         delta: bool = False,
@@ -253,7 +256,7 @@ def gauge(
         """
         Record a new value for a Gauge.
 
-        :param stat: The name of the stat to update.
+        :param metric_name: The name of the metric to update.
         :param value: The new value of stat, either a float or an int.
         :param rate: value between 0 and 1 that represents the sample rate at
             which the metric is going to be emitted.
@@ -271,24 +274,29 @@ def gauge(
                 full_name(prefix=self.prefix, name=back_compat_name), value, delta, tags
             )
 
-        if self.metrics_validator.test(stat):
-            self.metrics_map.set_gauge_value(full_name(prefix=self.prefix, name=stat), value, delta, tags)
+        if self.metrics_validator.test(metric_name):
+            full_metric_name = self.get_name(metric_name, tags)
+            self.metrics_map.set_gauge_value(full_metric_name, value, delta, tags)
 
     def timing(
         self,
-        stat: str,
-        dt: DeltaType,
+        metric_name: str,
+        dt: DeltaType | None,
         *,
         tags: Attributes = None,
     ) -> None:
         """OTel does not have a native timer, stored as a Gauge whose value is number of seconds elapsed."""
-        if self.metrics_validator.test(stat) and name_is_otel_safe(self.prefix, stat):
+        if dt is None:
+            log.warning("The duration (dt) cannot be None. Skipping timing update.")
+            return
+        if self.metrics_validator.test(metric_name):
+            full_metric_name = self.get_name(metric_name, tags)
             if isinstance(dt, datetime.timedelta):
                 if metrics_consistency_on:
                     dt = dt.total_seconds() * 1000.0
                 else:
                     dt = dt.total_seconds()
-            self.metrics_map.set_gauge_value(full_name(prefix=self.prefix, name=stat), float(dt), False, tags)
+            self.metrics_map.set_gauge_value(full_metric_name, float(dt), False, tags)
 
     def timer(
         self,
@@ -300,6 +308,18 @@ def timer(
         """Timer context manager returns the duration and can be cancelled."""
         return _OtelTimer(self, stat, tags)
 
+    def get_name(self, metric_name: str, tags: Attributes | None = None) -> str:
+        """Generate an OTel-safe metric name with the prefix and delimiter."""
+        if not metric_name:
+            raise InvalidStatsNameException("The stat name cannot be None or an empty string.")
+
+        base_metric_name = metric_name
+
+        if not name_is_otel_safe(self.prefix, base_metric_name):
+            raise ValueError(f"Metric name `{base_metric_name}` is not OTel-safe.")
+
+        return full_name(name=base_metric_name, prefix=self.prefix)
+
 
 class MetricsMap:
     """Stores Otel Instruments."""