Replace runtime regex with static cluster-aware query variants

Instead of rewriting system-table metric queries with a regex at runtime,
define static cluster-aware query variants and select them in get_queries()
based on single_endpoint_mode.

- Add cluster_aware_query() helper in utils.py that reuses the base query's
  columns by reference (no duplication) and appends the clickhouse_node tag.
- Define legacy variants in queries.py and advanced variants in
  advanced_queries (JSON-backed via load_match_query + SystemErrorsClusterAware).
- Remove import re, SYSTEM_TABLE_FROM_CLAUSE, and make_cluster_aware().

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: sangeetashivaji

clickhouse/datadog_checks/clickhouse/advanced_queries/__init__.py

@@ -59,7 +59,18 @@
 import os
 from typing import Any
 
-__all__ = ['SystemAsynchronousMetrics', 'SystemErrors', 'SystemEvents', 'SystemMetrics']
+from datadog_checks.clickhouse.utils import cluster_aware_query
+
+__all__ = [
+    'SystemAsynchronousMetrics',
+    'SystemAsynchronousMetricsClusterAware',
+    'SystemErrors',
+    'SystemErrorsClusterAware',
+    'SystemEvents',
+    'SystemEventsClusterAware',
+    'SystemMetrics',
+    'SystemMetricsClusterAware',
+]
 
 DATA_DIR = os.path.join(os.path.dirname(os.path.dirname(__file__)), 'data')
 
@@ -69,6 +80,10 @@
     'SystemAsynchronousMetrics': 'system_async_metrics',
 }
 
+# Suffix that selects the single-endpoint-mode variant of a bulk match query (e.g.
+# 'SystemEventsClusterAware') through __getattr__.
+CLUSTER_AWARE_SUFFIX = 'ClusterAware'
+
 _match_query_cache: dict[str, dict[str, Any]] = {}
 
 SystemErrors: dict[str, Any] = {
@@ -82,14 +97,18 @@
     ],
 }
 
+SystemErrorsClusterAware: dict[str, Any] = cluster_aware_query(
+    SystemErrors, 'value, name, code, remote', 'errors', where='WHERE value > 0'
+)
+
 
-def load_match_query(name: str) -> dict[str, Any]:
+def load_match_query(name: str, cluster_aware: bool = False) -> dict[str, Any]:
     """Read ``data/<name>.json`` and reconstitute the QueryManager-shaped dict."""
     try:
         with open(os.path.join(DATA_DIR, f'{name}.json'), encoding='utf-8') as f:
             spec = json.load(f)
         items = _expand_match_items(spec['items'], spec['prefix'])
-        return {
+        base = {
             'name': spec['name'],
             'query': spec['query'],
             'columns': [
@@ -102,13 +121,18 @@ def load_match_query(name: str) -> dict[str, Any]:
                 },
             ],
         }
+        if cluster_aware:
+            # value_column/match_column are logical labels, not SQL identifiers, so derive the
+            # real SELECT list and table from the (fixed-shape) generated query string.
+            select, _, tail = spec['query'].partition(' FROM system.')
+            table, _, where = tail.partition(' ')
+            return cluster_aware_query(base, select.removeprefix('SELECT '), table, where=where)
+        return base
     except (OSError, json.JSONDecodeError, KeyError, TypeError, AttributeError) as exc:
         raise RuntimeError(f'failed to load advanced query {name!r}') from exc
 
 
-def _expand_match_items(
-    compact: dict[str, list[str] | dict[str, str]], prefix: str
-) -> dict[str, dict[str, Any]]:
+def _expand_match_items(compact: dict[str, list[str] | dict[str, str]], prefix: str) -> dict[str, dict[str, Any]]:
     """Expand the compact ``{type: keys | {key: scale}}`` map to the per-entry dict shape."""
     merged: dict[str, dict[str, Any]] = {}
     for type_name, group in compact.items():
@@ -129,8 +153,10 @@ def warm_cache() -> None:
 
 
 def __getattr__(name: str) -> dict[str, Any]:
-    if name not in MATCH_QUERIES:
+    cluster_aware = name.endswith(CLUSTER_AWARE_SUFFIX)
+    base_name = name[: -len(CLUSTER_AWARE_SUFFIX)] if cluster_aware else name
+    if base_name not in MATCH_QUERIES:
         raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
     if name not in _match_query_cache:
-        _match_query_cache[name] = load_match_query(MATCH_QUERIES[name])
+        _match_query_cache[name] = load_match_query(MATCH_QUERIES[base_name], cluster_aware=cluster_aware)
     return _match_query_cache[name]

clickhouse/datadog_checks/clickhouse/clickhouse.py

@@ -1,7 +1,6 @@
 # (C) Datadog, Inc. 2019-present
 # All rights reserved
 # Licensed under a 3-clause BSD style license (see LICENSE)
-import re
 from string import Template
 from time import time
 
@@ -36,14 +35,6 @@
 # Database instance collection interval in seconds (not user-configurable)
 DATABASE_INSTANCE_COLLECTION_INTERVAL = 300
 
-# Tag added to per-node metrics when collecting from all replicas in single endpoint mode.
-CLUSTER_NODE_TAG = 'clickhouse_node'
-
-# Matches the FROM clause of the standard system-table metric queries so they can be retargeted at
-# clusterAllReplicas() and tagged per node in single endpoint mode. The leading whitespace is part of
-# the match so the per-node projection can be spliced in without leaving a gap before the comma.
-SYSTEM_TABLE_FROM_CLAUSE = re.compile(r'\s+FROM\s+system\.(?P<table>\w+)', re.IGNORECASE)
-
 
 class ClickhouseCheck(DatabaseCheck):
     __NAMESPACE__ = 'clickhouse'
@@ -299,13 +290,15 @@ def check(self, _):
     def get_queries(self) -> list[dict]:
         query_list = []
 
+        single = self._config.single_endpoint_mode
+
         if self._config.use_legacy_queries:
             query_list.extend(
                 [
-                    self.make_cluster_aware(queries.SystemMetrics),
-                    self.make_cluster_aware(queries.SystemEventsToDeprecate),
-                    self.make_cluster_aware(queries.SystemEvents),
-                    self.make_cluster_aware(queries.SystemAsynchronousMetrics),
+                    queries.SystemMetricsClusterAware if single else queries.SystemMetrics,
+                    queries.SystemEventsToDeprecateClusterAware if single else queries.SystemEventsToDeprecate,
+                    queries.SystemEventsClusterAware if single else queries.SystemEvents,
+                    queries.SystemAsynchronousMetricsClusterAware if single else queries.SystemAsynchronousMetrics,
                     queries.SystemParts,
                     queries.SystemReplicas,
                     queries.SystemDictionaries,
@@ -315,13 +308,17 @@ def get_queries(self) -> list[dict]:
         if self._config.use_advanced_queries:
             query_list.extend(
                 [
-                    self.make_cluster_aware(advanced_queries.SystemMetrics),
-                    self.make_cluster_aware(advanced_queries.SystemEvents),
-                    self.make_cluster_aware(advanced_queries.SystemAsynchronousMetrics),
+                    advanced_queries.SystemMetricsClusterAware if single else advanced_queries.SystemMetrics,
+                    advanced_queries.SystemEventsClusterAware if single else advanced_queries.SystemEvents,
+                    advanced_queries.SystemAsynchronousMetricsClusterAware
+                    if single
+                    else advanced_queries.SystemAsynchronousMetrics,
                 ]
             )
             if self.version_ge('21.3'):
-                query_list.append(self.make_cluster_aware(advanced_queries.SystemErrors))
+                query_list.append(
+                    advanced_queries.SystemErrorsClusterAware if single else advanced_queries.SystemErrors
+                )
 
         return query_list
 
@@ -457,25 +454,6 @@ def get_system_table(self, table_name):
             # Direct connection: Query the local system table directly
             return f"system.{table_name}"
 
-    def make_cluster_aware(self, query: dict) -> dict:
-        """Retarget a system-table query at clusterAllReplicas and tag each row per node."""
-        if not self._config.single_endpoint_mode:
-            return query
-
-        def rewrite(match):
-            cluster_table = self.get_system_table(match.group('table'))
-            return f', hostName() AS {CLUSTER_NODE_TAG} FROM {cluster_table}'
-
-        new_query, replaced = SYSTEM_TABLE_FROM_CLAUSE.subn(rewrite, query['query'], count=1)
-        if not replaced:
-            return query
-
-        return {
-            **query,
-            'query': new_query,
-            'columns': [*query['columns'], {'name': CLUSTER_NODE_TAG, 'type': 'tag'}],
-        }
-
     def ping_clickhouse(self):
         return self._client.ping()
 

clickhouse/datadog_checks/clickhouse/queries.py

@@ -1,7 +1,7 @@
 # (C) Datadog, Inc. 2019-present
 # All rights reserved
 # Licensed under a 3-clause BSD style license (see LICENSE)
-from .utils import compact_query
+from .utils import cluster_aware_query, compact_query
 
 # https://clickhouse.com/docs/operations/system-tables/metrics
 SystemMetrics = {
@@ -2059,3 +2059,14 @@
         {'name': 'dictionary.load', 'type': 'gauge'},
     ],
 }
+
+
+# Cluster-aware variants used in single endpoint mode: they read all replicas via
+# clusterAllReplicas() and tag each row per node so per-node counters stay distinct.
+# SystemParts/SystemReplicas/SystemDictionaries use GROUP BY and are intentionally excluded.
+SystemMetricsClusterAware = cluster_aware_query(SystemMetrics, 'value, metric', 'metrics')
+SystemEventsToDeprecateClusterAware = cluster_aware_query(SystemEventsToDeprecate, 'value, event', 'events')
+SystemEventsClusterAware = cluster_aware_query(SystemEvents, 'value, event', 'events')
+SystemAsynchronousMetricsClusterAware = cluster_aware_query(
+    SystemAsynchronousMetrics, 'value, metric', 'asynchronous_metrics'
+)

clickhouse/datadog_checks/clickhouse/utils.py

@@ -27,5 +27,22 @@ def compact_query(query):
     return re.sub(r'\n\s+', ' ', query.strip())
 
 
+# Tag added to per-node metrics when collecting from all replicas in single endpoint mode.
+CLUSTER_NODE_TAG = 'clickhouse_node'
+
+
+def cluster_aware_query(base: dict, select: str, table: str, where: str = '') -> dict:
+    """Build a cluster-aware variant that reads all replicas and tags each row per node."""
+    tail = f' {where}' if where else ''
+    return {
+        'name': base['name'],
+        'query': (
+            f"SELECT {select}, hostName() AS {CLUSTER_NODE_TAG} "
+            f"FROM clusterAllReplicas('default', system.{table}){tail}"
+        ),
+        'columns': [*base['columns'], {'name': CLUSTER_NODE_TAG, 'type': 'tag'}],
+    }
+
+
 def parse_version(version: str) -> list[int]:
     return [int(v) for v in version.split('.')]

clickhouse/tests/test_unit.py

@@ -383,42 +383,39 @@ def test_database_hostname_ignores_reported_hostname_override(reported_hostname,
         mock_resolve.assert_called_with(BASE_INSTANCE['server'])
 
 
-def test_make_cluster_aware_direct_connection_unchanged(instance):
-    """Direct connections keep the bare per-node query (SDBM-2746)."""
-    instance = {**instance, 'single_endpoint_mode': False}
-    check = ClickhouseCheck('clickhouse', {}, [instance])
-
-    assert check.make_cluster_aware(advanced_queries.SystemEvents) is advanced_queries.SystemEvents
-
-
-def test_make_cluster_aware_bulk_match_query(instance):
-    """Single endpoint mode tags system.events per node so counters don't flap across the LB (SDBM-2746)."""
-    instance = {**instance, 'single_endpoint_mode': True}
-    check = ClickhouseCheck('clickhouse', {}, [instance])
-
-    rewritten = check.make_cluster_aware(advanced_queries.SystemEvents)
+def test_cluster_aware_variant_bulk_match_query():
+    """The cluster-aware variant reads all replicas and tags system.events per node (SDBM-2746)."""
+    variant = advanced_queries.SystemEventsClusterAware
 
-    assert rewritten['query'] == (
+    assert variant['query'] == (
         "SELECT value, event, hostName() AS clickhouse_node FROM clusterAllReplicas('default', system.events)"
     )
-    assert rewritten['columns'][-1] == {'name': 'clickhouse_node', 'type': 'tag'}
-    # The original cached query dict must not be mutated.
+    assert variant['columns'][-1] == {'name': 'clickhouse_node', 'type': 'tag'}
+    # The base query dict must not be mutated by building the variant.
     assert advanced_queries.SystemEvents['query'] == 'SELECT value, event FROM system.events'
     assert all(column['name'] != 'clickhouse_node' for column in advanced_queries.SystemEvents['columns'])
 
 
-def test_make_cluster_aware_preserves_where_clause(instance):
-    """system.errors carries a WHERE clause that must survive the rewrite."""
-    instance = {**instance, 'single_endpoint_mode': True}
-    check = ClickhouseCheck('clickhouse', {}, [instance])
-
-    rewritten = check.make_cluster_aware(advanced_queries.SystemErrors)
+def test_cluster_aware_variant_preserves_where_clause():
+    """system.errors carries a WHERE clause that must survive in the cluster-aware variant."""
+    variant = advanced_queries.SystemErrorsClusterAware
 
-    assert rewritten['query'] == (
+    assert variant['query'] == (
         "SELECT value, name, code, remote, hostName() AS clickhouse_node "
         "FROM clusterAllReplicas('default', system.errors) WHERE value > 0"
     )
-    assert rewritten['columns'][-1] == {'name': 'clickhouse_node', 'type': 'tag'}
+    assert variant['columns'][-1] == {'name': 'clickhouse_node', 'type': 'tag'}
+
+
+def test_cluster_aware_variant_legacy_query():
+    """Legacy queries.py exposes matching static cluster-aware variants."""
+    variant = queries.SystemMetricsClusterAware
+
+    assert variant['query'] == (
+        "SELECT value, metric, hostName() AS clickhouse_node FROM clusterAllReplicas('default', system.metrics)"
+    )
+    assert variant['columns'][-1] == {'name': 'clickhouse_node', 'type': 'tag'}
+    assert queries.SystemMetrics['query'] == 'SELECT value, metric FROM system.metrics'
 
 
 @pytest.mark.parametrize('use_advanced_queries', [True, False])
@@ -443,3 +440,17 @@ def test_get_queries_tags_system_tables_per_node_in_single_endpoint_mode(instanc
     # system.parts/replicas/dictionaries use GROUP BY and are intentionally left untouched here.
     if not use_advanced_queries:
         assert any(q is queries.SystemParts for q in check.get_queries())
+
+
+@pytest.mark.parametrize('use_advanced_queries', [True, False])
+def test_get_queries_uses_base_queries_for_direct_connection(instance, use_advanced_queries):
+    instance = {
+        **instance,
+        'single_endpoint_mode': False,
+        'use_advanced_queries': use_advanced_queries,
+        'use_legacy_queries': not use_advanced_queries,
+    }
+    check = ClickhouseCheck('clickhouse', {}, [instance])
+    check._server_version = '24.8'
+
+    assert all('clusterAllReplicas' not in q['query'] for q in check.get_queries())