Repository: apache/airflow

Issue:
Enabling Datadog to tag metrics results in AttributeError
**Apache Airflow version**: 2.0.1
**Python version**: 3.8
**Cloud provider or hardware configuration**: AWS

**What happened**:
In order to add tags to [Airflow metrics,](https://airflow.apache.org/docs/apache-airflow/stable/logging-monitoring/metrics.html), it's required to set `AIRFLOW__METRICS__STATSD_DATADOG_ENABLED` to `True` and add tags in the `AIRFLOW__METRICS__STATSD_DATADOG_TAGS` variable. We were routing our statsd metrics to Datadog anyway, so this should theoretically have not changed anything other than the addition of any specified tags.
Setting the environment variable `AIRFLOW__METRICS__STATSD_DATADOG_ENABLED` to `True` (along with the other required statsd connection variables) results in the following error, which causes the process to terminate. This is from the scheduler, but this would apply anywhere that `Stats.timer()` is being called.
```
AttributeError: 'DogStatsd' object has no attribute 'timer'
return Timer(self.dogstatsd.timer(stat, *args, tags=tags, **kwargs))
File "/usr/local/lib/python3.8/site-packages/airflow/stats.py", line 345, in timer
return fn(_self, stat, *args, **kwargs)
File "/usr/local/lib/python3.8/site-packages/airflow/stats.py", line 233, in wrapper
timer = Stats.timer('scheduler.critical_section_duration')
File "/usr/local/lib/python3.8/site-packages/airflow/jobs/scheduler_job.py", line 1538, in _do_scheduling
num_queued_tis = self._do_scheduling(session)
File "/usr/local/lib/python3.8/site-packages/airflow/jobs/scheduler_job.py", line 1382, in _run_scheduler_loop
self._run_scheduler_loop()
File "/usr/local/lib/python3.8/site-packages/airflow/jobs/scheduler_job.py", line 1280, in _execute
Traceback (most recent call last):
```

**What you expected to happen**:
The same default Airflow metrics get sent by connecting to datadog, tagged with the metrics specified in `AIRFLOW__METRICS__STATSD_DATADOG_TAGS`.

**What do you think went wrong?**:
There is a bug in the implementation of the `Timer` method of `SafeDogStatsdLogger`. https://github.com/apache/airflow/blob/master/airflow/stats.py#L341-L347
`DogStatsd` has no method called `timer`. Instead it should be `timed`: https://datadogpy.readthedocs.io/en/latest/#datadog.dogstatsd.base.DogStatsd.timed

**How to reproduce it**:
Set the environment variables (or their respective config values) `AIRFLOW__METRICS__STATSD_ON`, `AIRFLOW__METRICS__STATSD_HOST`, `AIRFLOW__METRICS__STATSD_PORT`, and then set `AIRFLOW__METRICS__STATSD_DATADOG_ENABLED` to `True` and start up Airflow.


**Anything else we need to know**:
How often does this problem occur? Every time




Discussion:
Thanks for opening your first issue here! Be sure to follow the issue template!


Resolving patch:
diff --git a/airflow/stats.py b/airflow/stats.py
--- a/airflow/stats.py
+++ b/airflow/stats.py
@@ -343,7 +343,7 @@ def timer(self, stat=None, *args, tags=None, **kwargs):
         """Timer metric that can be cancelled"""
         if stat and self.allow_list_validator.test(stat):
             tags = tags or []
-            return Timer(self.dogstatsd.timer(stat, *args, tags=tags, **kwargs))
+            return Timer(self.dogstatsd.timed(stat, *args, tags=tags, **kwargs))
         return Timer()
 
 


Test patch:
diff --git a/tests/core/test_stats.py b/tests/core/test_stats.py
--- a/tests/core/test_stats.py
+++ b/tests/core/test_stats.py
@@ -31,17 +31,7 @@
 
 
 class CustomStatsd(statsd.StatsClient):
-    incr_calls = 0
-
-    def __init__(self, host=None, port=None, prefix=None):
-        super().__init__()
-
-    def incr(self, stat, count=1, rate=1):  # pylint: disable=unused-argument
-        CustomStatsd.incr_calls += 1
-
-    @classmethod
-    def _reset(cls):
-        cls.incr_calls = 0
+    pass
 
 
 class InvalidCustomStatsd:
@@ -50,25 +40,14 @@ class InvalidCustomStatsd:
     statsd.StatsClient.
     """
 
-    incr_calls = 0
-
     def __init__(self, host=None, port=None, prefix=None):
         pass
 
-    def incr(self, stat, count=1, rate=1):  # pylint: disable=unused-argument
-        InvalidCustomStatsd.incr_calls += 1
-
-    @classmethod
-    def _reset(cls):
-        cls.incr_calls = 0
-
 
 class TestStats(unittest.TestCase):
     def setUp(self):
-        self.statsd_client = Mock()
+        self.statsd_client = Mock(spec=statsd.StatsClient)
         self.stats = SafeStatsdLogger(self.statsd_client)
-        CustomStatsd._reset()
-        InvalidCustomStatsd._reset()
 
     def test_increment_counter_with_valid_name(self):
         self.stats.incr('test_stats_run')
@@ -86,49 +65,56 @@ def test_stat_name_must_only_include_allowed_characters(self):
         self.stats.incr('test/$tats')
         self.statsd_client.assert_not_called()
 
-    @conf_vars({('metrics', 'statsd_on'): 'True'})
-    @mock.patch("statsd.StatsClient")
-    def test_does_send_stats_using_statsd(self, mock_statsd):
-        importlib.reload(airflow.stats)
-        airflow.stats.Stats.incr("dummy_key")
-        mock_statsd.return_value.incr.assert_called_once_with('dummy_key', 1, 1)
+    def test_timer(self):
+        with self.stats.timer("dummy_timer"):
+            pass
+        self.statsd_client.timer.assert_called_once_with('dummy_timer')
 
-    @conf_vars({('metrics', 'statsd_on'): 'True'})
-    @mock.patch("datadog.DogStatsd")
-    def test_does_not_send_stats_using_dogstatsd(self, mock_dogstatsd):
-        importlib.reload(airflow.stats)
-        airflow.stats.Stats.incr("dummy_key")
-        mock_dogstatsd.return_value.assert_not_called()
+    def test_empty_timer(self):
+        with self.stats.timer():
+            pass
+        self.statsd_client.timer.assert_not_called()
 
-    @conf_vars(
-        {
-            ("metrics", "statsd_on"): "True",
-            ("metrics", "statsd_custom_client_path"): "tests.core.test_stats.CustomStatsd",
-        }
-    )
-    def test_load_custom_statsd_client(self):
+    def test_timing(self):
+        self.stats.timing("dummy_timer", 123)
+        self.statsd_client.timing.assert_called_once_with('dummy_timer', 123)
+
+    def test_gauge(self):
+        self.stats.gauge("dummy", 123)
+        self.statsd_client.gauge.assert_called_once_with('dummy', 123, 1, False)
+
+    def test_decr(self):
+        self.stats.decr("dummy")
+        self.statsd_client.decr.assert_called_once_with('dummy', 1, 1)
+
+    def test_enabled_by_config(self):
+        """Test that enabling this sets the right instance properties"""
+        with conf_vars({('metrics', 'statsd_on'): 'True'}):
+            importlib.reload(airflow.stats)
+            assert isinstance(airflow.stats.Stats.statsd, statsd.StatsClient)
+            assert not hasattr(airflow.stats.Stats, 'dogstatsd')
+        # Avoid side-effects
         importlib.reload(airflow.stats)
-        assert 'CustomStatsd' == type(airflow.stats.Stats.statsd).__name__  # noqa: E721
 
-    @conf_vars(
-        {
-            ("metrics", "statsd_on"): "True",
-            ("metrics", "statsd_custom_client_path"): "tests.core.test_stats.CustomStatsd",
-        }
-    )
-    def test_does_use_custom_statsd_client(self):
+    def test_load_custom_statsd_client(self):
+        with conf_vars(
+            {
+                ("metrics", "statsd_on"): "True",
+                ("metrics", "statsd_custom_client_path"): f"{__name__}.CustomStatsd",
+            }
+        ):
+            importlib.reload(airflow.stats)
+            assert isinstance(airflow.stats.Stats.statsd, CustomStatsd)
+        # Avoid side-effects
         importlib.reload(airflow.stats)
-        airflow.stats.Stats.incr("dummy_key")
-        assert airflow.stats.Stats.statsd.incr_calls == 1
 
-    @conf_vars(
-        {
-            ("metrics", "statsd_on"): "True",
-            ("metrics", "statsd_custom_client_path"): "tests.core.test_stats.InvalidCustomStatsd",
-        }
-    )
     def test_load_invalid_custom_stats_client(self):
-        with pytest.raises(
+        with conf_vars(
+            {
+                ("metrics", "statsd_on"): "True",
+                ("metrics", "statsd_custom_client_path"): f"{__name__}.InvalidCustomStatsd",
+            }
+        ), pytest.raises(
             AirflowConfigException,
             match=re.escape(
                 'Your custom Statsd client must extend the statsd.'
@@ -137,15 +123,15 @@ def test_load_invalid_custom_stats_client(self):
         ):
             importlib.reload(airflow.stats)
             airflow.stats.Stats.incr("dummy_key")
-
-    def tearDown(self) -> None:
-        # To avoid side-effect
         importlib.reload(airflow.stats)
 
 
 class TestDogStats(unittest.TestCase):
     def setUp(self):
-        self.dogstatsd_client = Mock()
+        pytest.importorskip('datadog')
+        from datadog import DogStatsd
+
+        self.dogstatsd_client = Mock(spec=DogStatsd)
         self.dogstatsd = SafeDogStatsdLogger(self.dogstatsd_client)
 
     def test_increment_counter_with_valid_name_with_dogstatsd(self):
@@ -166,48 +152,72 @@ def test_stat_name_must_only_include_allowed_characters_with_dogstatsd(self):
         self.dogstatsd.incr('test/$tats')
         self.dogstatsd_client.assert_not_called()
 
-    @conf_vars({('metrics', 'statsd_datadog_enabled'): 'True'})
-    @mock.patch("datadog.DogStatsd")
-    def test_does_send_stats_using_dogstatsd_when_dogstatsd_on(self, mock_dogstatsd):
-        importlib.reload(airflow.stats)
-        airflow.stats.Stats.incr("dummy_key")
-        mock_dogstatsd.return_value.increment.assert_called_once_with(
+    def test_does_send_stats_using_dogstatsd_when_dogstatsd_on(self):
+        self.dogstatsd.incr("dummy_key")
+        self.dogstatsd_client.increment.assert_called_once_with(
             metric='dummy_key', sample_rate=1, tags=[], value=1
         )
 
-    @conf_vars({('metrics', 'statsd_datadog_enabled'): 'True'})
-    @mock.patch("datadog.DogStatsd")
-    def test_does_send_stats_using_dogstatsd_with_tags(self, mock_dogstatsd):
-        importlib.reload(airflow.stats)
-        airflow.stats.Stats.incr("dummy_key", 1, 1, ['key1:value1', 'key2:value2'])
-        mock_dogstatsd.return_value.increment.assert_called_once_with(
+    def test_does_send_stats_using_dogstatsd_with_tags(self):
+        self.dogstatsd.incr("dummy_key", 1, 1, ['key1:value1', 'key2:value2'])
+        self.dogstatsd_client.increment.assert_called_once_with(
             metric='dummy_key', sample_rate=1, tags=['key1:value1', 'key2:value2'], value=1
         )
 
-    @conf_vars({('metrics', 'statsd_on'): 'True', ('metrics', 'statsd_datadog_enabled'): 'True'})
-    @mock.patch("datadog.DogStatsd")
-    def test_does_send_stats_using_dogstatsd_when_statsd_and_dogstatsd_both_on(self, mock_dogstatsd):
-        importlib.reload(airflow.stats)
-        airflow.stats.Stats.incr("dummy_key")
-        mock_dogstatsd.return_value.increment.assert_called_once_with(
+    def test_does_send_stats_using_dogstatsd_when_statsd_and_dogstatsd_both_on(self):
+        self.dogstatsd.incr("dummy_key")
+        self.dogstatsd_client.increment.assert_called_once_with(
             metric='dummy_key', sample_rate=1, tags=[], value=1
         )
 
-    @conf_vars({('metrics', 'statsd_on'): 'True', ('metrics', 'statsd_datadog_enabled'): 'True'})
-    @mock.patch("statsd.StatsClient")
-    def test_does_not_send_stats_using_statsd_when_statsd_and_dogstatsd_both_on(self, mock_statsd):
+    def test_timer(self):
+        with self.dogstatsd.timer("dummy_timer"):
+            pass
+        self.dogstatsd_client.timed.assert_called_once_with('dummy_timer', tags=[])
+
+    def test_empty_timer(self):
+        with self.dogstatsd.timer():
+            pass
+        self.dogstatsd_client.timed.assert_not_called()
+
+    def test_timing(self):
+        self.dogstatsd.timing("dummy_timer", 123)
+        self.dogstatsd_client.timing.assert_called_once_with(metric='dummy_timer', value=123, 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=[])
+
+    def test_decr(self):
+        self.dogstatsd.decr("dummy")
+        self.dogstatsd_client.decrement.assert_called_once_with(
+            metric='dummy', sample_rate=1, value=1, tags=[]
+        )
+
+    def test_enabled_by_config(self):
+        """Test that enabling this sets the right instance properties"""
+        from datadog import DogStatsd
+
+        with conf_vars({('metrics', 'statsd_datadog_enabled'): 'True'}):
+            importlib.reload(airflow.stats)
+            assert isinstance(airflow.stats.Stats.dogstatsd, DogStatsd)
+            assert not hasattr(airflow.stats.Stats, 'statsd')
+        # Avoid side-effects
         importlib.reload(airflow.stats)
-        airflow.stats.Stats.incr("dummy_key")
-        mock_statsd.return_value.assert_not_called()
 
-    def tearDown(self) -> None:
-        # To avoid side-effect
+    def test_does_not_send_stats_using_statsd_when_statsd_and_dogstatsd_both_on(self):
+        from datadog import DogStatsd
+
+        with conf_vars({('metrics', 'statsd_on'): 'True', ('metrics', 'statsd_datadog_enabled'): 'True'}):
+            importlib.reload(airflow.stats)
+            assert isinstance(airflow.stats.Stats.dogstatsd, DogStatsd)
+            assert not hasattr(airflow.stats.Stats, 'statsd')
         importlib.reload(airflow.stats)
 
 
 class TestStatsWithAllowList(unittest.TestCase):
     def setUp(self):
-        self.statsd_client = Mock()
+        self.statsd_client = Mock(spec=statsd.StatsClient)
         self.stats = SafeStatsdLogger(self.statsd_client, AllowListValidator("stats_one, stats_two"))
 
     def test_increment_counter_with_allowed_key(self):
@@ -225,7 +235,10 @@ def test_not_increment_counter_if_not_allowed(self):
 
 class TestDogStatsWithAllowList(unittest.TestCase):
     def setUp(self):
-        self.dogstatsd_client = Mock()
+        pytest.importorskip('datadog')
+        from datadog import DogStatsd
+
+        self.dogstatsd_client = Mock(speck=DogStatsd)
         self.dogstats = SafeDogStatsdLogger(self.dogstatsd_client, AllowListValidator("stats_one, stats_two"))
 
     def test_increment_counter_with_allowed_key(self):