Add global custom queries feature for postgres (#17993)

* Add global custom queries for postgres * Add changelog entry * Update config models * Update instance config model
DataDog · Jul 22, 2024 · 20858d0 · 20858d0
1 parent 4d47709
commit 20858d0
Show file tree

Hide file tree

Showing 8 changed files with 188 additions and 1 deletion.
diff --git a/postgres/assets/configuration/spec.yaml b/postgres/assets/configuration/spec.yaml
@@ -11,6 +11,27 @@ files:
       value:
         example: false
         type: boolean
+    - name: global_custom_queries
+      description: |
+        See `custom_queries` defined below.
+
+        Global custom queries can be applied to all instances using the
+        `use_global_custom_queries` setting at the instance level.
+      value:
+        type: array
+        items:
+          type: object
+          properties: []
+        example:
+            - metric_prefix: postgresql
+              query: <QUERY>
+              columns:
+                - name: <COLUNMS_1_NAME>
+                  type: <COLUMNS_1_TYPE>
+                - name: <COLUNMS_2_NAME>
+                  type: <COLUMNS_2_TYPE>
+              tags:
+                - <TAG_KEY>:<TAG_VALUE>
     - template: init_config/default
   - template: instances
     options:
@@ -326,6 +347,16 @@ files:
                   type: <COLUMNS_2_TYPE>
               tags:
                 - <TAG_KEY>:<TAG_VALUE>
+    - name: use_global_custom_queries
+      description: |
+        How `global_custom_queries` should be used for this instance. There are 3 options:
+
+        1. true - `global_custom_queries` override `custom_queries`.
+        2. false - `custom_queries` override `global_custom_queries`.
+        3. extend - `global_custom_queries` are used in addition to any `custom_queries`.
+      value:
+        type: string
+        example: extend
     - name: database_autodiscovery
       description: |
         Define the configuration for database autodiscovery.

diff --git a/postgres/changelog.d/17993.added b/postgres/changelog.d/17993.added
@@ -0,0 +1 @@
+Add global custom queries for postgres
diff --git a/postgres/datadog_checks/postgres/config.py b/postgres/datadog_checks/postgres/config.py
@@ -99,6 +99,11 @@ def __init__(self, instance, init_config, check):
         if is_affirmative(instance.get('collect_default_database', True)):
             self.ignore_databases = [d for d in self.ignore_databases if d != 'postgres']
         self.custom_queries = instance.get('custom_queries', [])
+        self.use_global_custom_queries = instance.get('use_global_custom_queries', 'extend')
+        if self.use_global_custom_queries == 'extend':
+            self.custom_queries.extend(init_config.get('global_custom_queries', []))
+        elif is_affirmative(self.use_global_custom_queries):
+            self.custom_queries = init_config.get('global_custom_queries', [])
         self.tag_replication_role = is_affirmative(instance.get('tag_replication_role', True))
         self.custom_metrics = self._get_custom_metrics(instance.get('custom_metrics', []))
         self.max_relations = int(instance.get('max_relations', 300))

diff --git a/postgres/datadog_checks/postgres/config_models/defaults.py b/postgres/datadog_checks/postgres/config_models/defaults.py
@@ -138,3 +138,7 @@ def instance_table_count_limit():
 
 def instance_tag_replication_role():
     return False
+
+
+def instance_use_global_custom_queries():
+    return 'extend'
diff --git a/postgres/datadog_checks/postgres/config_models/instance.py b/postgres/datadog_checks/postgres/config_models/instance.py
@@ -249,6 +249,7 @@ class InstanceConfig(BaseModel):
     table_count_limit: Optional[int] = None
     tag_replication_role: Optional[bool] = None
     tags: Optional[tuple[str, ...]] = None
+    use_global_custom_queries: Optional[str] = None
     username: str
 
     @model_validator(mode='before')

diff --git a/postgres/datadog_checks/postgres/config_models/shared.py b/postgres/datadog_checks/postgres/config_models/shared.py
@@ -9,7 +9,8 @@
 
 from __future__ import annotations
 
-from typing import Optional
+from types import MappingProxyType
+from typing import Any, Optional
 
 from pydantic import BaseModel, ConfigDict, field_validator, model_validator
 
@@ -25,6 +26,7 @@ class SharedConfig(BaseModel):
         arbitrary_types_allowed=True,
         frozen=True,
     )
+    global_custom_queries: Optional[tuple[MappingProxyType[str, Any], ...]] = None
     propagate_agent_tags: Optional[bool] = None
     service: Optional[str] = None
 

diff --git a/postgres/datadog_checks/postgres/data/conf.yaml.example b/postgres/datadog_checks/postgres/data/conf.yaml.example
@@ -8,6 +8,23 @@ init_config:
     #
     # propagate_agent_tags: false
 
+    ## @param global_custom_queries - list of mappings - optional
+    ## See `custom_queries` defined below.
+    ##
+    ## Global custom queries can be applied to all instances using the
+    ## `use_global_custom_queries` setting at the instance level.
+    #
+    # global_custom_queries:
+    #   - metric_prefix: postgresql
+    #     query: <QUERY>
+    #     columns:
+    #     - name: <COLUNMS_1_NAME>
+    #       type: <COLUMNS_1_TYPE>
+    #     - name: <COLUNMS_2_NAME>
+    #       type: <COLUMNS_2_TYPE>
+    #     tags:
+    #     - <TAG_KEY>:<TAG_VALUE>
+
     ## @param service - string - optional
     ## Attach the tag `service:<SERVICE>` to every metric, event, and service check emitted by this integration.
     ##
@@ -258,6 +275,15 @@ instances:
     #     tags:
     #     - <TAG_KEY>:<TAG_VALUE>
 
+    ## @param use_global_custom_queries - string - optional - default: extend
+    ## How `global_custom_queries` should be used for this instance. There are 3 options:
+    ##
+    ## 1. true - `global_custom_queries` override `custom_queries`.
+    ## 2. false - `custom_queries` override `global_custom_queries`.
+    ## 3. extend - `global_custom_queries` are used in addition to any `custom_queries`.
+    #
+    # use_global_custom_queries: extend
+
     ## Define the configuration for database autodiscovery.
     ## Complete this section if you want to auto-discover databases on this host
     ## instead of specifying each using dbname.

diff --git a/postgres/tests/test_custom_metrics.py b/postgres/tests/test_custom_metrics.py
@@ -65,3 +65,120 @@ def test_custom_queries(aggregator, pg_instance):
 
         aggregator.assert_metric('custom.num', value=value, tags=custom_tags + ['query:custom'])
         aggregator.assert_metric('another_custom_one.num', value=value, tags=custom_tags + ['query:another_custom_one'])
+
+
+@pytest.mark.integration
+@pytest.mark.usefixtures('dd_environment')
+def test_both_global_and_instance_custom_queries(aggregator, pg_instance):
+    pg_instance.update(
+        {
+            'custom_queries': [
+                {
+                    'metric_prefix': 'custom',
+                    'query': "SELECT letter, num FROM (VALUES (97, 'a'), (98, 'b'), (99, 'c')) AS t (num,letter)",
+                    'columns': [{'name': 'customtag', 'type': 'tag'}, {'name': 'num', 'type': 'gauge'}],
+                    'tags': ['query:custom'],
+                },
+            ],
+            'use_global_custom_queries': 'extend',
+        }
+    )
+    pg_init_config = {
+        'global_custom_queries': [
+            {
+                'metric_prefix': 'global_custom',
+                'query': "SELECT letter, num FROM (VALUES (97, 'a'), (98, 'b'), (99, 'c')) AS t (num,letter)",
+                'columns': [{'name': 'customtag', 'type': 'tag'}, {'name': 'num', 'type': 'gauge'}],
+                'tags': ['query:global_custom'],
+            },
+        ]
+    }
+    postgres_check = PostgreSql('postgres', pg_init_config, [pg_instance])
+    postgres_check.check(pg_instance)
+    tags = _get_expected_tags(postgres_check, pg_instance, with_db=True)
+
+    for tag in ('a', 'b', 'c'):
+        value = ord(tag)
+        custom_tags = [f'customtag:{tag}']
+        custom_tags.extend(tags)
+
+        aggregator.assert_metric('custom.num', value=value, tags=custom_tags + ['query:custom'])
+        aggregator.assert_metric('global_custom.num', value=value, tags=custom_tags + ['query:global_custom'])
+
+
+@pytest.mark.integration
+@pytest.mark.usefixtures('dd_environment')
+def test_only_global_custom_queries(aggregator, pg_instance):
+    pg_instance.update(
+        {
+            'custom_queries': [
+                {
+                    'metric_prefix': 'custom',
+                    'query': "SELECT letter, num FROM (VALUES (97, 'a'), (98, 'b'), (99, 'c')) AS t (num,letter)",
+                    'columns': [{'name': 'customtag', 'type': 'tag'}, {'name': 'num', 'type': 'gauge'}],
+                    'tags': ['query:custom'],
+                },
+            ],
+            'use_global_custom_queries': 'true',
+        }
+    )
+    pg_init_config = {
+        'global_custom_queries': [
+            {
+                'metric_prefix': 'global_custom',
+                'query': "SELECT letter, num FROM (VALUES (97, 'a'), (98, 'b'), (99, 'c')) AS t (num,letter)",
+                'columns': [{'name': 'customtag', 'type': 'tag'}, {'name': 'num', 'type': 'gauge'}],
+                'tags': ['query:global_custom'],
+            },
+        ]
+    }
+    postgres_check = PostgreSql('postgres', pg_init_config, [pg_instance])
+    postgres_check.check(pg_instance)
+    tags = _get_expected_tags(postgres_check, pg_instance, with_db=True)
+
+    for tag in ('a', 'b', 'c'):
+        value = ord(tag)
+        custom_tags = [f'customtag:{tag}']
+        custom_tags.extend(tags)
+
+        aggregator.assert_metric('custom.num', value=value, tags=custom_tags + ['query:custom'], count=0)
+        aggregator.assert_metric('global_custom.num', value=value, tags=custom_tags + ['query:global_custom'])
+
+
+@pytest.mark.integration
+@pytest.mark.usefixtures('dd_environment')
+def test_only_instance_custom_queries(aggregator, pg_instance):
+    pg_instance.update(
+        {
+            'custom_queries': [
+                {
+                    'metric_prefix': 'custom',
+                    'query': "SELECT letter, num FROM (VALUES (97, 'a'), (98, 'b'), (99, 'c')) AS t (num,letter)",
+                    'columns': [{'name': 'customtag', 'type': 'tag'}, {'name': 'num', 'type': 'gauge'}],
+                    'tags': ['query:custom'],
+                },
+            ],
+            'use_global_custom_queries': 'false',
+        }
+    )
+    pg_init_config = {
+        'global_custom_queries': [
+            {
+                'metric_prefix': 'global_custom',
+                'query': "SELECT letter, num FROM (VALUES (97, 'a'), (98, 'b'), (99, 'c')) AS t (num,letter)",
+                'columns': [{'name': 'customtag', 'type': 'tag'}, {'name': 'num', 'type': 'gauge'}],
+                'tags': ['query:global_custom'],
+            },
+        ]
+    }
+    postgres_check = PostgreSql('postgres', pg_init_config, [pg_instance])
+    postgres_check.check(pg_instance)
+    tags = _get_expected_tags(postgres_check, pg_instance, with_db=True)
+
+    for tag in ('a', 'b', 'c'):
+        value = ord(tag)
+        custom_tags = [f'customtag:{tag}']
+        custom_tags.extend(tags)
+
+        aggregator.assert_metric('custom.num', value=value, tags=custom_tags + ['query:custom'])
+        aggregator.assert_metric('global_custom.num', value=value, tags=custom_tags + ['query:global_custom'], count=0)