[DBMON-6881] Move clickhouse onto the DatabaseCheck job registry and cancellation lifecycle

clickhouse/changelog.d/24934.fixed

@@ -0,0 +1 @@
+Manage the DBM async jobs through the ``DatabaseCheck`` registry.

clickhouse/datadog_checks/clickhouse/clickhouse.py

@@ -106,64 +106,55 @@ def __init__(self, name, init_config, instances):
             ca_cert=self._config.tls_ca_cert,
         )
 
-        # Initialize DBM components if enabled
-        self._init_dbm_components()
-
-    def _init_dbm_components(self):
-        """Initialize DBM components based on typed configuration."""
-        # Initialize query metrics (from system.query_log - analogous to pg_stat_statements)
-        if self._config.dbm and self._config.query_metrics.enabled:
-            self.statement_metrics = ClickhouseStatementMetrics(self, self._config.query_metrics)
-        else:
-            self.statement_metrics = None
-
-        # Initialize query samples (from system.processes - analogous to pg_stat_activity).
-        # The async insert buffer snapshot collapses into this job (sharing its connection and
-        # loop) instead of running as its own DBMAsyncJob, which would open another concurrent
-        # connection against the check's capped DBM connection pool.
-        if self._config.dbm and (
-            self._config.query_samples.enabled or self._config.collect_pending_async_inserts.enabled
-        ):
-            self.statement_samples = ClickhouseStatementSamples(
-                self, self._config.query_samples, self._config.collect_pending_async_inserts
+        self.statement_metrics: ClickhouseStatementMetrics | None = None
+        self.statement_samples: ClickhouseStatementSamples | None = None
+        self.query_completions: ClickhouseQueryCompletions | None = None
+        self.query_errors: ClickhouseQueryErrors | None = None
+        self.table_metrics: ClickhouseTableMetrics | None = None
+        self.metadata: ClickhouseMetadata | None = None
+        self.parts_and_merges: ClickhousePartsAndMerges | None = None
+        self._register_async_jobs()
+
+    def _register_async_jobs(self):
+        """Build and register the async jobs enabled by this check's configuration."""
+        if not self._config.dbm:
+            return
+
+        # Query metrics (from system.query_log)
+        if self._config.query_metrics.enabled:
+            self.statement_metrics = self.register_async_job(
+                ClickhouseStatementMetrics(self, self._config.query_metrics)
             )
-        else:
-            self.statement_samples = None
-
-        # Initialize query completions (from system.query_log - completed queries).
-        # The async insert flush log collection collapses into this job (shares its connection and loop),
-        # so its config is passed in here rather than run as its own DBMAsyncJob, which would add another
-        # concurrent connection to the check's capped DBM connection pool.
-        if self._config.dbm and (self._config.query_completions.enabled or self._config.collect_async_inserts.enabled):
-            self.query_completions = ClickhouseQueryCompletions(
-                self, self._config.query_completions, self._config.collect_async_inserts
+
+        # Query samples (from system.processes) and pending async inserts (system.asynchronous_inserts)
+        if self._config.query_samples.enabled or self._config.collect_pending_async_inserts.enabled:
+            self.statement_samples = self.register_async_job(
+                ClickhouseStatementSamples(self, self._config.query_samples, self._config.collect_pending_async_inserts)
             )
-        else:
-            self.query_completions = None
 
-        # Initialize query errors (from system.query_log - failed queries)
-        if self._config.dbm and self._config.query_errors.enabled:
-            self.query_errors = ClickhouseQueryErrors(self, self._config.query_errors)
-        else:
-            self.query_errors = None
+        # Completed queries and async insert flushes (from system.query_log and system.asynchronous_insert_log)
+        if self._config.query_completions.enabled or self._config.collect_async_inserts.enabled:
+            self.query_completions = self.register_async_job(
+                ClickhouseQueryCompletions(self, self._config.query_completions, self._config.collect_async_inserts)
+            )
 
-        # Initialize schema metrics (per-table size and per-view refresh gauges)
-        if self._config.dbm and self._config.schema_metrics.enabled:
-            self.table_metrics = ClickhouseTableMetrics(self, self._config.schema_metrics)
-        else:
-            self.table_metrics = None
+        # Failed queries (from system.query_log)
+        if self._config.query_errors.enabled:
+            self.query_errors = self.register_async_job(ClickhouseQueryErrors(self, self._config.query_errors))
 
-        # Initialize schema collection (catalog metadata for Schema Explorer)
-        if self._config.dbm and self._config.collect_schemas.enabled:
-            self.metadata = ClickhouseMetadata(self)
-        else:
-            self.metadata = None
+        # Schema metrics (from system.tables and system.view_refreshes)
+        if self._config.schema_metrics.enabled:
+            self.table_metrics = self.register_async_job(ClickhouseTableMetrics(self, self._config.schema_metrics))
 
-        # Initialize parts and merges monitoring (from system.parts, merges, mutations, replication_queue)
-        if self._config.dbm and self._config.parts_and_merges.enabled:
-            self.parts_and_merges = ClickhousePartsAndMerges(self, self._config.parts_and_merges)
-        else:
-            self.parts_and_merges = None
+        # Schema collection (from system.tables and system.columns)
+        if self._config.collect_schemas.enabled:
+            self.metadata = self.register_async_job(ClickhouseMetadata(self))
+
+        # Parts and merges (from system.parts, merges, mutations, replication_queue)
+        if self._config.parts_and_merges.enabled:
+            self.parts_and_merges = self.register_async_job(
+                ClickhousePartsAndMerges(self, self._config.parts_and_merges)
+            )
 
     def _add_core_tags(self):
         """
@@ -275,33 +266,7 @@ def check(self, _):
         # Send database instance metadata
         self._send_database_instance_metadata()
 
-        # Run query metrics collection if DBM is enabled (from system.query_log)
-        if self.statement_metrics:
-            self.statement_metrics.run_job_loop(self.tags)
-
-        # Run query samples collection if DBM is enabled (from system.processes)
-        if self.statement_samples:
-            self.statement_samples.run_job_loop(self.tags)
-
-        # Run query completions if DBM is enabled (from system.query_log)
-        if self.query_completions:
-            self.query_completions.run_job_loop(self.tags)
-
-        # Run query errors if DBM is enabled (from system.query_log - failed queries)
-        if self.query_errors:
-            self.query_errors.run_job_loop(self.tags)
-
-        # Run schema metrics (per-table size and per-view refresh gauges) if enabled
-        if self.table_metrics:
-            self.table_metrics.run_job_loop(self.tags)
-
-        # Run schema collection if enabled
-        if self.metadata:
-            self.metadata.run_job_loop(self.tags)
-
-        # Run parts and merges monitoring if enabled
-        if self.parts_and_merges:
-            self.parts_and_merges.run_job_loop(self.tags)
+        self.run_async_jobs(self.tags)
 
     def get_queries(self) -> list[dict]:
         query_list = []
@@ -603,55 +568,19 @@ def create_dbm_client(self):
             self.log.warning(error)
             raise
 
-    def cancel(self):
-        """
-        Cancel DBM async jobs and clean up connections.
-        This is called when the check is being shut down.
-        """
-        self.log.debug("Cancelling ClickHouse check and cleaning up connections")
-
-        # Cancel DBM async jobs
-        if self.statement_metrics:
-            self.statement_metrics.cancel()
-        if self.statement_samples:
-            self.statement_samples.cancel()
-        if self.query_completions:
-            self.query_completions.cancel()
-        if self.query_errors:
-            self.query_errors.cancel()
-        if self.table_metrics:
-            self.table_metrics.cancel()
-        if self.metadata:
-            self.metadata.cancel()
-        if self.parts_and_merges:
-            self.parts_and_merges.cancel()
-
-        # Wait for job loops to finish
-        if self.statement_metrics and self.statement_metrics._job_loop_future:
-            self.statement_metrics._job_loop_future.result()
-        if self.statement_samples and self.statement_samples._job_loop_future:
-            self.statement_samples._job_loop_future.result()
-        if self.query_completions and self.query_completions._job_loop_future:
-            self.query_completions._job_loop_future.result()
-        if self.query_errors and self.query_errors._job_loop_future:
-            self.query_errors._job_loop_future.result()
-        if self.table_metrics and self.table_metrics._job_loop_future:
-            self.table_metrics._job_loop_future.result()
-        if self.metadata and self.metadata._job_loop_future:
-            self.metadata._job_loop_future.result()
-        if self.parts_and_merges and self.parts_and_merges._job_loop_future:
-            self.parts_and_merges._job_loop_future.result()
-
-        # Close main client
+    def shutdown(self) -> None:
+        """Close the main client and release the shared connection pool."""
+        self._query_manager = None
+        self.health = None
         if self._client:
             try:
                 self._client.close()
             except Exception as e:
                 self.log.debug("Error closing main client: %s", e)
             self._client = None
 
-        # Clear the shared pool manager
-        # Note: urllib3 pool connections are automatically closed when idle
+        # urllib3 pool connections are closed automatically once idle, so dropping the manager is
+        # enough. The jobs' dedicated clients share it, and they are shut down before this runs.
         self._pool_manager = None
 
     def version_lt(self, version: str) -> bool:

clickhouse/datadog_checks/clickhouse/metadata.py

@@ -40,9 +40,12 @@ def __init__(self, check: ClickhouseCheck):
         self._schema_collector = ClickhouseSchemaCollector(check)
         self._schema_collector._cancel_event = self._cancel_event
 
-    def cancel(self):
-        super(ClickhouseMetadata, self).cancel()
-        self._schema_collector.close()
+    def shutdown(self) -> None:
+        if self._schema_collector is not None:
+            self._schema_collector.close()
+        # The collector holds the check too, so dropping it here releases both.
+        self._schema_collector = None
+        self._check = None
 
     @tracked_method(agent_check_getter=agent_check_getter)
     def run_job(self):

clickhouse/datadog_checks/clickhouse/parts_and_merges.py

@@ -266,9 +266,9 @@ def __init__(self, check: ClickhouseCheck, config: PartsAndMerges):
         }
         self._obfuscate_options = to_native_string(json.dumps(obfuscate_options))
 
-    def cancel(self):
-        super(ClickhousePartsAndMerges, self).cancel()
+    def shutdown(self) -> None:
         self._close_db_client()
+        self._check = None
 
     def _close_db_client(self):
         if self._db_client:
@@ -282,6 +282,8 @@ def _get_debug_tags(self) -> list[str]:
         return list(self._tags_no_db) if self._tags_no_db else []
 
     def _execute_query(self, query: str) -> list:
+        if self._cancel_event.is_set():
+            raise Exception("Job loop cancelled. Aborting query.")
         if self._db_client is None:
             self._db_client = self._check.create_dbm_client()
         try:

clickhouse/datadog_checks/clickhouse/query_completions.py

@@ -148,6 +148,11 @@ def __init__(self, check: ClickhouseCheck, config, flush_config: CollectAsyncIns
         self._flush_checkpoint = NodeCheckpoint(self, FLUSH_CHECKPOINT_CACHE_KEY, self._flush_collection_interval)
         self._last_flush_collection_time = 0.0
 
+    def shutdown(self) -> None:
+        super().shutdown()
+        # Holds the check and a bound method of this job, so dropping it releases both.
+        self._explain_plans = None
+
     @tracked_method(agent_check_getter=agent_check_getter)
     def _collect_and_submit(self):
         """

clickhouse/datadog_checks/clickhouse/query_log_job.py

@@ -377,10 +377,10 @@ def __init__(
         }
         self._obfuscate_options = to_native_string(json.dumps(obfuscate_options))
 
-    def cancel(self):
-        """Cancel the job and clean up the dedicated client."""
-        super().cancel()
+    def shutdown(self) -> None:
+        """Close the dedicated client, once the job loop has stopped."""
         self._close_db_client()
+        self._check = None
 
     def _close_db_client(self):
         """Close the dedicated database client if it exists."""

clickhouse/datadog_checks/clickhouse/statement_samples.py

@@ -193,10 +193,10 @@ def __init__(
         # Set once system.asynchronous_inserts is found to be missing, so we skip collection
         self._buffer_unavailable = False
 
-    def cancel(self):
-        """Cancel the job and clean up the dedicated client."""
-        super(ClickhouseStatementSamples, self).cancel()
+    def shutdown(self) -> None:
+        """Close the dedicated client, once the job loop has stopped."""
         self._close_db_client()
+        self._check = None
 
     def _close_db_client(self):
         """Close the dedicated database client if it exists."""
@@ -223,6 +223,9 @@ def _get_active_queries(self):
         all nodes in the cluster.
         For self-hosted: Queries only the local node's system.processes.
         """
+        if self._cancel_event.is_set():
+            raise Exception("Job loop cancelled. Aborting query.")
+
         start_time = time.time()
 
         try:
@@ -383,6 +386,9 @@ def _get_active_connections(self):
         For ClickHouse Cloud: Uses clusterAllReplicas to aggregate across all nodes.
         For self-hosted: Aggregates only the local node's connections.
         """
+        if self._cancel_event.is_set():
+            raise Exception("Job loop cancelled. Aborting query.")
+
         try:
             start_time = time.time()
 
@@ -534,6 +540,10 @@ def _query_buffer_snapshot(self) -> list[dict]:
             asynchronous_inserts_table=buffer_table,
             max_samples_per_collection=self._buffer_max_samples_per_collection,
         )
+
+        if self._cancel_event.is_set():
+            raise Exception("Job loop cancelled. Aborting query.")
+
         try:
             if self._db_client is None:
                 self._db_client = self._check.create_dbm_client()

clickhouse/datadog_checks/clickhouse/table_metrics.py

@@ -82,9 +82,9 @@ def __init__(self, check: ClickhouseCheck, config: SchemaMetrics):
         self._view_refreshes_permission_logged = False
         self._view_refreshes_skip = False
 
-    def cancel(self):
-        super(ClickhouseTableMetrics, self).cancel()
+    def shutdown(self) -> None:
         self._close_db_client()
+        self._check = None
 
     def _close_db_client(self):
         if self._db_client:
@@ -95,6 +95,8 @@ def _close_db_client(self):
             self._db_client = None
 
     def _execute_query(self, query: str) -> list:
+        if self._cancel_event.is_set():
+            raise Exception("Job loop cancelled. Aborting query.")
         if self._db_client is None:
             self._db_client = self._check.create_dbm_client()
             self._db_client.set_client_setting('max_execution_time', self._collection_interval)
@@ -108,6 +110,10 @@ def _execute_query(self, query: str) -> list:
     @tracked_method(agent_check_getter=agent_check_getter)
     def run_job(self):
         self._emit_table_size_gauges()
+        # _emit_table_size_gauges() swallows the cancellation raised by _execute_query, and the
+        # view refresh collection queries the check's client directly, so it needs its own check.
+        if self._cancel_event.is_set():
+            return
         self._collect_view_refresh_metrics()
 
     def _emit_table_size_gauges(self) -> None:

clickhouse/pyproject.toml

@@ -27,7 +27,7 @@ classifiers = [
     "Private :: Do Not Upload",
 ]
 dependencies = [
-    "datadog-checks-base>=37.42.0",
+    "datadog-checks-base>=38.1.0",
 ]
 dynamic = [
     "version",

clickhouse/tests/test_metadata.py

@@ -399,16 +399,27 @@ def test_collect_all_chunks_share_collection_started_at(check):
     assert len(started_ats) == 1
 
 
-def test_cancel_closes_db_client(check):
+def test_shutdown_closes_db_client(check):
+    collector = check.metadata._schema_collector
     fake_client = mock.MagicMock()
-    check.metadata._schema_collector._db_client = fake_client
+    collector._db_client = fake_client
 
-    check.metadata.cancel()
+    check.metadata.shutdown()
 
-    assert check.metadata._schema_collector._db_client is None
+    assert collector._db_client is None
     fake_client.close.assert_called_once()
 
 
+def test_shutdown_is_idempotent(check):
+    """DBMAsyncJob.shutdown() overrides must tolerate a second call.
+
+    shutdown_async_jobs() does not isolate teardown failures per job, so a job that raises here
+    would stop the jobs after it from releasing their own clients.
+    """
+    check.metadata.shutdown()
+    check.metadata.shutdown()
+
+
 def test_combined_query_dedupes_replicas_before_limit(check):
     _capture_payloads(check)
     with _capture_all_queries(check.metadata._schema_collector) as seen_queries: