[airflow] 04/23: Correctly send timing metrics when using dogstatsd (fix schedule_delay metric) (#19973)

[je...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

jedcunningham pushed a commit to branch v2-2-test
in repository https://gitbox.apache.org/repos/asf/airflow.git

commit b05722e8b98963f65cd5f615d4390bb6377ba660
Author: Lutz Ostkamp <35...@users.noreply.github.com>
AuthorDate: Wed Dec 15 11:42:14 2021 +0100

    Correctly send timing metrics when using dogstatsd (fix schedule_delay metric) (#19973)
    
    (cherry picked from commit 5d405d9cda0b88909e6b726769381044477f4678)
---
 airflow/stats.py         | 9 ++++++---
 tests/core/test_stats.py | 5 +++++
 2 files changed, 11 insertions(+), 3 deletions(-)

diff --git a/airflow/stats.py b/airflow/stats.py
index 0a7004d..d1c4da6 100644
--- a/airflow/stats.py
+++ b/airflow/stats.py
@@ -16,13 +16,14 @@
 # specific language governing permissions and limitations
 # under the License.
 
+import datetime
 import logging
 import socket
 import string
 import textwrap
 import time
 from functools import wraps
-from typing import TYPE_CHECKING, Callable, Optional, TypeVar, cast
+from typing import TYPE_CHECKING, Callable, List, Optional, TypeVar, Union, cast
 
 from airflow.configuration import conf
 from airflow.exceptions import AirflowConfigException, InvalidStatsNameException
@@ -65,7 +66,7 @@ class StatsLogger(Protocol):
         """Gauge stat"""
 
     @classmethod
-    def timing(cls, stat: str, dt) -> None:
+    def timing(cls, stat: str, dt: Union[float, datetime.timedelta]) -> None:
         """Stats timing"""
 
     @classmethod
@@ -331,10 +332,12 @@ class SafeDogStatsdLogger:
         return None
 
     @validate_stat
-    def timing(self, stat, dt, tags=None):
+    def timing(self, stat, dt: Union[float, datetime.timedelta], tags: Optional[List[str]] = None):
         """Stats timing"""
         if self.allow_list_validator.test(stat):
             tags = tags or []
+            if isinstance(dt, datetime.timedelta):
+                dt = dt.total_seconds()
             return self.dogstatsd.timing(metric=stat, value=dt, tags=tags)
         return None
 
diff --git a/tests/core/test_stats.py b/tests/core/test_stats.py
index 83169e2..c401a2f 100644
--- a/tests/core/test_stats.py
+++ b/tests/core/test_stats.py
@@ -181,9 +181,14 @@ class TestDogStats(unittest.TestCase):
         self.dogstatsd_client.timed.assert_not_called()
 
     def test_timing(self):
+        import datetime
+
         self.dogstatsd.timing("dummy_timer", 123)
         self.dogstatsd_client.timing.assert_called_once_with(metric='dummy_timer', value=123, tags=[])
 
+        self.dogstatsd.timing("dummy_timer", datetime.timedelta(seconds=123))
+        self.dogstatsd_client.timing.assert_called_with(metric='dummy_timer', value=123.0, tags=[])
+
     def test_gauge(self):
         self.dogstatsd.gauge("dummy", 123)
         self.dogstatsd_client.gauge.assert_called_once_with(metric='dummy', sample_rate=1, value=123, tags=[])