Enable only_custom_queries for sqlserver

edengorevoy · edengorevoy · commit c36a70c5c319 · 2025-09-10T15:57:54.000-04:00
diff --git a/sqlserver/datadog_checks/sqlserver/config.py b/sqlserver/datadog_checks/sqlserver/config.py
@@ -40,6 +40,7 @@ def __init__(self, init_config, instance, log):
 
         self.proc: str = instance.get('stored_procedure')
         self.custom_metrics: list[dict] = init_config.get('custom_metrics', []) or []
+        self.only_custom_queries: bool = is_affirmative(instance.get('only_custom_queries', False))
         self.ignore_missing_database = is_affirmative(instance.get("ignore_missing_database", False))
         if self.ignore_missing_database:
             self.log.warning(
diff --git a/sqlserver/datadog_checks/sqlserver/sqlserver.py b/sqlserver/datadog_checks/sqlserver/sqlserver.py
@@ -821,6 +821,7 @@ def _check_database_conns(self):
             self._check_connections_by_use_db()
 
     def check(self, _):
+        self.log.error("[EG] blooob")
         if self.do_check:
             self.load_static_information()
             # configure custom queries for the check
@@ -830,11 +831,15 @@ def check(self, _):
                     self, self.execute_query_raw, tags=self.tag_manager.get_tags(), hostname=self.reported_hostname
                 )
                 self._query_manager.compile_queries()
+
             self._send_database_instance_metadata()
+
             if self._config.proc:
                 self.do_stored_procedure_check()
             else:
                 self.collect_metrics()
+
+            self.log.error("[EG] I am right here!!!! blub")
             if self._config.autodiscovery and self._config.autodiscovery_db_service_check:
                 self._check_database_conns()
             if self._config.dbm_enabled:
@@ -852,6 +857,7 @@ def check(self, _):
                         handler.run_job_loop(self.tag_manager.get_tags())
                     except Exception as e:
                         self.log.error("Error running XE session handler for %s: %s", handler.session_name, e)
+
         else:
             self.log.debug("Skipping check")
 
@@ -926,52 +932,56 @@ def log_missing_metric(self, metric_name, major_version, engine_version):
         else:
             self.log.warning("%s metrics are not supported on Azure engine version: %s", metric_name, engine_version)
 
-    def collect_metrics(self):
-        """Fetch the metrics from all the associated database tables."""
+    # queries for default integration metrics from the database
+    def load_basic_metrics(self, cursor):
+        # initiate autodiscovery or if the server was down at check __init__ key could be missing.
+        if self.autodiscover_databases(cursor) or not self.instance_metrics:
+            self._make_metric_list_to_collect(self._config.custom_metrics)
+
+        instance_results = {}
+        engine_edition = self.static_info_cache.get(STATIC_INFO_ENGINE_EDITION, "")
+        # Execute the `fetch_all` operations first to minimize the database calls
+        for cls, metric_names in self.instance_per_type_metrics.items():
+            if not metric_names:
+                instance_results[cls] = None, None
+            else:
+                try:
+                    db_names = [d.name for d in self.databases] or [
+                        self.instance.get("database", self.connection.DEFAULT_DATABASE)
+                    ]
+                    metric_cls = getattr(metrics, cls)
+                    with tracked_query(self, operation=metric_cls.OPERATION_NAME):
+                        rows, cols = metric_cls.fetch_all_values(
+                            cursor,
+                            list(metric_names),
+                            self.log,
+                            databases=db_names,
+                            engine_edition=engine_edition,
+                        )
+                except Exception as e:
+                    self.log.error("Error running `fetch_all` for metrics %s - skipping.  Error: %s", cls, e)
+                    rows, cols = None, None
 
-        with self.connection.open_managed_default_connection():
-            with self.connection.get_managed_cursor() as cursor:
-                # initiate autodiscovery or if the server was down at check __init__ key could be missing.
-                if self.autodiscover_databases(cursor) or not self.instance_metrics:
-                    self._make_metric_list_to_collect(self._config.custom_metrics)
+                instance_results[cls] = rows, cols
 
-                instance_results = {}
-                engine_edition = self.static_info_cache.get(STATIC_INFO_ENGINE_EDITION, "")
-                # Execute the `fetch_all` operations first to minimize the database calls
-                for cls, metric_names in self.instance_per_type_metrics.items():
-                    if not metric_names:
-                        instance_results[cls] = None, None
+        for metric in self.instance_metrics:
+            key = metric.__class__.__name__
+            if key not in instance_results:
+                self.log.warning("No %s metrics found, skipping", str(key))
+            else:
+                rows, cols = instance_results[key]
+                if rows is not None:
+                    if key == "SqlIncrFractionMetric":
+                        metric.fetch_metric(rows, cols, self.sqlserver_incr_fraction_metric_previous_values)
                     else:
-                        try:
-                            db_names = [d.name for d in self.databases] or [
-                                self.instance.get("database", self.connection.DEFAULT_DATABASE)
-                            ]
-                            metric_cls = getattr(metrics, cls)
-                            with tracked_query(self, operation=metric_cls.OPERATION_NAME):
-                                rows, cols = metric_cls.fetch_all_values(
-                                    cursor,
-                                    list(metric_names),
-                                    self.log,
-                                    databases=db_names,
-                                    engine_edition=engine_edition,
-                                )
-                        except Exception as e:
-                            self.log.error("Error running `fetch_all` for metrics %s - skipping.  Error: %s", cls, e)
-                            rows, cols = None, None
-
-                        instance_results[cls] = rows, cols
+                        metric.fetch_metric(rows, cols)
 
-                for metric in self.instance_metrics:
-                    key = metric.__class__.__name__
-                    if key not in instance_results:
-                        self.log.warning("No %s metrics found, skipping", str(key))
-                    else:
-                        rows, cols = instance_results[key]
-                        if rows is not None:
-                            if key == "SqlIncrFractionMetric":
-                                metric.fetch_metric(rows, cols, self.sqlserver_incr_fraction_metric_previous_values)
-                            else:
-                                metric.fetch_metric(rows, cols)
+    def collect_metrics(self):
+        """Fetch the metrics from all the associated database tables."""
+        with self.connection.open_managed_default_connection():
+            if not self._config.only_custom_queries:
+                with self.connection.get_managed_cursor() as cursor:
+                    self.load_basic_metrics(cursor)
 
             # Neither pyodbc nor adodbapi are able to read results of a query if the number of rows affected
             # statement are returned as part of the result set, so we disable for the entire connection
@@ -980,14 +990,15 @@ def collect_metrics(self):
             with self.connection.get_managed_cursor() as cursor:
                 cursor.execute("SET NOCOUNT ON")
             try:
-                # restore the current database after executing dynamic queries
-                # this is to ensure the current database context is not changed
-                with self.connection.restore_current_database_context():
-                    if self.database_metrics:
-                        for database_metric in self.database_metrics:
-                            database_metric.execute()
-
-                # reuse connection for any custom queries
+                if not self._config.only_custom_queries:
+                    # restore the current database after executing dynamic queries
+                    # this is to ensure the current database context is not changed
+                    with self.connection.restore_current_database_context():
+                        if self.database_metrics:
+                            for database_metric in self.database_metrics:
+                                database_metric.execute()
+
+                # reuse the connection for custom queries
                 self._query_manager.execute()
             finally:
                 with self.connection.get_managed_cursor() as cursor:
diff --git a/sqlserver/tests/test_integration.py b/sqlserver/tests/test_integration.py
@@ -506,6 +506,59 @@ def test_custom_queries(aggregator, dd_run_check, instance_docker, custom_query,
         aggregator.assert_metric(metric_name, **kwargs)
 
 
+@pytest.mark.integration
+@pytest.mark.usefixtures('dd_environment')
+@pytest.mark.parametrize(
+    "custom_query, assert_metrics",
+    [
+        (
+            {
+                'query': "SELECT letter, num FROM (VALUES (97, 'a'), (98, 'b')) AS t (num,letter)",
+                'columns': [{'name': 'customtag', 'type': 'tag'}, {'name': 'num', 'type': 'gauge'}],
+                'tags': ['query:custom'],
+            },
+            [
+                ("sqlserver.num", {"value": 97, "tags": ["customtag:a", "query:custom"]}),
+                ("sqlserver.num", {"value": 98, "tags": ["customtag:b", "query:custom"]}),
+            ],
+        ),
+        (
+            {
+                'query': "EXEC exampleProcWithoutNocount",
+                'columns': [{'name': 'value', 'type': 'gauge'}],
+                'tags': ['hello:there'],
+            },
+            [
+                ("sqlserver.value", {"value": 1, "tags": ["hello:there"]}),
+            ],
+        ),
+    ],
+)
+def test_custom_queries_only(aggregator, dd_run_check, instance_docker, custom_query, assert_metrics):
+    instance = copy(instance_docker)
+    instance['custom_queries'] = [custom_query]
+    instance['only_custom_queries'] = True
+    instance['procedure_metrics'] = {'enabled': False}
+
+    check = SQLServer(CHECK_NAME, {}, [instance])
+    dd_run_check(check)
+
+    for metric_name, kwargs in assert_metrics:
+        kwargs = copy(kwargs)
+        kwargs['tags'] = (
+            check._config.tags
+            + [
+                "database_hostname:{}".format("stubbed.hostname"),
+                "database_instance:{}".format("stubbed.hostname"),
+                "ddagenthostname:{}".format("stubbed.hostname"),
+                "dd.internal.resource:database_instance:{}".format("stubbed.hostname"),
+                "sqlserver_servername:{}".format(check.static_info_cache.get(STATIC_INFO_SERVERNAME)),
+            ]
+            + kwargs.get('tags', [])
+        )
+        aggregator.assert_metric(metric_name, **kwargs)
+
+
 @pytest.mark.integration
 @pytest.mark.usefixtures('dd_environment')
 def test_load_static_information(aggregator, dd_run_check, instance_docker):
