Support metadata and service checks for DB utility

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

@@ -6,6 +6,7 @@
 import re
 
 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
@@ -83,6 +84,18 @@ def match(sources, value, **kwargs):
     return match
 
 
+def get_service_check(transformers, column_name, **modifiers):
+    # Do work in a separate function to avoid having to `del` a bunch of variables
+    status_map = _compile_service_check_statuses(modifiers)
+
+    service_check_method = transformers['__service_check'](transformers, column_name, **modifiers)
+
+    def service_check(_, value, **kwargs):
+        service_check_method(_, status_map.get(value, ServiceCheck.UNKNOWN), **kwargs)
+
+    return service_check
+
+
 def get_expression(transformers, name, **modifiers):
     available_sources = modifiers.pop('sources')
 
@@ -169,11 +182,38 @@ def percent(sources, **kwargs):
     'monotonic_gauge': get_monotonic_gauge,
     'tag': get_tag,
     'match': get_match,
+    'service_check': get_service_check,
 }
 
 EXTRA_TRANSFORMERS = {'expression': get_expression, 'percent': get_percent}
 
 
+def _compile_service_check_statuses(modifiers):
+    status_map = modifiers.pop('status_map', None)
+    if status_map is None:
+        raise ValueError('the `status_map` parameter is required')
+    elif not isinstance(status_map, dict):
+        raise ValueError('the `status_map` parameter must be a mapping')
+    elif not status_map:
+        raise ValueError('the `status_map` parameter must not be empty')
+
+    for value, status_string in list(status_map.items()):
+        if not isinstance(status_string, str):
+            raise ValueError(
+                'status `{}` for value `{}` of parameter `status_map` is not a string'.format(status_string, value)
+            )
+
+        status = getattr(ServiceCheck, status_string.upper(), None)
+        if status is None:
+            raise ValueError(
+                'invalid status `{}` for value `{}` of parameter `status_map`'.format(status_string, value)
+            )
+
+        status_map[value] = status
+
+    return status_map
+
+
 def _compile_match_items(transformers, modifiers):
     items = modifiers.pop('items', None)
     if items is None:

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

@@ -11,6 +11,10 @@
     'rate': 'rate',
     'histogram': 'histogram',
     'historate': 'historate',
+    'set_metadata': 'metadata',
+    # These submission methods require more configuration than just a name
+    # and a value and therefore must be defined as a custom transformer.
+    'service_check': '__service_check',
 }
 
 

datadog_checks_base/tests/test_db.py

@@ -25,6 +25,8 @@ def create_query_manager(*args, **kwargs):
         executor = mock_executor()
 
     check = kwargs.pop('check', None) or AgentCheck('test', {}, [{}])
+    check.check_id = 'test:instance'
+
     return QueryManager(check, executor, [Query(arg) for arg in args], **kwargs)
 
 
@@ -806,6 +808,111 @@ def test_percent_total_not_source(self):
         ):
             query_manager.compile_queries()
 
+    def test_service_check_no_status_map(self):
+        query_manager = create_query_manager(
+            {
+                'name': 'test query',
+                'query': 'foo',
+                'columns': [{'name': 'test.foo', 'type': 'service_check'}],
+                'tags': ['test:bar'],
+            },
+            executor=mock_executor([['known']]),
+            tags=['test:foo'],
+        )
+
+        with pytest.raises(
+            ValueError,
+            match=(
+                '^error compiling type `service_check` for column test.foo of test query: '
+                'the `status_map` parameter is required$'
+            ),
+        ):
+            query_manager.compile_queries()
+
+    def test_service_check_status_map_not_dict(self):
+        query_manager = create_query_manager(
+            {
+                'name': 'test query',
+                'query': 'foo',
+                'columns': [{'name': 'test.foo', 'type': 'service_check', 'status_map': 5}],
+                'tags': ['test:bar'],
+            },
+            executor=mock_executor([['known']]),
+            tags=['test:foo'],
+        )
+
+        with pytest.raises(
+            ValueError,
+            match=(
+                '^error compiling type `service_check` for column test.foo of test query: '
+                'the `status_map` parameter must be a mapping$'
+            ),
+        ):
+            query_manager.compile_queries()
+
+    def test_service_check_status_map_empty(self):
+        query_manager = create_query_manager(
+            {
+                'name': 'test query',
+                'query': 'foo',
+                'columns': [{'name': 'test.foo', 'type': 'service_check', 'status_map': {}}],
+                'tags': ['test:bar'],
+            },
+            executor=mock_executor([['known']]),
+            tags=['test:foo'],
+        )
+
+        with pytest.raises(
+            ValueError,
+            match=(
+                '^error compiling type `service_check` for column test.foo of test query: '
+                'the `status_map` parameter must not be empty$'
+            ),
+        ):
+            query_manager.compile_queries()
+
+    def test_service_check_status_map_status_not_string(self):
+        query_manager = create_query_manager(
+            {
+                'name': 'test query',
+                'query': 'foo',
+                'columns': [{'name': 'test.foo', 'type': 'service_check', 'status_map': {'known': 0}}],
+                'tags': ['test:bar'],
+            },
+            executor=mock_executor([['known']]),
+            tags=['test:foo'],
+        )
+
+        with pytest.raises(
+            ValueError,
+            match=(
+                '^error compiling type `service_check` for column test.foo of test query: '
+                'status `0` for value `known` of parameter `status_map` is not a string$'
+            ),
+        ):
+            query_manager.compile_queries()
+
+    def test_service_check_status_map_status_invalid(self):
+        query_manager = create_query_manager(
+            {
+                'name': 'test query',
+                'query': 'foo',
+                'columns': [{'name': 'test.foo', 'type': 'service_check', 'status_map': {'known': '0k'}}],
+                'tags': ['test:bar'],
+            },
+            executor=mock_executor([['known']]),
+            tags=['test:foo'],
+        )
+
+        with pytest.raises(
+            ValueError,
+            match=(
+                '^error compiling type `service_check` for column test.foo of test query: '
+                'invalid status `0k` for value `known` of parameter `status_map`$'
+            ),
+        ):
+            query_manager.compile_queries()
+
 
 class TestSubmission:
     @pytest.mark.parametrize(
@@ -1028,6 +1135,34 @@ def test_extra_transformer_error(self, caplog, aggregator):
 
         aggregator.assert_all_metrics_covered()
 
+    def test_metadata(self, aggregator, datadog_agent):
+        query_manager = create_query_manager(
+            {
+                'name': 'test query',
+                'query': 'foo',
+                'columns': [{'name': 'version', 'type': 'metadata'}],
+                'tags': ['test:bar'],
+            },
+            executor=mock_executor([['1.2.3-rc.4+5']]),
+            tags=['test:foo'],
+        )
+        query_manager.compile_queries()
+        query_manager.execute()
+
+        version_metadata = {
+            'version.major': '1',
+            'version.minor': '2',
+            'version.patch': '3',
+            'version.release': 'rc.4',
+            'version.build': '5',
+            'version.raw': '1.2.3-rc.4+5',
+            'version.scheme': 'semver',
+        }
+
+        datadog_agent.assert_metadata('test:instance', version_metadata)
+        datadog_agent.assert_metadata_count(len(version_metadata))
+        aggregator.assert_all_metrics_covered()
+
 
 class TestColumnTransformers:
     def test_tag_boolean(self, aggregator):
@@ -1213,6 +1348,44 @@ def test_match_none(self, aggregator):
 
         aggregator.assert_all_metrics_covered()
 
+    def test_service_check_known(self, aggregator):
+        query_manager = create_query_manager(
+            {
+                'name': 'test query',
+                'query': 'foo',
+                'columns': [
+                    {'name': 'test.foo', 'type': 'service_check', 'status_map': {'known': 'ok'}, 'message': 'baz'},
+                ],
+                'tags': ['test:bar'],
+            },
+            executor=mock_executor([['known']]),
+            tags=['test:foo'],
+        )
+        query_manager.compile_queries()
+        query_manager.execute()
+
+        aggregator.assert_service_check('test.foo', 0, message='baz', tags=['test:foo', 'test:bar'])
+        aggregator.assert_all_metrics_covered()
+
+    def test_service_check_unknown(self, aggregator):
+        query_manager = create_query_manager(
+            {
+                'name': 'test query',
+                'query': 'foo',
+                'columns': [
+                    {'name': 'test.foo', 'type': 'service_check', 'status_map': {'known': 'ok'}, 'message': 'baz'},
+                ],
+                'tags': ['test:bar'],
+            },
+            executor=mock_executor([['unknown']]),
+            tags=['test:foo'],
+        )
+        query_manager.compile_queries()
+        query_manager.execute()
+
+        aggregator.assert_service_check('test.foo', 3, message='baz', tags=['test:foo', 'test:bar'])
+        aggregator.assert_all_metrics_covered()
+
 
 class TestExtraTransformers:
     def test_expression(self, aggregator):