[dbm] Fix logic for only_custom_queries parameter (#21708)

* fix logic for only_custom_queries parameter

* fix custom queries test

* add changelog

* run ddev test format

Author: dalextorres

postgres/changelog.d/21708.fixed

@@ -0,0 +1 @@
+Fix logic for only_custom_queries configuration option.

postgres/datadog_checks/postgres/postgres.py

@@ -1099,17 +1099,21 @@ def check(self, _):
 
             self.log.debug("Running check against version %s: is_aurora: %s", str(self.version), str(self.is_aurora))
             self._emit_running_metric()
-            self._collect_stats(tags)
+
+            if not self._config.only_custom_queries:
+                self._collect_stats(tags)
+                if self._config.dbm:
+                    self.statement_metrics.run_job_loop(tags)
+                    self.statement_samples.run_job_loop(tags)
+                    self.metadata_samples.run_job_loop(tags)
+                if self._config.collect_wal_metrics:
+                    # collect wal metrics for pg < 10, disabled by enabled
+                    self._collect_wal_metrics()
+
             if self._query_manager.queries:
                 self._query_manager.executor = functools.partial(self.execute_query_raw, db=self.db)
                 self._query_manager.execute(extra_tags=tags)
-            if self._config.dbm:
-                self.statement_metrics.run_job_loop(tags)
-                self.statement_samples.run_job_loop(tags)
-                self.metadata_samples.run_job_loop(tags)
-            if self._config.collect_wal_metrics:
-                # collect wal metrics for pg < 10, disabled by enabled
-                self._collect_wal_metrics()
+
         except Exception as e:
             self.log.exception("Unable to collect postgres metrics.")
             self._clean_state()

postgres/tests/test_custom_metrics.py

@@ -141,3 +141,35 @@ def test_only_instance_custom_queries(aggregator, pg_instance, dd_run_check, int
 
         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)
+
+
+@pytest.mark.integration
+@pytest.mark.usefixtures('dd_environment')
+def test_only_custom_queries(aggregator, pg_instance, dd_run_check, integration_check):
+    pg_instance.update(
+        {
+            'only_custom_queries': True,
+            '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'],
+                },
+            ],
+        }
+    )
+    postgres_check = integration_check(pg_instance)
+    dd_run_check(postgres_check)
+    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'])
+
+    running_tags = _get_expected_tags(postgres_check, pg_instance)
+    aggregator.assert_metric('postgresql.running', count=1, value=1, tags=running_tags)
+    aggregator.assert_all_metrics_covered()