Add ability to submit time deltas to database query utility

datadog_checks_base/datadog_checks/base/utils/db/transform.py

@@ -4,12 +4,13 @@
 from __future__ import division
 
 import re
+from datetime import datetime
 
 from ... import is_affirmative
 from ...constants import ServiceCheck
 from .. import constants
 from ..common import compute_percent, total_time_to_temporal_percent
-from .utils import create_extra_transformer
+from .utils import create_extra_transformer, normalize_datetime
 
 # Used for the user-defined `expression`s
 ALLOWED_GLOBALS = {
@@ -96,6 +97,28 @@ def service_check(_, value, **kwargs):
     return service_check
 
 
+def get_time_elapsed(transformers, column_name, **modifiers):
+    time_format = modifiers.pop('format', 'native')
+    if not isinstance(time_format, str):
+        raise ValueError('the `format` parameter must be a string')
+
+    gauge = transformers['gauge'](transformers, column_name, **modifiers)
+
+    if time_format == 'native':
+
+        def time_elapsed(_, value, **kwargs):
+            value = normalize_datetime(value)
+            gauge(_, (datetime.now(value.tzinfo) - value).total_seconds(), **kwargs)
+
+    else:
+
+        def time_elapsed(_, value, **kwargs):
+            value = normalize_datetime(datetime.strptime(value, time_format))
+            gauge(_, (datetime.now(value.tzinfo) - value).total_seconds(), **kwargs)
+
+    return time_elapsed
+
+
 def get_expression(transformers, name, **modifiers):
     available_sources = modifiers.pop('sources')
 
@@ -183,6 +206,7 @@ def percent(sources, **kwargs):
     'tag': get_tag,
     'match': get_match,
     'service_check': get_service_check,
+    'time_elapsed': get_time_elapsed,
 }
 
 EXTRA_TRANSFORMERS = {'expression': get_expression, 'percent': get_percent}

datadog_checks_base/datadog_checks/base/utils/db/utils.py

@@ -1,6 +1,7 @@
 # (C) Datadog, Inc. 2019-present
 # All rights reserved
 # Licensed under a 3-clause BSD style license (see LICENSE)
+from datetime import timezone
 from itertools import chain
 
 # AgentCheck methods to transformer name e.g. set_metadata -> metadata
@@ -52,3 +53,11 @@ def transformer(sources, **kwargs):
         transformer = column_transformer
 
     return transformer
+
+
+def normalize_datetime(dt):
+    # Prevent naive datetime objects
+    if dt.tzinfo is None:
+        dt = dt.replace(tzinfo=timezone.utc)
+
+    return dt

datadog_checks_base/tests/test_db.py

@@ -2,6 +2,7 @@
 # All rights reserved
 # Licensed under a 3-clause BSD style license (see LICENSE)
 import logging
+from datetime import datetime, timedelta, timezone, tzinfo
 
 import pytest
 
@@ -30,6 +31,17 @@ def create_query_manager(*args, **kwargs):
     return QueryManager(check, executor, [Query(arg) for arg in args], **kwargs)
 
 
+class EST(tzinfo):
+    def utcoffset(self, dt):
+        return timedelta(hours=-5) + self.dst(dt)
+
+    def dst(self, dt):
+        return timedelta(0)
+
+    def tzname(self, dt):
+        return 'Eastern Standard Time'
+
+
 class TestQueryResultIteration:
     def test_executor_empty_result(self):
         query_manager = create_query_manager()
@@ -913,6 +925,25 @@ def test_service_check_status_map_status_invalid(self):
         ):
             query_manager.compile_queries()
 
+    def test_time_elapsed_format_not_string(self):
+        query_manager = create_query_manager(
+            {
+                'name': 'test query',
+                'query': 'foo',
+                'columns': [{'name': 'test.foo', 'type': 'time_elapsed', 'format': 5}],
+                'tags': ['test:bar'],
+            }
+        )
+
+        with pytest.raises(
+            ValueError,
+            match=(
+                '^error compiling type `time_elapsed` for column test.foo of test query: '
+                'the `format` parameter must be a string$'
+            ),
+        ):
+            query_manager.compile_queries()
+
 
 class TestSubmission:
     @pytest.mark.parametrize(
@@ -1386,6 +1417,110 @@ def test_service_check_unknown(self, aggregator):
         aggregator.assert_service_check('test.foo', 3, message='baz', tags=['test:foo', 'test:bar'])
         aggregator.assert_all_metrics_covered()
 
+    def test_time_elapsed_native(self, aggregator):
+        query_manager = create_query_manager(
+            {
+                'name': 'test query',
+                'query': 'foo',
+                'columns': [
+                    {'name': 'test', 'type': 'tag'},
+                    {'name': 'test.foo', 'type': 'time_elapsed', 'format': 'native'},
+                ],
+                'tags': ['test:bar'],
+            },
+            executor=mock_executor([['tag1', datetime.now(timezone.utc) + timedelta(hours=-1)]]),
+            tags=['test:foo'],
+        )
+        query_manager.compile_queries()
+        query_manager.execute()
+
+        assert 'test.foo' in aggregator._metrics
+        assert len(aggregator._metrics) == 1
+        assert len(aggregator._metrics['test.foo']) == 1
+        m = aggregator._metrics['test.foo'][0]
+
+        assert 3599 < m.value < 3601
+        assert m.type == aggregator.GAUGE
+        assert m.tags == ['test:foo', 'test:bar', 'test:tag1']
+
+    def test_time_elapsed_native_default(self, aggregator):
+        query_manager = create_query_manager(
+            {
+                'name': 'test query',
+                'query': 'foo',
+                'columns': [{'name': 'test', 'type': 'tag'}, {'name': 'test.foo', 'type': 'time_elapsed'}],
+                'tags': ['test:bar'],
+            },
+            executor=mock_executor([['tag1', datetime.now(timezone.utc) + timedelta(hours=-1)]]),
+            tags=['test:foo'],
+        )
+        query_manager.compile_queries()
+        query_manager.execute()
+
+        assert 'test.foo' in aggregator._metrics
+        assert len(aggregator._metrics) == 1
+        assert len(aggregator._metrics['test.foo']) == 1
+        m = aggregator._metrics['test.foo'][0]
+
+        assert 3599 < m.value < 3601
+        assert m.type == aggregator.GAUGE
+        assert m.tags == ['test:foo', 'test:bar', 'test:tag1']
+
+    def test_time_elapsed_format(self, aggregator):
+        time_format = '%Y-%m-%dT%H-%M-%S%z'
+        query_manager = create_query_manager(
+            {
+                'name': 'test query',
+                'query': 'foo',
+                'columns': [
+                    {'name': 'test', 'type': 'tag'},
+                    {'name': 'test.foo', 'type': 'time_elapsed', 'format': time_format},
+                ],
+                'tags': ['test:bar'],
+            },
+            executor=mock_executor(
+                [['tag1', (datetime.now(timezone.utc) + timedelta(hours=-1)).strftime(time_format)]]
+            ),
+            tags=['test:foo'],
+        )
+        query_manager.compile_queries()
+        query_manager.execute()
+
+        assert 'test.foo' in aggregator._metrics
+        assert len(aggregator._metrics) == 1
+        assert len(aggregator._metrics['test.foo']) == 1
+        m = aggregator._metrics['test.foo'][0]
+
+        assert 3599 < m.value < 3601
+        assert m.type == aggregator.GAUGE
+        assert m.tags == ['test:foo', 'test:bar', 'test:tag1']
+
+    def test_time_elapsed_datetime_naive(self, aggregator):
+        query_manager = create_query_manager(
+            {
+                'name': 'test query',
+                'query': 'foo',
+                'columns': [
+                    {'name': 'test', 'type': 'tag'},
+                    {'name': 'test.foo', 'type': 'time_elapsed', 'format': 'native'},
+                ],
+                'tags': ['test:bar'],
+            },
+            executor=mock_executor([['tag1', datetime.now() + timedelta(hours=-1)]]),
+            tags=['test:foo'],
+        )
+        query_manager.compile_queries()
+        query_manager.execute()
+
+        assert 'test.foo' in aggregator._metrics
+        assert len(aggregator._metrics) == 1
+        assert len(aggregator._metrics['test.foo']) == 1
+        m = aggregator._metrics['test.foo'][0]
+
+        assert 3599 < m.value < 3601
+        assert m.type == aggregator.GAUGE
+        assert m.tags == ['test:foo', 'test:bar', 'test:tag1']
+
 
 class TestExtraTransformers:
     def test_expression(self, aggregator):