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

mysql/changelog.d/24936.fixed

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

mysql/datadog_checks/mysql/activity.py

@@ -247,6 +247,7 @@ def _check_version(self):
     def _collect_activity(self):
         # type: () -> None
         # do not emit any dd.internal metrics for DBM specific check code
+        self._raise_if_cancelled()
         tags = [t for t in self._tags if not t.startswith('dd.internal')]
         with closing(self._get_db_connection().cursor(CommenterDictCursor)) as cursor:
             rows = self._get_activity(cursor)
@@ -289,6 +290,7 @@ def _get_activity_query(self):
     @tracked_method(agent_check_getter=agent_check_getter, track_result_length=True)
     def _get_activity(self, cursor):
         # type: (pymysql.cursor) -> List[Dict[str]]
+        self._raise_if_cancelled()
         query = self._get_activity_query()
         self._log.debug("Running activity query [%s]", query)
         cursor.execute(query)
@@ -432,6 +434,11 @@ def _json_event_encoding(o):
             return int(o.total_seconds())
         raise TypeError
 
+    def shutdown(self) -> None:
+        self._close_db_conn()
+        self._check = None
+        self._connection_args_provider = None
+
     def _close_db_conn(self):
         # type: () -> None
         if self._db:

mysql/datadog_checks/mysql/databases_data.py

@@ -144,11 +144,11 @@ def __init__(self, mysql_metadata, check, config):
             config.schemas_config.get('max_execution_time', self.DEFAULT_MAX_EXECUTION_TIME), collection_interval
         )
 
-    def shut_down(self):
-        self._data_submitter.submit()
-
     def _cursor_run(self, cursor, query, params=None):
         """Run the query, log it, and emit a metric on database error."""
+        cancel_event = getattr(self._metadata, '_cancel_event', None)
+        if cancel_event is not None and cancel_event.is_set():
+            raise Exception("Job loop cancelled. Aborting query.")
         try:
             params_repr = "({} params)".format(len(params)) if isinstance(params, list) else params
             self._log.debug("Running query [{}] params={}".format(query, params_repr))
@@ -279,7 +279,7 @@ def _fetch_for_databases(self, db_infos, cursor):
                     )
                 )
                 return
-            except Exception as e:
+            except pymysql.DatabaseError as e:
                 self._log.error(
                     "While executing fetch database data for database {}, the following exception occurred {}".format(
                         db_info['name'], e

mysql/datadog_checks/mysql/metadata.py

@@ -110,6 +110,12 @@ def get_db_connection(self):
             self._db.ping()
         return self._db
 
+    def shutdown(self) -> None:
+        self._close_db_conn()
+        self._check = None
+        self._connection_args_provider = None
+        self._databases_data = None
+
     def _close_db_conn(self):
         if self._db:
             try:
@@ -123,6 +129,7 @@ def _cursor_run(self, cursor, query, params=None):
         """
         Run and log the query. If provided, obfuscated params are logged in place of the regular params.
         """
+        self._raise_if_cancelled()
         try:
             self._log.debug("Running query [{}] params={}".format(query, params))
             cursor.execute(query, params)
@@ -158,9 +165,6 @@ def run_job(self):
                                 These may be unavailable until the error is resolved. The error - {}""".format(e)
                 )
 
-    def shut_down(self):
-        self._databases_data.shut_down()
-
     @tracked_method(agent_check_getter=attrgetter('_check'))
     def report_mysql_metadata(self):
         settings = []

mysql/datadog_checks/mysql/mysql.py

@@ -149,15 +149,11 @@ def __init__(self, name, init_config, instances):
             and self.cloud_metadata['aws']['managed_authentication'].get('enabled', False)
         )
 
-        # Pass function reference and managed auth flag to async jobs
-        self._statement_metrics = MySQLStatementMetrics(
-            self, self._config, self._get_connection_args, self._uses_aws_managed_auth
-        )
-        self._statement_samples = MySQLStatementSamples(
-            self, self._config, self._get_connection_args, self._uses_aws_managed_auth
-        )
-        self._mysql_metadata = MySQLMetadata(self, self._config, self._get_connection_args, self._uses_aws_managed_auth)
-        self._query_activity = MySQLActivity(self, self._config, self._get_connection_args, self._uses_aws_managed_auth)
+        self.statement_metrics = None
+        self.statement_samples = None
+        self.mysql_metadata = None
+        self.query_activity = None
+        self._register_async_jobs()
         self._index_metrics = MySqlIndexMetrics(self._config)
         # _database_instance_emitted: limit the collection and transmission of the database instance metadata
         self._database_instance_emitted = TTLCache(
@@ -171,6 +167,30 @@ def __init__(self, name, init_config, instances):
 
         self._submit_initialization_health_event()
 
+    def shutdown(self) -> None:
+        """Release the resources this check holds for its whole lifetime."""
+        self._query_manager = None
+        self._runtime_queries_cached = None
+        self.health = None
+
+    def _register_async_jobs(self):
+        """Build and register the async jobs enabled by this check's configuration."""
+        if not self._config.dbm_enabled:
+            return
+
+        self.statement_metrics = self.register_async_job(
+            MySQLStatementMetrics(self, self._config, self._get_connection_args, self._uses_aws_managed_auth)
+        )
+        self.statement_samples = self.register_async_job(
+            MySQLStatementSamples(self, self._config, self._get_connection_args, self._uses_aws_managed_auth)
+        )
+        self.mysql_metadata = self.register_async_job(
+            MySQLMetadata(self, self._config, self._get_connection_args, self._uses_aws_managed_auth)
+        )
+        self.query_activity = self.register_async_job(
+            MySQLActivity(self, self._config, self._get_connection_args, self._uses_aws_managed_auth)
+        )
+
     def _submit_initialization_health_event(self):
         try:
             # Handle the config validation result after we've set tags so those tags are included in the health event
@@ -398,10 +418,7 @@ def check(self, _):
 
                 if self._config.dbm_enabled:
                     dbm_tags = list(set(self.service_check_tags) | set(tags))
-                    self._statement_metrics.run_job_loop(dbm_tags)
-                    self._statement_samples.run_job_loop(dbm_tags)
-                    self._query_activity.run_job_loop(dbm_tags)
-                    self._mysql_metadata.run_job_loop(dbm_tags)
+                    self.run_async_jobs(dbm_tags)
 
                 # keeping track of these:
                 self._put_qcache_stats()
@@ -416,12 +433,6 @@ def check(self, _):
             self._conn = None
             self._report_warnings()
 
-    def cancel(self):
-        self._statement_samples.cancel()
-        self._statement_metrics.cancel()
-        self._query_activity.cancel()
-        self._mysql_metadata.cancel()
-
     def _new_query_executor(self, queries):
         return QueryExecutor(
             self.execute_query_raw,

mysql/datadog_checks/mysql/statement_samples.py

@@ -290,6 +290,11 @@ def _read_version_info(self):
                 self._global_status_table = "performance_schema.global_status"
             self._version_processed = True
 
+    def shutdown(self) -> None:
+        self._close_db_conn()
+        self._check = None
+        self._connection_args_provider = None
+
     def _close_db_conn(self):
         if self._db:
             try:
@@ -330,6 +335,7 @@ def _cursor_run(self, cursor, query, params=None, obfuscated_params=None, obfusc
         """
         Run and log the query. If provided, obfuscated params are logged in place of the regular params.
         """
+        self._raise_if_cancelled()
         try:
             logged_query = obfuscated_query if obfuscated_query else query
             self._log.debug("Running query [%s] %s", logged_query, obfuscated_params if obfuscated_params else params)

mysql/datadog_checks/mysql/statements.py

@@ -147,6 +147,11 @@ def __init__(self, check, config, connection_args_provider, uses_managed_auth=Fa
             ttl=self._config.statement_rows_cache_ttl,
         )
 
+    def shutdown(self) -> None:
+        self._close_db_conn()
+        self._check = None
+        self._connection_args_provider = None
+
     def _close_db_conn(self):
         if self._db:
             try:
@@ -236,6 +241,7 @@ def _collect_per_statement_metrics(self, tags):
         return rows
 
     def _get_statement_count(self, tags):
+        self._raise_if_cancelled()
         with closing(self._get_db_connection().cursor(CommenterDictCursor)) as cursor:
             cursor.execute("SELECT count(*) AS count from performance_schema.events_statements_summary_by_digest")
 
@@ -343,6 +349,7 @@ def _query_summary_per_statement(self):
                 LIMIT 10000
                 """
 
+        self._raise_if_cancelled()
         with closing(self._get_db_connection().cursor(CommenterDictCursor)) as cursor:
             args = [self._last_seen] if only_query_recent_statements else None
             cursor.execute(sql_statement_summary, args)

mysql/datadog_checks/mysql/util.py

@@ -100,6 +100,7 @@ class ManagedAuthConnectionMixin:
         self._uses_managed_auth (bool)
         self._db_created_at (float, timestamp)
         self._db (connection or None)
+        self._cancel_event (event.Event, used to abort queries if the Agent has unscheduled this check)
 
     Subclasses must implement:
         _close_db_conn() - closes self._db
@@ -113,6 +114,11 @@ def _should_reconnect_for_managed_auth(self):
             return False
         return (time.time() - self._db_created_at) >= self.MANAGED_AUTH_RECONNECT_INTERVAL
 
+    def _raise_if_cancelled(self):
+        """Abort before a query if the Agent has unscheduled this check."""
+        if self._cancel_event.is_set():
+            raise Exception("Job loop cancelled. Aborting query.")
+
     def _get_db_connection(self):
         """Get or create database connection, reconnecting periodically for managed auth."""
         if self._should_reconnect_for_managed_auth():

mysql/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",

mysql/tests/test_metadata.py

@@ -3,11 +3,13 @@
 # Licensed under a 3-clause BSD style license (see LICENSE)
 
 import re
+from unittest import mock
 
 import pytest
 from packaging.version import parse as parse_version
 
 from datadog_checks.mysql import MySql
+from datadog_checks.mysql.databases_data import DatabasesData
 
 from . import common
 from .common import MYSQL_FLAVOR, MYSQL_REPLICATION, MYSQL_VERSION_PARSED
@@ -62,6 +64,18 @@ def normalize_values(actual_payload):
                             )
 
 
+@pytest.mark.unit
+def test_schema_collection_aborts_query_when_cancelled(dbm_instance):
+    check = MySql(common.CHECK_NAME, {}, instances=[dbm_instance])
+    check.mysql_metadata.cancel()
+    databases_data = DatabasesData(check.mysql_metadata, check, check._config)
+    cursor = mock.MagicMock()
+
+    with pytest.raises(Exception, match='cancelled'):
+        databases_data._fetch_for_databases([{'name': 'db1'}, {'name': 'db2'}], cursor)
+    cursor.execute.assert_not_called()
+
+
 @pytest.mark.integration
 @pytest.mark.usefixtures('dd_environment')
 def test_collect_mysql_settings(aggregator, dbm_instance, dd_run_check):
@@ -77,32 +91,33 @@ def test_collect_mysql_settings(aggregator, dbm_instance, dd_run_check):
     assert len(event["metadata"]) > 0
 
 
-@pytest.mark.integration
-@pytest.mark.usefixtures('dd_environment')
-def test_metadata_collection_interval_and_enabled(dbm_instance):
-    dbm_instance['schemas_collection'] = {"enabled": True, "collection_interval": 101}
-    dbm_instance['collect_settings'] = {"enabled": False, "collection_interval": 100}
-
-    mysql_check = MySql(common.CHECK_NAME, {}, instances=[dbm_instance])
-    assert mysql_check._mysql_metadata.enabled
-    assert mysql_check._mysql_metadata.collection_interval == 101
-    dbm_instance['schemas_collection'] = {"enabled": False, "collection_interval": 101}
-    dbm_instance['collect_settings'] = {"enabled": True, "collection_interval": 102}
-
-    mysql_check = MySql(common.CHECK_NAME, {}, instances=[dbm_instance])
-    assert mysql_check._mysql_metadata.enabled
-    assert mysql_check._mysql_metadata.collection_interval == 102
-
-    dbm_instance['schemas_collection'] = {"enabled": True, "collection_interval": 101}
-    dbm_instance['collect_settings'] = {"enabled": True, "collection_interval": 102}
-
-    mysql_check = MySql(common.CHECK_NAME, {}, instances=[dbm_instance])
-    assert mysql_check._mysql_metadata.enabled
-    assert mysql_check._mysql_metadata.collection_interval == 101
-    dbm_instance['schemas_collection'] = {"enabled": False}
-    dbm_instance['collect_settings'] = {"enabled": False}
-    mysql_check = MySql(common.CHECK_NAME, {}, instances=[dbm_instance])
-    assert not mysql_check._mysql_metadata.enabled
+@pytest.mark.unit
+@pytest.mark.parametrize(
+    'schemas_collection, collect_settings, expected_enabled, expected_interval',
+    [
+        ({"enabled": True, "collection_interval": 101}, {"enabled": False}, True, 101),
+        ({"enabled": False}, {"enabled": True, "collection_interval": 102}, True, 102),
+        (
+            {"enabled": True, "collection_interval": 101},
+            {"enabled": True, "collection_interval": 102},
+            True,
+            101,
+        ),
+        ({"enabled": False}, {"enabled": False}, False, None),
+    ],
+    ids=['schemas-only', 'settings-only', 'both', 'neither'],
+)
+def test_metadata_collection_interval_and_enabled(
+    dbm_instance, schemas_collection, collect_settings, expected_enabled, expected_interval
+):
+    dbm_instance['schemas_collection'] = schemas_collection
+    dbm_instance['collect_settings'] = collect_settings
+
+    metadata = MySql(common.CHECK_NAME, {}, instances=[dbm_instance]).mysql_metadata
+
+    assert metadata.enabled is expected_enabled
+    if expected_interval is not None:
+        assert metadata.collection_interval == expected_interval
 
 
 @pytest.mark.integration