Add incremental query metrics primitives to the base package (DataDog#24791)

* Copy the Postgres incremental query metrics primitives into base

MySQL needs DeltaDetector and ObfuscationLookup too. Copy them verbatim
first so the rest of the stack reads as a diff against the reviewed
Postgres code; git diff --no-index against the originals prints nothing.

Tests are lifted from postgres/tests/test_statements_v2.py with their
bodies unchanged. Postgres keeps its own copies until a follow-up PR.

* Generalize ObfuscationLookup over the statement key type

The cache never interprets its key, so PgssKey was incidental; MySQL
identifies a statement by a digest string. Make the class generic over K
and rename queryid_map_size to key_map_size.

The docstring no longer claims a rejection is permanent. That holds for
a MySQL digest, not in general, so the caller decides.

* Generalize DeltaDetector over the row key

Take a key callable instead of reading queryid, dbid and userid off each
row. PgssKey is gone and DeltaResult now exposes changed_keys and
vanished_keys.

The two DeltaResult field docstrings were also attached to the wrong
fields.

* Stop mutating the caller's rows when collapsing duplicates

The first row of each duplicate group was stored by reference and then
summed into, so the caller's snapshot came back rewritten: given rows
with calls=8 and calls=7, the first was left holding 15. Copy instead.

* Add a maxsize property that trims when the cache shrinks

Callers resize the cache from a server setting on every collection, and
with no public setter Postgres assigns _maxsize directly. That skips
trimming, so lowering the bound evicts nothing until a later populate.

* Return obfuscation failures from populate

A statement the obfuscator rejects was skipped silently and stayed a
miss, so every collection in which it changed fetched and re-attempted
it. Hand the failed keys back so callers can negative-cache them.

* Extract a module-level obfuscate_statement

Obfuscation was only reachable through the cache, which assumes a key
determines its text. MySQL's prepared_statements_instances is keyed on a
reusable address, so its rows must be obfuscated afresh each cycle.

* Add resolve_obfuscations to own the cache miss path

Postgres and MySQL wrap the cache in the same sequence, whose ordering
constraints fail quietly when broken. TextDisposition also makes the
integration say whether a rejected text is permanent (DDIGNORE, EXPLAIN)
or transient (<insufficient privilege>), which a chain of continue
statements blurs.

Emits no telemetry: counts come back in ResolveStats for the caller to
report under its own metric names.

* Add changelog entry

Co-authored-by: Cursor <cursoragent@cursor.com>

* Drive cache retention from the live keys, not the vanished ones

resolve_obfuscations bound the delta key and the cache key to one type
variable, but MySQL keys counters on (schema, digest) and caches on the
digest alone. Projecting vanished_keys down to a digest drops entries
another schema is still running, and since eviction precedes the lookup
the digest is re-fetched in the same cycle. Against a live MySQL 8.4 with
one statement shared by two schemas, the obvious projection re-fetched
the shared digest on 4 of 8 collections; the live set makes that 1.

So evict() becomes retain(live_keys) and a caller can only misstate what
is present, which shows up in the hit rate, rather than silently discard
an entry that is still needed. Sweeping the whole cache costs 0.55 ms at
maxsize=10000, and the count comes back as ResolveStats.dropped.

Retention no longer depends on the delta, so DeltaResult.vanished_keys
goes, along with a set difference the detector computed twice.

The docstrings claimed eviction stopped a returning key being served a
result cached against its previous incarnation. Neither source can do
that: a queryid and a digest are both derived from the normalized
statement, so a key cannot name two texts. Retention reclaims memory, and
saying so is what makes it clear it has to run on quiet collections too.

Co-authored-by: Cursor <cursoragent@cursor.com>

* Group the primitives into a query_metrics package and rename for clarity

The two modules were named for their mechanisms rather than their subject, which
left the vocabulary inconsistent with the rest of DBM and made the classification
enum read as an instruction to the base library rather than a description of what
the integration saw.

- db/query_metrics/ now holds the set, split into stats, obfuscation, cache and
  resolver. The package re-exports the public names lazily, mirroring db/__init__.py,
  so the file layout is not part of the interface.
- DeltaDetector.compute() becomes QueryStats.diff(snapshot), returning a Delta.
  "Query stats" is what every source calls itself (pg_stat_statements,
  dm_exec_query_stats, $queryStats), and it keeps stats (what the database exposes)
  distinct from metrics (what we emit). derivative_rows becomes rows, and
  metric_columns becomes counter_columns to match.
- TextDisposition becomes TextKind, with STATEMENT/EXCLUDED/UNAVAILABLE replacing
  CACHE/IGNORE/SKIP. Integrations now report what a text turned out to be and the
  resolver owns the caching policy those kinds imply, rather than each integration
  re-deriving whether a rejection is permanent. The kinds split on whether the text
  is the statement's own and, if not, whether that can change; text that is
  permanently unavailable would need a kind of its own.

Co-authored-by: Cursor <cursoragent@cursor.com>

* Keep derivative_rows as the name of the delta rows

rows is an actual pg_stat_statements counter column, so Delta.rows produced
expressions like delta.rows[0]['rows'] at the places worth reading carefully.

Co-authored-by: Cursor <cursoragent@cursor.com>

* Add DBM as a codeowner for the query_metrics package

* Reclaim obfuscation results whose last key is gone

retain() dropped the key mappings but left the results they named, and the
results are what hold the obfuscated text, so nothing that costs memory was
reclaimed. Because only an overflow evicts a result, the stranded ones
accumulated until the result tier sat at maxsize however small the live set
was, and then displaced results of statements still live but not recently run.
With four keys per signature and two live signatures, a 50-entry cache filled
from cycle 48 and evicted a statement present in every snapshot.

Keys also leave by LRU trimming, so the sweep runs on every retain rather than
only when a stale key was dropped. At maxsize=10000 with nothing shared it
costs 0.66 ms, against 0.40 ms for the stale key set it follows.

Co-authored-by: Cursor <cursoragent@cursor.com>

---------

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: eric-weaver

.github/CODEOWNERS

@@ -154,6 +154,8 @@ datadog_checks_base/tests/**/test_util.py           @DataDog/agent-integrations
 datadog_checks_base/tests/**/test_db_sql.py         @DataDog/database-monitoring-agent @DataDog/agent-integrations
 **/base/utils/db/statement_metrics.py               @DataDog/database-monitoring-agent @DataDog/agent-integrations
 datadog_checks_base/tests/**/test_db_statements.py  @DataDog/database-monitoring-agent @DataDog/agent-integrations
+**/base/utils/db/query_metrics/                     @DataDog/database-monitoring-agent @DataDog/agent-integrations
+datadog_checks_base/tests/**/query_metrics/         @DataDog/database-monitoring-agent @DataDog/agent-integrations
 /postgres/                                          @DataDog/database-monitoring-agent @DataDog/agent-integrations
 /postgres/*.md                                      @DataDog/database-monitoring @DataDog/agent-integrations @DataDog/documentation
 /postgres/manifest.json                             @DataDog/database-monitoring @DataDog/agent-integrations @DataDog/documentation

datadog_checks_base/changelog.d/24791.added

@@ -0,0 +1 @@
+Add QueryStats, ObfuscationLookup and resolve_obfuscations for incremental query metrics collection.

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

@@ -0,0 +1,6 @@
+# (C) Datadog, Inc. 2026-present
+# All rights reserved
+# Licensed under a 3-clause BSD style license (see LICENSE)
+import lazy_loader
+
+__getattr__, __dir__, __all__ = lazy_loader.attach_stub(__name__, __file__)

datadog_checks_base/datadog_checks/base/utils/db/query_metrics/__init__.pyi

@@ -0,0 +1,19 @@
+# (C) Datadog, Inc. 2026-present
+# All rights reserved
+# Licensed under a 3-clause BSD style license (see LICENSE)
+from .cache import ObfuscationLookup
+from .obfuscation import ObfuscationResult, obfuscate_statement
+from .resolver import ResolveResult, ResolveStats, TextKind, resolve_obfuscations
+from .stats import Delta, QueryStats
+
+__all__ = [
+    'Delta',
+    'ObfuscationLookup',
+    'ObfuscationResult',
+    'QueryStats',
+    'ResolveResult',
+    'ResolveStats',
+    'TextKind',
+    'obfuscate_statement',
+    'resolve_obfuscations',
+]

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

@@ -0,0 +1,247 @@
+# (C) Datadog, Inc. 2026-present
+# All rights reserved
+# Licensed under a 3-clause BSD style license (see LICENSE)
+from __future__ import annotations
+
+import logging
+from collections import OrderedDict
+from collections.abc import Hashable
+
+from .obfuscation import ObfuscationResult, obfuscate_statement
+
+logger = logging.getLogger(__name__)
+
+
+class ObfuscationLookup[K: Hashable]:
+    """LRU cache mapping statement identity keys to obfuscated query results.
+
+    ``K`` is whatever identifies a statement in the source table: a
+    ``(queryid, dbid, userid)`` triple for ``pg_stat_statements``, a digest string for MySQL's
+    ``events_statements_summary_by_digest``, and so on. The cache never interprets the key.
+
+    Caching is only sound where the key functionally determines the text, which is why both of
+    those work: a queryid is derived from the normalized parse tree and a digest from the
+    normalized token stream, so neither can name two statements. A source whose identity outlives
+    the text it named, such as MySQL's ``prepared_statements_instances``, must go through
+    :func:`~.obfuscation.obfuscate_statement` instead.
+
+    A lookup resolves a key to one of three outcomes:
+
+    - hit: the obfuscated result is cached, avoiding both the text fetch and FFI obfuscation.
+      Stored as two tiers, key -> query_signature -> result, so multiple keys sharing a
+      query_signature share one result.
+    - miss: nothing is cached for the key; the caller must fetch its text and pass it to
+      :meth:`populate` to obfuscate, store, and discard the raw text.
+    - ignored: the key is known to resolve to nothing usable. These are neither hit nor miss;
+      lookup skips them so they are never fetched again.
+
+    The cache does not decide what is non-cacheable; :meth:`mark_ignored` is driven by
+    :func:`~.resolver.resolve_obfuscations`, which owns that policy. The cache owns only storage
+    and lifecycle. All three tiers are LRU-bounded by ``maxsize``, and :meth:`retain` drops the
+    keys that have left the source table, along with any result no live key still names, so entries
+    do not sit at that bound long after the statements they describe.
+    """
+
+    def __init__(self, maxsize: int, obfuscate_options: str, log_unobfuscated_queries: bool = False):
+        self._maxsize = maxsize
+        self._obfuscate_options = obfuscate_options
+        self._log_unobfuscated_queries = log_unobfuscated_queries
+
+        self._key_to_sig: OrderedDict[K, str] = OrderedDict()
+        self._sig_to_result: OrderedDict[str, ObfuscationResult] = OrderedDict()
+        # Negative cache: keys we have learned resolve to nothing cacheable.
+        self._ignored_keys: OrderedDict[K, None] = OrderedDict()
+
+        self._hits = 0
+        self._misses = 0
+
+    @property
+    def maxsize(self) -> int:
+        return self._maxsize
+
+    @maxsize.setter
+    def maxsize(self, value: int):
+        self._maxsize = value
+        self._trim()
+
+    @property
+    def key_map_size(self) -> int:
+        return len(self._key_to_sig)
+
+    @property
+    def signature_map_size(self) -> int:
+        return len(self._sig_to_result)
+
+    @property
+    def ignored_map_size(self) -> int:
+        return len(self._ignored_keys)
+
+    @property
+    def hits(self) -> int:
+        return self._hits
+
+    @property
+    def misses(self) -> int:
+        return self._misses
+
+    def reset_stats(self):
+        self._hits = 0
+        self._misses = 0
+
+    def lookup(self, keys: set[K]) -> tuple[dict[K, ObfuscationResult], set[K]]:
+        """Return (hits, misses) for the given statement keys.
+
+        Keys in the negative cache are excluded from both: they are neither a hit
+        (no result to return) nor a miss (must not be re-fetched).
+        """
+        hits: dict[K, ObfuscationResult] = {}
+        misses: set[K] = set()
+        ignored = 0
+
+        for key in keys:
+            if key in self._ignored_keys:
+                self._ignored_keys.move_to_end(key)
+                ignored += 1
+                continue
+            sig = self._key_to_sig.get(key)
+            if sig is not None:
+                self._key_to_sig.move_to_end(key)
+                result = self._sig_to_result.get(sig)
+                if result is not None:
+                    self._sig_to_result.move_to_end(sig)
+                    self._hits += 1
+                    hits[key] = result
+                    continue
+            self._misses += 1
+            misses.add(key)
+
+        logger.debug(
+            "lookup: requested=%d hits=%d misses=%d ignored=%d key_map=%d sig_map=%d ignored_map=%d",
+            len(keys),
+            len(hits),
+            len(misses),
+            ignored,
+            len(self._key_to_sig),
+            len(self._sig_to_result),
+            len(self._ignored_keys),
+        )
+        return hits, misses
+
+    def mark_ignored(self, keys: set[K]) -> None:
+        """Record keys that resolve to nothing usable so future lookups skip them.
+
+        Entries are forgotten via :meth:`retain` once their key disappears from the source table.
+        """
+        for key in keys:
+            # Drop any stale positive mapping so an ignored key can never resurface as a
+            # hit (e.g. if its signature is later repopulated by another key after this
+            # negative entry is LRU-trimmed).
+            self._key_to_sig.pop(key, None)
+            self._ignored_keys[key] = None
+            self._ignored_keys.move_to_end(key)
+        if keys:
+            self._trim_ignored()
+            logger.debug("mark_ignored: added=%d ignored_map=%d", len(keys), len(self._ignored_keys))
+
+    def populate(self, raw_texts: dict[K, str]) -> tuple[dict[K, ObfuscationResult], set[K]]:
+        """Obfuscate raw texts and store the results.
+
+        Returns (results, failures), where failures are the keys whose text could not be
+        obfuscated. Obfuscation depends only on the text, so a failure will recur for as long as
+        the key keeps resolving to that text; callers that know the text is stable should pass
+        these to :meth:`mark_ignored` rather than re-fetching them every collection.
+        """
+        results: dict[K, ObfuscationResult] = {}
+        failures: set[K] = set()
+
+        for key, raw_text in raw_texts.items():
+            result = self._obfuscate_single(raw_text)
+            if result is None:
+                failures.add(key)
+                continue
+
+            self._key_to_sig[key] = result.query_signature
+            self._trim_keys()
+
+            if result.query_signature not in self._sig_to_result:
+                self._sig_to_result[result.query_signature] = result
+                self._trim_sig()
+            else:
+                self._sig_to_result.move_to_end(result.query_signature)
+
+            results[key] = result
+
+        logger.debug(
+            "populate: input=%d obfuscated=%d failed=%d key_map=%d sig_map=%d",
+            len(raw_texts),
+            len(results),
+            len(failures),
+            len(self._key_to_sig),
+            len(self._sig_to_result),
+        )
+        return results, failures
+
+    def retain(self, live_keys: set[K]) -> int:
+        """Forget all state, positive and negative, for keys absent from *live_keys*.
+
+        Returns how many keys were dropped.
+
+        Callers pass the keys currently in the source table rather than the ones that left, so a
+        caller whose cache key is a projection of a wider counter key cannot report a key as gone
+        while another live row still needs it. MySQL is the case in point: its counters are keyed
+        on ``(schema, digest)`` but one digest has one text, so a digest is only finished once no
+        schema references it.
+
+        This reclaims memory rather than protecting correctness. Because a key determines its
+        text, an entry that outlives its key is unreachable rather than wrong. It does need to run
+        on every collection, including quiet ones, or the cache sits at ``maxsize`` indefinitely.
+
+        Both tiers are pruned, and a result goes once no live key names it. The results are what
+        hold the obfuscated text, so pruning the key mappings alone would reclaim nothing worth
+        reclaiming; and because only an overflow evicts a result, any left stranded accumulate
+        until the tier is full however small the live set is, and then displace the results of
+        statements that are still live but have not run recently.
+        """
+        stale = (self._key_to_sig.keys() | self._ignored_keys.keys()) - live_keys
+        for key in stale:
+            self._key_to_sig.pop(key, None)
+            self._ignored_keys.pop(key, None)
+
+        # Results are shared by every key with the same signature, so one is reclaimable only once
+        # no live key names it. Keys also leave by LRU trimming, so this sweeps on every call
+        # rather than only when retain itself dropped one.
+        orphaned = self._sig_to_result.keys() - set(self._key_to_sig.values())
+        for signature in orphaned:
+            del self._sig_to_result[signature]
+
+        if stale or orphaned:
+            logger.debug(
+                "retain: live=%d dropped=%d orphaned=%d key_map=%d sig_map=%d ignored_map=%d",
+                len(live_keys),
+                len(stale),
+                len(orphaned),
+                len(self._key_to_sig),
+                len(self._sig_to_result),
+                len(self._ignored_keys),
+            )
+        return len(stale)
+
+    def _obfuscate_single(self, raw_text: str) -> ObfuscationResult | None:
+        return obfuscate_statement(raw_text, self._obfuscate_options, self._log_unobfuscated_queries)
+
+    def _trim(self):
+        self._trim_keys()
+        self._trim_sig()
+        self._trim_ignored()
+
+    def _trim_keys(self):
+        while len(self._key_to_sig) > self._maxsize:
+            self._key_to_sig.popitem(last=False)
+
+    def _trim_sig(self):
+        while len(self._sig_to_result) > self._maxsize:
+            self._sig_to_result.popitem(last=False)
+
+    def _trim_ignored(self):
+        while len(self._ignored_keys) > self._maxsize:
+            self._ignored_keys.popitem(last=False)

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

@@ -0,0 +1,51 @@
+# (C) Datadog, Inc. 2026-present
+# All rights reserved
+# Licensed under a 3-clause BSD style license (see LICENSE)
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+
+from datadog_checks.base.utils.db.sql import compute_sql_signature
+from datadog_checks.base.utils.db.utils import obfuscate_sql_with_metadata
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True, slots=True)
+class ObfuscationResult:
+    obfuscated_query: str
+    query_signature: str
+    tables: list[str] | None
+    commands: list[str] | None
+    comments: list[str] | None
+
+
+def obfuscate_statement(
+    raw_text: str, obfuscate_options: str, log_unobfuscated_queries: bool = False
+) -> ObfuscationResult | None:
+    """Obfuscate one statement via the FFI, returning None if it cannot be obfuscated.
+
+    Exposed separately from :class:`~.cache.ObfuscationLookup` for statement sources whose identity
+    does not determine their text, which therefore cannot be cached. MySQL's
+    ``prepared_statements_instances`` is one: it is keyed on a reusable memory address, so a
+    recycled instance can carry unrelated text and every row has to be obfuscated afresh.
+    """
+    try:
+        statement = obfuscate_sql_with_metadata(raw_text, obfuscate_options)
+    except Exception as e:
+        if log_unobfuscated_queries:
+            logger.warning("Failed to obfuscate query=[%s] | err=[%s]", raw_text, e)
+        else:
+            logger.debug("Failed to obfuscate query | err=[%s]", e)
+        return None
+
+    obfuscated_query = statement['query']
+    metadata = statement['metadata']
+    return ObfuscationResult(
+        obfuscated_query=obfuscated_query,
+        query_signature=compute_sql_signature(obfuscated_query),
+        tables=metadata.get('tables', None),
+        commands=metadata.get('commands', None),
+        comments=metadata.get('comments', None),
+    )