[DBMON-6812] Add async job registry to DatabaseCheck (DataDog#24442)

* Add job registration to DatabaseCheck

* Add changelog

* Update variable to avoid collision with Postgres existing usage

* Add type hints to new public functions

* Remove None handling

Author: eric-weaver

datadog_checks_base/changelog.d/24442.added

@@ -0,0 +1 @@
+Add an async job registry to `DatabaseCheck` so DBM integrations can register `DBMAsyncJob`s additively and run, cancel, or shut them all down through a single entry point (`register_async_job`, `run_async_jobs`, `cancel_async_jobs`, `shutdown_async_jobs`), along with a `DBMAsyncJob.shutdown()` hook for releasing lifetime-scoped resources on unschedule.

datadog_checks_base/datadog_checks/base/checks/db.py

@@ -4,14 +4,22 @@
 
 from abc import abstractmethod
 from string import Template
+from typing import TYPE_CHECKING, Dict, List
 
 from datadog_checks.base.agent import datadog_agent
 from datadog_checks.base.utils.db.utils import TagManager
 
 from . import AgentCheck
 
+if TYPE_CHECKING:
+    from datadog_checks.base.utils.db.utils import DBMAsyncJob
+
 
 class DatabaseCheck(AgentCheck):
+    """
+    Base class for Database Monitoring (DBM) integrations.
+    """
+
     #: Authoritative DBM platform identifier for this integration.
     #: Subclasses should set this explicitly; it is the value surfaced by
     #: :attr:`dbms` and used across DBM payloads, metric name prefixes and async jobs.
@@ -23,6 +31,49 @@ def __init__(self, *args, **kwargs):
         self._database_identifier = None
         self._dbms_fallback_warning_logged = False
         self.tag_manager = TagManager()
+        #: Async jobs owned by this check, keyed by job name, populated via
+        #: :meth:`register_async_job`.
+        self._async_job_registry: Dict[str, "DBMAsyncJob"] = {}
+
+    def register_async_job(self, job: "DBMAsyncJob") -> "DBMAsyncJob":
+        """
+        Register ``job`` under its ``job_name`` so the check manages its lifecycle, and return it
+        unchanged.
+
+        Registering a job whose name matches an already-registered job replaces it. Raises
+        ``ValueError`` if the job has no name.
+        """
+        if job.job_name is None:
+            raise ValueError("Cannot register an async job without a job_name")
+        self._async_job_registry[job.job_name] = job
+        return job
+
+    def run_async_jobs(self, tags: List[str]) -> None:
+        """Run each registered job's loop, forwarding ``tags`` to every job."""
+        for job in self._async_job_registry.values():
+            job.run_job_loop(tags)
+
+    def cancel_async_jobs(self) -> None:
+        """
+        Signal every registered job to stop, without waiting for loops to finish or releasing
+        resources.
+
+        Safe to call while ``check()`` is running. Follow with :meth:`shutdown_async_jobs` to wait
+        for the loops and release resources.
+        """
+        for job in self._async_job_registry.values():
+            job.cancel()
+
+    def shutdown_async_jobs(self) -> None:
+        """
+        Wait for every registered job's loop to finish (:meth:`~DBMAsyncJob.wait_for_completion`)
+        and run its teardown (:meth:`~DBMAsyncJob.shutdown`).
+
+        Must not run concurrently with ``check()``.
+        """
+        for job in self._async_job_registry.values():
+            job.wait_for_completion()
+            job.shutdown()
 
     def database_monitoring_query_sample(self, raw_event: str):
         self.event_platform_event(raw_event, "dbm-samples")

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

@@ -345,12 +345,26 @@ def __init__(
         if self._features is None:
             self._features = [None]
 
+    @property
+    def job_name(self) -> Optional[str]:
+        """The job's name"""
+        return self._job_name
+
     def cancel(self):
         """
         Send a signal to cancel the job loop asynchronously.
         """
         self._cancel_event.set()
 
+    def wait_for_completion(self) -> None:
+        """
+        Block until the job loop has finished running, then clear its future. No-op if the loop is
+        not running. Typically called after :meth:`cancel` to wait for the loop to stop.
+        """
+        if self._job_loop_future:
+            self._job_loop_future.result()
+            self._job_loop_future = None
+
     def run_job_loop(self, tags):
         """
         :param tags:
@@ -467,6 +481,8 @@ def _job_loop(self):
                         )
         finally:
             self._log.info("[%s] Shutting down job loop", self._job_tags_str)
+            # Runs on every loop exit, including the inactivity stop above, after which the loop may
+            # restart on the next check run. For one-time teardown on unschedule, override shutdown().
             if self._shutdown_callback:
                 self._shutdown_callback()
 
@@ -498,6 +514,19 @@ def _run_job_traced(self):
     def run_job(self):
         raise NotImplementedError()
 
+    def shutdown(self) -> None:
+        """
+        Release resources the job holds for its whole lifetime, such as dedicated DB connections or
+        clients.
+
+        Called once when the owning check is unscheduled, after the loop has stopped. The default
+        is a no-op; override to close long-lived resources, and keep the implementation idempotent.
+
+        Unlike ``shutdown_callback``, which runs on every loop exit and may be followed by a
+        restart, this runs only during final teardown.
+        """
+        pass
+
 
 @contextlib.contextmanager
 def tracked_query(check, operation, tags=None):

datadog_checks_base/tests/base/checks/test_database_check.py

@@ -1,12 +1,14 @@
 # (C) Datadog, Inc. 2026-present
 # All rights reserved
 # Licensed under a 3-clause BSD style license (see LICENSE)
+from concurrent.futures.thread import ThreadPoolExecutor
 from unittest import mock
 
 import pytest
 
 from datadog_checks.base.checks.db import DatabaseCheck
 from datadog_checks.base.stubs.datadog_agent import datadog_agent
+from datadog_checks.base.utils.db.utils import DBMAsyncJob
 
 
 class FakeDatabaseCheck(DatabaseCheck):
@@ -23,6 +25,43 @@ def cloud_metadata(self):
         return {}
 
 
+class RegistryTestJob(DBMAsyncJob):
+    """Minimal DBMAsyncJob used to exercise the DatabaseCheck async job registry."""
+
+    def __init__(self, check, enabled=True, job_name="test-job"):
+        super().__init__(
+            check,
+            enabled=enabled,
+            dbms="test-dbms",
+            rate_limit=10,
+            max_sleep_chunk_s=0.1,
+            job_name=job_name,
+        )
+        self.shutdown_calls = 0
+
+    def shutdown(self):
+        self.shutdown_calls += 1
+
+    def run_job(self):
+        pass
+
+
+@pytest.fixture
+def registry_check():
+    check = FakeDatabaseCheck("test", {}, [{}])
+    yield check
+    # Stop any registered jobs so their loops don't outlive the test.
+    check.cancel_async_jobs()
+    check.shutdown_async_jobs()
+
+
+@pytest.fixture(autouse=True)
+def stop_orphaned_threads():
+    # Recreate the shared executor per test so job loops don't leak across tests.
+    DBMAsyncJob.executor.shutdown(wait=True)
+    DBMAsyncJob.executor = ThreadPoolExecutor()
+
+
 def test_agent_hostname_resolves_once_and_caches():
     check = FakeDatabaseCheck("test", {}, [{}])
     # The hostname comes from an FFI call, so it should only be resolved once and cached.
@@ -89,3 +128,64 @@ def database_identifier_params(self):
     if tags_after is not None:
         check.tag_manager.set_tags_from_list(tags_after, replace=True)
         assert check.database_identifier == expected
+
+
+@pytest.mark.parametrize("register_twice", [False, True], ids=["single", "duplicate_instance"])
+def test_register_async_job_adds_and_dedupes(registry_check, register_twice):
+    job = RegistryTestJob(registry_check)
+    assert registry_check.register_async_job(job) is job
+    if register_twice:
+        assert registry_check.register_async_job(job) is job
+    assert registry_check._async_job_registry == {"test-job": job}
+
+
+def test_register_async_job_replaces_job_with_same_name(registry_check):
+    first = RegistryTestJob(registry_check, job_name="query-metrics")
+    second = RegistryTestJob(registry_check, job_name="query-metrics")
+
+    registry_check.register_async_job(first)
+    registry_check.register_async_job(second)
+
+    assert registry_check._async_job_registry == {"query-metrics": second}
+
+
+def test_register_async_job_requires_job_name(registry_check):
+    with pytest.raises(ValueError):
+        registry_check.register_async_job(RegistryTestJob(registry_check, job_name=None))
+
+
+@pytest.mark.parametrize("enabled", [True, False], ids=["enabled", "disabled"])
+def test_run_async_jobs_starts_only_enabled_jobs(registry_check, enabled):
+    job = registry_check.register_async_job(RegistryTestJob(registry_check, enabled=enabled))
+
+    registry_check.run_async_jobs([])
+
+    # Only enabled jobs get a running loop; disabled jobs are skipped by run_job_loop.
+    assert (job._job_loop_future is not None) == enabled
+
+
+def test_cancel_async_jobs_signals_without_touching_futures(registry_check):
+    job = registry_check.register_async_job(RegistryTestJob(registry_check))
+    registry_check.run_async_jobs([])
+    assert job._job_loop_future is not None
+
+    registry_check.cancel_async_jobs()
+
+    # cancel_async_jobs only sets the cancel event; the future stays in place for shutdown to await.
+    assert job._cancel_event.is_set()
+    assert job._job_loop_future is not None
+
+
+@pytest.mark.parametrize("started", [True, False], ids=["loop_started", "loop_not_started"])
+def test_shutdown_async_jobs_tears_down_and_calls_shutdown(registry_check, started):
+    job = registry_check.register_async_job(RegistryTestJob(registry_check))
+    if started:
+        registry_check.run_async_jobs([])
+        assert job._job_loop_future is not None
+        registry_check.cancel_async_jobs()
+
+    registry_check.shutdown_async_jobs()
+
+    # The future is cleared and shutdown runs once, whether or not a loop was started.
+    assert job._job_loop_future is None
+    assert job.shutdown_calls == 1