Isolate SQL Server plan lookup failures

datadog_checks_base/changelog.d/25091.added

@@ -0,0 +1 @@
+Add RateLimitingTTLCache.would_acquire to check admission without consuming a key's budget.

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

@@ -150,13 +150,23 @@ class RateLimitingTTLCache(TTLCache):
     TTLCache wrapper used for rate limiting by key
     """
 
-    def acquire(self, key):
+    def would_acquire(self, key):
         """
-        :return: True if the key has not yet reached its rate limit
+        :return: True if :meth:`acquire` would admit ``key`` right now, without consuming its budget
+
+        Use this when the work guarded by the rate limit can fail: check with ``would_acquire`` first,
+        do the work, and only ``acquire`` once it succeeded. Calling ``acquire`` up front would mark the
+        key as seen even when the work failed, suppressing retries until the entry expires.
         """
         if len(self) >= self.maxsize:
             return False
-        if key in self:
+        return key not in self
+
+    def acquire(self, key):
+        """
+        :return: True if the key has not yet reached its rate limit
+        """
+        if not self.would_acquire(key):
             return False
         self[key] = True
         return True

datadog_checks_base/tests/base/utils/db/test_util.py

@@ -983,3 +983,27 @@ def test_case_sensitivity_without_normalization(self):
 
         expected_tags = sorted(['test_key:UPPERCASE-VALUE', 'env:PRODUCTION', 'KEYLESS-TAG-UPPERCASE'])
         assert sorted(tag_manager.get_tags()) == expected_tags
+
+
+def test_rate_limiting_ttl_cache_would_acquire_does_not_consume_budget():
+    """
+    A failed attempt must not burn the key's budget.
+
+    Without `would_acquire`, callers had to `acquire` before doing the work the limit guards; if that
+    work then failed, the key was already marked seen and the retry was suppressed until the entry
+    expired -- so a transient error silently cost a full TTL of collection.
+    """
+    cache = RateLimitingTTLCache(maxsize=2, ttl=100)
+
+    assert cache.would_acquire('a') is True
+    assert cache.would_acquire('a') is True, "checking admission must be repeatable"
+    assert len(cache) == 0, "checking admission must not store the key"
+
+    assert cache.acquire('a') is True
+    assert cache.would_acquire('a') is False
+    assert cache.acquire('a') is False
+
+    # maxsize is honored by both, so a full cache admits nothing.
+    assert cache.acquire('b') is True
+    assert cache.would_acquire('c') is False
+    assert cache.acquire('c') is False

sqlserver/changelog.d/25091.fixed

@@ -0,0 +1 @@
+Isolate execution plan lookup failures so query metrics collection can continue.

sqlserver/datadog_checks/sqlserver/connection_errors.py

@@ -24,6 +24,20 @@ class SQLConnectionError(Exception):
     pass
 
 
+def _expected_db_exceptions() -> tuple[type[Exception], ...]:
+    exceptions: list[type[Exception]] = [SQLConnectionError]
+    if pyodbc is not None:
+        exceptions.append(pyodbc.Error)
+    if adodbapi is not None:
+        exceptions.append(adodbapi.DatabaseError)
+    return tuple(exceptions)
+
+
+# Database errors a DBM async job should report as a warning rather than a crash. Which entries are
+# present depends on which drivers are installed, so this is built once at import time.
+EXPECTED_DB_EXCEPTIONS = _expected_db_exceptions()
+
+
 class ConnectionErrorCode(Enum):
     """
     Denotes the various reasons a connection might fail.

sqlserver/datadog_checks/sqlserver/data_observability.py

@@ -14,7 +14,7 @@
 from datadog_checks.base.utils.serialization import json
 
 from .connection import split_sqlserver_host_port
-from .connection_errors import SQLConnectionError
+from .connection_errors import EXPECTED_DB_EXCEPTIONS
 from .utils import raise_if_cancelled
 
 try:
@@ -77,13 +77,6 @@ def query(self) -> Query:
         return self.scheduled_query.query
 
 
-_EXPECTED_DB_EXCEPTIONS: list[type[Exception]] = [SQLConnectionError]
-if pyodbc is not None:
-    _EXPECTED_DB_EXCEPTIONS.append(pyodbc.Error)
-if adodbapi is not None:
-    _EXPECTED_DB_EXCEPTIONS.append(adodbapi.DatabaseError)
-
-
 class SqlServerDataObservability(DBMAsyncJob):
     def __init__(self, check: SQLServer, config: InstanceConfig):
         self._check = check
@@ -98,7 +91,7 @@ def __init__(self, check: SQLServer, config: InstanceConfig):
             enabled=config.data_observability.enabled,
             dbms=check.dbms,
             min_collection_interval=config.min_collection_interval,
-            expected_db_exceptions=tuple(_EXPECTED_DB_EXCEPTIONS),
+            expected_db_exceptions=EXPECTED_DB_EXCEPTIONS,
             job_name="data-observability",
         )
         # Filter bad queries on check construction.
@@ -214,7 +207,7 @@ def _execute_single_query(self, cursor: Any, query_spec: Query) -> dict[str, Any
                 'error': None,
             }
         except Exception as e:
-            if not _EXPECTED_DB_EXCEPTIONS or not isinstance(e, tuple(_EXPECTED_DB_EXCEPTIONS)):
+            if not EXPECTED_DB_EXCEPTIONS or not isinstance(e, EXPECTED_DB_EXCEPTIONS):
                 raise
             duration = time.time() - start
             self._log.warning(
@@ -380,15 +373,15 @@ def run_job(self):
                     for due in group:
                         try:
                             self._run_due_query(due, base_tags, conn_dbname)
-                        except tuple(_EXPECTED_DB_EXCEPTIONS) as e:
+                        except EXPECTED_DB_EXCEPTIONS as e:
                             self._log.warning(
                                 "Failed to execute monitor_id=%d on db_name=%s; will retry next poll: %s",
                                 due.query.monitor_id,
                                 conn_dbname,
                                 e,
                             )
                             self._queue_for_retry(due, base_tags)
-            except tuple(_EXPECTED_DB_EXCEPTIONS) as e:
+            except EXPECTED_DB_EXCEPTIONS as e:
                 # Opening the shared connection itself failed before any query in this
                 # group ran, so all of them are safe to retry next poll.
                 self._log.warning(

sqlserver/datadog_checks/sqlserver/statements.py

@@ -23,6 +23,7 @@
 from datadog_checks.base.utils.serialization import json
 from datadog_checks.base.utils.tracking import tracked_method
 from datadog_checks.sqlserver.config import SQLServerConfig
+from datadog_checks.sqlserver.connection_errors import EXPECTED_DB_EXCEPTIONS
 from datadog_checks.sqlserver.const import STATIC_INFO_ENGINE_EDITION, STATIC_INFO_VERSION
 from datadog_checks.sqlserver.utils import is_azure_sql_database, needs_comment_recovery, raise_if_cancelled
 
@@ -263,7 +264,7 @@ def __init__(self, check, config: SQLServerConfig):
             check,
             run_sync=is_affirmative(self._config.statement_metrics_config.get('run_sync', False)),
             enabled=is_affirmative(self._config.statement_metrics_config.get('enabled', True)),
-            expected_db_exceptions=(),
+            expected_db_exceptions=EXPECTED_DB_EXCEPTIONS,
             min_collection_interval=self._config.min_collection_interval,
             dbms=check.dbms,
             rate_limit=1 / float(collection_interval),
@@ -637,9 +638,32 @@ def _collect_plans(self, rows, cursor, deadline):
             # we use the plan handle
             if row['is_proc'] or row['is_encrypted']:
                 plan_key = row['plan_handle']
-            if self._seen_plans_ratelimiter.acquire(plan_key):
-                raise_if_cancelled(self._cancel_event)
+            # Check admission without consuming the plan's budget: a failed lookup below must not
+            # suppress this plan until the rate limiter entry expires.
+            if not self._seen_plans_ratelimiter.would_acquire(plan_key):
+                continue
+            raise_if_cancelled(self._cancel_event)
+            try:
                 raw_plan, is_plan_encrypted = self._load_plan(row['plan_handle'], cursor)
+            except Exception as e:
+                # A connection closed during cancellation may surface as a database error.
+                raise_if_cancelled(self._cancel_event)
+                self.log.debug(
+                    "Failed to load plan | query_signature=[%s] query_hash=[%s] query_plan_hash=[%s] "
+                    "plan_handle=[%s] err=[%s]",
+                    row['query_signature'],
+                    row['query_hash'],
+                    row['query_plan_hash'],
+                    row['plan_handle'],
+                    e,
+                )
+                self._check.count(
+                    "dd.sqlserver.statements.error",
+                    1,
+                    **self._check.debug_stats_kwargs(tags=["error:load-plan-{}".format(type(e))]),
+                )
+                continue
+            if self._seen_plans_ratelimiter.acquire(plan_key):
                 obfuscated_plan = None
                 collection_errors = []
 

sqlserver/tests/test_statements.py

@@ -1456,3 +1456,79 @@ def _mock_collect_plans(*_args, **_kwargs):
         mock_collect_plans.assert_called_once()
     else:
         mock_collect_plans.assert_not_called()
+
+
+def _plan_row(suffix: str) -> dict[str, object]:
+    return {
+        'query_signature': f'query-{suffix}',
+        'query_hash': f'query-hash-{suffix}',
+        'query_plan_hash': f'plan-hash-{suffix}',
+        'plan_handle': f'plan-handle-{suffix}',
+        'text': 'SELECT 1',
+        'dd_tables': [],
+        'dd_commands': ['SELECT'],
+        'dd_comments': [],
+        'database_name': 'master',
+        'is_proc': False,
+        'is_encrypted': False,
+        'procedure_signature': None,
+        'procedure_name': None,
+    }
+
+
+@pytest.mark.unit
+def test_plan_lookup_failure_allows_later_rows_and_retry(aggregator, instance_docker):
+    """A failed plan lookup does not stop later rows or suppress its retry."""
+    instance_docker['dbm'] = True
+    instance_docker['query_metrics'] = {
+        'enabled': True,
+        'run_sync': True,
+        'enforce_collection_interval_deadline': False,
+    }
+    check = SQLServer(CHECK_NAME, {}, [instance_docker])
+    failed_row = _plan_row('failed')
+    later_row = _plan_row('later')
+    plan = ('<ShowPlanXML/>', False)
+
+    with mock.patch.object(
+        check.statement_metrics,
+        '_load_plan',
+        side_effect=[RuntimeError('plan lookup timed out'), plan, plan],
+    ) as load_plan:
+        first_pass = list(check.statement_metrics._collect_plans([failed_row, later_row], mock.Mock(), float('inf')))
+        retry_pass = list(check.statement_metrics._collect_plans([failed_row], mock.Mock(), float('inf')))
+
+    assert [event['db']['query_signature'] for event in first_pass] == [later_row['query_signature']]
+    assert [event['db']['query_signature'] for event in retry_pass] == [failed_row['query_signature']]
+    assert [call.args[0] for call in load_plan.call_args_list] == [
+        failed_row['plan_handle'],
+        later_row['plan_handle'],
+        failed_row['plan_handle'],
+    ]
+    aggregator.assert_metric(
+        'dd.sqlserver.statements.error',
+        value=1,
+        tags=check.debug_tags() + ["error:load-plan-<class 'RuntimeError'>"],
+    )
+
+
+@pytest.mark.unit
+def test_plan_lookup_failure_during_cancellation_propagates(instance_docker):
+    """Cancellation during a failing plan lookup still aborts plan collection."""
+    instance_docker['dbm'] = True
+    instance_docker['query_metrics'] = {
+        'enabled': True,
+        'run_sync': True,
+        'enforce_collection_interval_deadline': False,
+    }
+    check = SQLServer(CHECK_NAME, {}, [instance_docker])
+
+    def cancel_during_lookup(*_args):
+        check.statement_metrics._cancel_event.set()
+        raise RuntimeError('connection closed')
+
+    with (
+        mock.patch.object(check.statement_metrics, '_load_plan', side_effect=cancel_during_lookup),
+        pytest.raises(Exception, match='Job loop cancelled'),
+    ):
+        list(check.statement_metrics._collect_plans([_plan_row('cancelled')], mock.Mock(), float('inf')))