diff --git a/mysql/changelog.d/24936.fixed b/mysql/changelog.d/24936.fixed new file mode 100644 index 0000000000000..ac3bcee0c5c12 --- /dev/null +++ b/mysql/changelog.d/24936.fixed @@ -0,0 +1 @@ +Manage the DBM async jobs through the ``DatabaseCheck`` registry. diff --git a/mysql/datadog_checks/mysql/activity.py b/mysql/datadog_checks/mysql/activity.py index 1eeed7a3faa31..b076a47bc5f23 100644 --- a/mysql/datadog_checks/mysql/activity.py +++ b/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: diff --git a/mysql/datadog_checks/mysql/databases_data.py b/mysql/datadog_checks/mysql/databases_data.py index c26f7a0c988d8..47bca9907af40 100644 --- a/mysql/datadog_checks/mysql/databases_data.py +++ b/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 diff --git a/mysql/datadog_checks/mysql/metadata.py b/mysql/datadog_checks/mysql/metadata.py index 3cabbad3bf609..d653b9cfa086e 100644 --- a/mysql/datadog_checks/mysql/metadata.py +++ b/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 = [] diff --git a/mysql/datadog_checks/mysql/mysql.py b/mysql/datadog_checks/mysql/mysql.py index 16af130515f57..b53b3fc6edca9 100644 --- a/mysql/datadog_checks/mysql/mysql.py +++ b/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, diff --git a/mysql/datadog_checks/mysql/statement_samples.py b/mysql/datadog_checks/mysql/statement_samples.py index 3efcf5dbdc5e5..8e339b731304f 100644 --- a/mysql/datadog_checks/mysql/statement_samples.py +++ b/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) diff --git a/mysql/datadog_checks/mysql/statements.py b/mysql/datadog_checks/mysql/statements.py index 06506129e7dbd..a42341c4c03b6 100644 --- a/mysql/datadog_checks/mysql/statements.py +++ b/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) diff --git a/mysql/datadog_checks/mysql/util.py b/mysql/datadog_checks/mysql/util.py index f278934945ffc..0d767ddb236fe 100644 --- a/mysql/datadog_checks/mysql/util.py +++ b/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(): diff --git a/mysql/pyproject.toml b/mysql/pyproject.toml index 55eb750dcc390..4bb27f16a16f3 100644 --- a/mysql/pyproject.toml +++ b/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", diff --git a/mysql/tests/test_metadata.py b/mysql/tests/test_metadata.py index fc7014f2418e0..c368f8616c141 100644 --- a/mysql/tests/test_metadata.py +++ b/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 diff --git a/mysql/tests/test_query_activity.py b/mysql/tests/test_query_activity.py index 6ae8ecf010509..f3e8d4905673a 100644 --- a/mysql/tests/test_query_activity.py +++ b/mysql/tests/test_query_activity.py @@ -164,7 +164,7 @@ def _run_blocking(conn): assert blocked_row['event_timer_end'], "missing event timer end" assert blocked_row['query_truncated'] == expected_query_truncated - if check._query_activity._should_collect_blocking_queries(): + if check.query_activity._should_collect_blocking_queries(): assert len(activity['mysql_activity']) >= 2, "should have collected at least two activity payloads" captured_idle_blocker = False for activity in dbm_activity: @@ -278,10 +278,10 @@ def _run_query_with_mock_obfuscator(conn, _query): def test_get_estimated_row_size_bytes(dbm_instance, file): check = MySql(CHECK_NAME, {}, [dbm_instance]) test_activity = _load_test_activity_json(file) - actual_size = len(json.dumps(test_activity, default=check._query_activity._json_event_encoding)) + actual_size = len(json.dumps(test_activity, default=check.query_activity._json_event_encoding)) computed_size = 0 for a in test_activity: - computed_size += check._query_activity._get_estimated_row_size_bytes(a) + computed_size += check.query_activity._get_estimated_row_size_bytes(a) assert abs((actual_size - computed_size) / float(actual_size)) <= 0.10 @@ -323,7 +323,7 @@ def _create_time_in_picoseconds(date_obj): ) def test_sort_key(dbm_instance, rows, expected_rows): check = MySql(CHECK_NAME, {}, [dbm_instance]) - output = sorted(rows, key=lambda r: check._query_activity._sort_key(r)) + output = sorted(rows, key=lambda r: check.query_activity._sort_key(r)) assert output == expected_rows @@ -386,7 +386,7 @@ def test_truncate_on_max_size_bytes(dbm_instance, datadog_agent, rows, expected_ check = MySql(CHECK_NAME, {}, [dbm_instance]) with mock.patch.object(datadog_agent, 'obfuscate_sql', passthrough=True) as mock_agent: mock_agent.side_effect = "something" - result_rows = check._query_activity._normalize_rows(rows) + result_rows = check.query_activity._normalize_rows(rows) assert len(result_rows) == expected_len for index, user in enumerate(expected_users): assert result_rows[index]['processlist_user'] == user @@ -441,7 +441,7 @@ def test_normalize_rows_with_null_event_timers(dbm_instance, datadog_agent, rows check = MySql(CHECK_NAME, {}, [dbm_instance]) with mock.patch.object(datadog_agent, 'obfuscate_sql', passthrough=True) as mock_agent: mock_agent.side_effect = "something" - result_rows = check._query_activity._normalize_rows(rows) + result_rows = check.query_activity._normalize_rows(rows) assert [row['processlist_user'] for row in result_rows] == expected_users @@ -590,57 +590,7 @@ def test_activity_collection_rate_limit(aggregator, dd_run_check, dbm_instance): f"expected at most ~{max_expected} (with 2x tolerance: {max_expected * 2})" ) - # Verify the rate limiter is configured with the expected rate - expected_rate = 1.0 / collection_interval - assert check._query_activity._rate_limiter.rate_limit_s == expected_rate - assert check._query_activity.collection_interval == collection_interval - - -@pytest.mark.integration -@pytest.mark.usefixtures('dd_environment') -@pytest.mark.parametrize("activity_enabled", [True, False]) -def test_async_job_enabled(dd_run_check, dbm_instance, activity_enabled): - dbm_instance['query_activity'] = {'enabled': activity_enabled, 'run_sync': False} - check = MySql(CHECK_NAME, {}, [dbm_instance]) - dd_run_check(check) - check.cancel() - if activity_enabled: - assert check._query_activity._job_loop_future is not None - check._query_activity._job_loop_future.result() - else: - assert check._query_activity._job_loop_future is None - - -@pytest.mark.integration -@pytest.mark.usefixtures('dd_environment') -def test_async_job_inactive_stop(aggregator, dd_run_check, dbm_instance): - dbm_instance['query_activity']['run_sync'] = False - check = MySql(CHECK_NAME, {}, [dbm_instance]) - dd_run_check(check) - check._query_activity._job_loop_future.result() - aggregator.assert_metric( - "dd.mysql.async_job.inactive_stop", - tags=_expected_dbm_job_err_tags(dbm_instance, check), - hostname='', - ) - - -@pytest.mark.integration -@pytest.mark.usefixtures('dd_environment') -def test_async_job_cancel(aggregator, dd_run_check, dbm_instance): - dbm_instance['query_activity']['run_sync'] = False - check = MySql(CHECK_NAME, {}, [dbm_instance]) - dd_run_check(check) - check.cancel() - # wait for it to stop and make sure it doesn't throw any exceptions - check._query_activity._job_loop_future.result() - assert not check._query_activity._job_loop_future.running(), "activity thread should be stopped" - # if the thread doesn't start until after the cancel signal is set then the db connection will never - # be created in the first place - aggregator.assert_metric( - "dd.mysql.async_job.cancel", - tags=_expected_dbm_job_err_tags(dbm_instance, check), - ) + assert check.query_activity.collection_interval == collection_interval @pytest.mark.integration @@ -659,7 +609,7 @@ def test_events_wait_current_disabled(dbm_instance, dd_run_check, root_conn, agg dd_run_check(check) # force query activity to run once, expect it to exit immediately with a warning - check._query_activity.run_job() + check.query_activity.run_job() dbm_activity = aggregator.get_event_platform_events("dbm-activity") assert check.events_wait_current_enabled is False assert check.warnings == [ @@ -681,7 +631,7 @@ def test_events_wait_current_disabled(dbm_instance, dd_run_check, root_conn, agg dd_run_check(check) check.warnings.clear() assert check.events_wait_current_enabled is True - check._query_activity.run_job() + check.query_activity.run_job() check.cancel() dbm_activity = aggregator.get_event_platform_events("dbm-activity") assert check.warnings == [] @@ -721,7 +671,7 @@ def test_events_wait_current_disabled_no_warning_azure_flexible_server( dd_run_check(check) # force query activity to run once, expect it to collect nothing - check._query_activity.run_job() + check.query_activity.run_job() dbm_activity = aggregator.get_event_platform_events("dbm-activity") assert check.events_wait_current_enabled is False @@ -729,24 +679,6 @@ def test_events_wait_current_disabled_no_warning_azure_flexible_server( assert not dbm_activity, "should not have collected any activity" -# the inactive job metrics are emitted from the main integrations -# directly to metrics-intake, so they should also be properly tagged with a resource -def _expected_dbm_job_err_tags(dbm_instance, check): - _tags = dbm_instance['tags'] + ( - 'database_hostname:stubbed.hostname', - 'database_instance:stubbed.hostname', - 'job:query-activity', - 'port:{}'.format(PORT), - 'dd.internal.resource:database_instance:stubbed.hostname', - 'dbms_flavor:{}'.format(MYSQL_FLAVOR.lower()), - ) - if MYSQL_FLAVOR.lower() in ('mysql', 'percona'): - _tags += ("server_uuid:{}".format(check.server_uuid),) - if MYSQL_REPLICATION == 'classic': - _tags += ('cluster_uuid:{}'.format(check.cluster_uuid), 'replication_role:primary') - return _tags - - @pytest.mark.integration @pytest.mark.usefixtures('dd_environment') def test_if_deadlock_metric_is_collected(aggregator, dd_run_check, dbm_instance): diff --git a/mysql/tests/test_statements.py b/mysql/tests/test_statements.py index d6bae7062a19a..271d9ae5ac83f 100644 --- a/mysql/tests/test_statements.py +++ b/mysql/tests/test_statements.py @@ -84,7 +84,7 @@ def test_statement_samples_enabled_config(dbm_instance, statement_samples_key, s dbm_instance.pop(k, None) dbm_instance[statement_samples_key] = {'enabled': statement_samples_enabled} mysql_check = MySql(common.CHECK_NAME, {}, instances=[dbm_instance]) - assert mysql_check._statement_samples._enabled == statement_samples_enabled + assert mysql_check.statement_samples._enabled == statement_samples_enabled @pytest.mark.integration @@ -348,13 +348,13 @@ def obfuscate_sql(query, options=None): with ( mock.patch.object(datadog_agent, 'obfuscate_sql', passthrough=True) as mock_agent, - mock.patch.object(mysql_check._statement_metrics, '_get_statement_count'), - mock.patch.object(mysql_check._statement_metrics, '_query_summary_per_statement', side_effect=rows_by_run), + mock.patch.object(mysql_check.statement_metrics, '_get_statement_count'), + mock.patch.object(mysql_check.statement_metrics, '_query_summary_per_statement', side_effect=rows_by_run), ): mock_agent.side_effect = obfuscate_sql - assert mysql_check._statement_metrics._collect_per_statement_metrics([]) == [] - rows = mysql_check._statement_metrics._collect_per_statement_metrics([]) + assert mysql_check.statement_metrics._collect_per_statement_metrics([]) == [] + rows = mysql_check.statement_metrics._collect_per_statement_metrics([]) assert len(rows) == 1 assert rows[0]['query_signature'] == query_signature @@ -397,13 +397,13 @@ def obfuscate_sql(query, options=None): with ( mock.patch.object(datadog_agent, 'obfuscate_sql', passthrough=True) as mock_agent, - mock.patch.object(mysql_check._statement_metrics, '_get_statement_count'), - mock.patch.object(mysql_check._statement_metrics, '_query_summary_per_statement', side_effect=rows_by_run), + mock.patch.object(mysql_check.statement_metrics, '_get_statement_count'), + mock.patch.object(mysql_check.statement_metrics, '_query_summary_per_statement', side_effect=rows_by_run), ): mock_agent.side_effect = obfuscate_sql - assert mysql_check._statement_metrics._collect_per_statement_metrics([]) == [] - rows = mysql_check._statement_metrics._collect_per_statement_metrics([]) + assert mysql_check.statement_metrics._collect_per_statement_metrics([]) == [] + rows = mysql_check.statement_metrics._collect_per_statement_metrics([]) assert len(rows) == 1 assert rows[0]['query_signature'] == query_signature @@ -444,13 +444,13 @@ def obfuscate_sql(query, options=None): with ( mock.patch.object(datadog_agent, 'obfuscate_sql', passthrough=True) as mock_agent, - mock.patch.object(mysql_check._statement_metrics, '_get_statement_count'), - mock.patch.object(mysql_check._statement_metrics, '_query_summary_per_statement', side_effect=rows_by_run), + mock.patch.object(mysql_check.statement_metrics, '_get_statement_count'), + mock.patch.object(mysql_check.statement_metrics, '_query_summary_per_statement', side_effect=rows_by_run), ): mock_agent.side_effect = obfuscate_sql - assert mysql_check._statement_metrics._collect_per_statement_metrics([]) == [] - rows = mysql_check._statement_metrics._collect_per_statement_metrics([]) + assert mysql_check.statement_metrics._collect_per_statement_metrics([]) == [] + rows = mysql_check.statement_metrics._collect_per_statement_metrics([]) # the reused id is a different statement => baselined, not merged with the stale row assert rows == [] @@ -643,7 +643,7 @@ def test_statement_samples_collect( mysql_check = MySql(common.CHECK_NAME, {}, [dbm_instance]) if explain_strategy: - mysql_check._statement_samples._preferred_explain_strategies = [explain_strategy] + mysql_check.statement_samples._preferred_explain_strategies = [explain_strategy] with ( mock.patch.object( @@ -662,7 +662,7 @@ def test_statement_samples_collect( logger.debug("running first check") dd_run_check(mysql_check) aggregator.reset() - mysql_check._statement_samples._init_caches() + mysql_check.statement_samples._init_caches() # we deliberately want to keep the connection open for the duration of the test to ensure # the query remains in the events_statements_current and events_statements_history tables @@ -729,10 +729,6 @@ def test_statement_samples_collect( assert event['timestamp'] is not None assert time.time() - event['timestamp'] < 60 # ensure the timestamp is recent - # we avoid closing these in a try/finally block in order to maintain the connections in case we want to - # debug the test with --pdb - mysql_check._statement_samples._close_db_conn() - @pytest.mark.parametrize( "statement,schema,expected_warnings", @@ -770,9 +766,9 @@ def test_missing_explain_procedure(dbm_instance, dd_run_check, aggregator, state # explain plans dbm_instance['query_samples']['enabled'] = False mysql_check = MySql(common.CHECK_NAME, {}, [dbm_instance]) - mysql_check._statement_samples._preferred_explain_strategies = ['PROCEDURE'] - mysql_check._statement_samples._tags = [] - mysql_check._statement_samples._tags_str = '' + mysql_check.statement_samples._preferred_explain_strategies = ['PROCEDURE'] + mysql_check.statement_samples._tags = [] + mysql_check.statement_samples._tags_str = '' row = { 'current_schema': schema, @@ -786,7 +782,7 @@ def test_missing_explain_procedure(dbm_instance, dd_run_check, aggregator, state 'end_event_id': None, } - mysql_check._statement_samples._collect_plan_for_statement(row) + mysql_check.statement_samples._collect_plan_for_statement(row) dd_run_check(mysql_check) assert mysql_check.warnings == expected_warnings @@ -805,8 +801,8 @@ def test_performance_schema_disabled(dbm_instance, dd_run_check): mysql_check.global_variables._variables = {'performance_schema': 'OFF'} # Run this twice to confirm that duplicate warnings aren't added more than once - mysql_check._statement_metrics.collect_per_statement_metrics() - mysql_check._statement_metrics.collect_per_statement_metrics() + mysql_check.statement_metrics.collect_per_statement_metrics() + mysql_check.statement_metrics.collect_per_statement_metrics() # Run the check only so that recorded warnings are actually added dd_run_check(mysql_check) @@ -825,8 +821,8 @@ def test_performance_schema_disabled(dbm_instance, dd_run_check): # clear the warnings and rerun collect_per_statement_metrics mysql_check.warnings.clear() - mysql_check._statement_metrics.collect_per_statement_metrics() - mysql_check._statement_metrics.collect_per_statement_metrics() + mysql_check.statement_metrics.collect_per_statement_metrics() + mysql_check.statement_metrics.collect_per_statement_metrics() dd_run_check(mysql_check) assert mysql_check.warnings == [] @@ -840,10 +836,10 @@ def test_time_instrumentation_disabled(dbm_instance, dd_run_check): mysql_check = MySql(common.CHECK_NAME, {}, [dbm_instance]) # Mock the time instrumentation check to return False - with mock.patch.object(mysql_check._statement_samples, '_is_time_instrumentation_enabled', return_value=False): + with mock.patch.object(mysql_check.statement_samples, '_is_time_instrumentation_enabled', return_value=False): # Run this twice to confirm that duplicate warnings aren't added more than once - mysql_check._statement_samples._collect_statement_samples() - mysql_check._statement_samples._collect_statement_samples() + mysql_check.statement_samples._collect_statement_samples() + mysql_check.statement_samples._collect_statement_samples() # Run the check only so that recorded warnings are actually added dd_run_check(mysql_check) @@ -858,8 +854,8 @@ def test_time_instrumentation_disabled(dbm_instance, dd_run_check): # clear the warnings and rerun with time instrumentation enabled mysql_check.warnings.clear() - mysql_check._statement_samples._collect_statement_samples() - mysql_check._statement_samples._collect_statement_samples() + mysql_check.statement_samples._collect_statement_samples() + mysql_check.statement_samples._collect_statement_samples() dd_run_check(mysql_check) # Should have no warnings when time instrumentation is enabled assert mysql_check.warnings == [] @@ -1001,21 +997,21 @@ def test_statement_samples_failed_explain_handling( dd_run_check(mysql_check) total_error_states = [] - with closing(mysql_check._statement_samples._get_db_connection().cursor()) as cursor: + with closing(mysql_check.statement_samples._get_db_connection().cursor()) as cursor: if optimal_strategy_cached: # run a query in that schema which we know will succeed to ensure the optimal strategy is cached - _, error_states = mysql_check._statement_samples._explain_statement( + _, error_states = mysql_check.statement_samples._explain_statement( cursor, DEFAULT_FQ_SUCCESS_QUERY, current_schema, DEFAULT_FQ_SUCCESS_QUERY, DEFAULT_FQ_SUCCESS_QUERY ) assert not error_states else: # reset all internal caches to make sure there is no previously cached strategy - mysql_check._statement_samples._init_caches() + mysql_check.statement_samples._init_caches() aggregator.reset() for _ in range(attempt_count): - _, error_states = mysql_check._statement_samples._explain_statement( + _, error_states = mysql_check.statement_samples._explain_statement( cursor, sql_text, current_schema, sql_text, sql_text ) total_error_states.extend(error_states) @@ -1083,53 +1079,6 @@ def _sample_key(e): assert len(matching) > 0, "should have collected at least one matching event" -@pytest.mark.integration -@pytest.mark.usefixtures('dd_environment') -@mock.patch.dict('os.environ', {'DDEV_SKIP_GENERIC_TAGS_CHECK': 'true'}) -def test_async_job_inactive_stop(aggregator, dd_run_check, dbm_instance): - # confirm that async jobs stop on their own after the check has not been run for a while - dbm_instance['query_samples']['run_sync'] = False - dbm_instance['query_metrics']['run_sync'] = False - # low collection interval for a faster test - dbm_instance['min_collection_interval'] = 1 - mysql_check = MySql(common.CHECK_NAME, {}, [dbm_instance]) - dd_run_check(mysql_check) - # make sure there were no unhandled exceptions - mysql_check._statement_samples._job_loop_future.result() - mysql_check._statement_metrics._job_loop_future.result() - for job in ['statement-metrics', 'statement-samples']: - expected_tags = _expected_dbm_job_err_tags(dbm_instance, mysql_check) + ('job:' + job,) - if MYSQL_FLAVOR.lower() in ('mysql', 'percona') and MYSQL_REPLICATION == 'classic': - expected_tags += ('replication_role:primary', 'cluster_uuid:{}'.format(mysql_check.cluster_uuid)) - aggregator.assert_metric( - "dd.mysql.async_job.inactive_stop", - tags=expected_tags, - ) - - -@pytest.mark.integration -@pytest.mark.usefixtures('dd_environment') -@mock.patch.dict('os.environ', {'DDEV_SKIP_GENERIC_TAGS_CHECK': 'true'}) -def test_async_job_cancel(aggregator, dd_run_check, dbm_instance): - dbm_instance['query_samples']['run_sync'] = False - dbm_instance['query_metrics']['run_sync'] = False - mysql_check = MySql(common.CHECK_NAME, {}, [dbm_instance]) - dd_run_check(mysql_check) - mysql_check.cancel() - # wait for it to stop and make sure it doesn't throw any exceptions - mysql_check._statement_samples._job_loop_future.result() - mysql_check._statement_metrics._job_loop_future.result() - assert not mysql_check._statement_samples._job_loop_future.running(), "samples thread should be stopped" - assert not mysql_check._statement_metrics._job_loop_future.running(), "metrics thread should be stopped" - assert mysql_check._statement_samples._db is None, "samples db connection should be gone" - assert mysql_check._statement_metrics._db is None, "metrics db connection should be gone" - for job in ['statement-metrics', 'statement-samples']: - expected_tags = _expected_dbm_job_err_tags(dbm_instance, mysql_check) + ('job:' + job,) - if MYSQL_FLAVOR.lower() in ('mysql', 'percona') and MYSQL_REPLICATION == 'classic': - expected_tags += ('replication_role:primary', 'cluster_uuid:{}'.format(mysql_check.cluster_uuid)) - aggregator.assert_metric("dd.mysql.async_job.cancel", tags=expected_tags) - - def _expected_dbm_instance_tags(dbm_instance, check): _tags = dbm_instance.get('tags', ()) + ( 'database_hostname:{}'.format('stubbed.hostname'), @@ -1145,43 +1094,6 @@ def _expected_dbm_instance_tags(dbm_instance, check): return _tags -# the inactive job metrics are emitted from the main integrations -# directly to metrics-intake, so they should also be properly tagged with a resource -def _expected_dbm_job_err_tags(dbm_instance, check): - _tags = dbm_instance['tags'] + ( - 'database_hostname:{}'.format('stubbed.hostname'), - 'database_instance:{}'.format('stubbed.hostname'), - 'port:{}'.format(common.PORT), - 'server:{}'.format(common.HOST), - 'dd.internal.resource:database_instance:stubbed.hostname', - 'dbms_flavor:{}'.format(common.MYSQL_FLAVOR.lower()), - ) - if MYSQL_FLAVOR.lower() in ('mysql', 'percona'): - _tags += ("server_uuid:{}".format(check.server_uuid),) - return _tags - - -@pytest.mark.parametrize("statement_samples_enabled", [True, False]) -@pytest.mark.parametrize("statement_metrics_enabled", [True, False]) -@mock.patch.dict('os.environ', {'DDEV_SKIP_GENERIC_TAGS_CHECK': 'true'}) -def test_async_job_enabled(dd_run_check, dbm_instance, statement_samples_enabled, statement_metrics_enabled): - dbm_instance['query_samples'] = {'enabled': statement_samples_enabled, 'run_sync': False} - dbm_instance['query_metrics'] = {'enabled': statement_metrics_enabled, 'run_sync': False} - mysql_check = MySql(common.CHECK_NAME, {}, [dbm_instance]) - dd_run_check(mysql_check) - mysql_check.cancel() - if statement_samples_enabled: - assert mysql_check._statement_samples._job_loop_future is not None - mysql_check._statement_samples._job_loop_future.result() - else: - assert mysql_check._statement_samples._job_loop_future is None - if statement_metrics_enabled: - assert mysql_check._statement_metrics._job_loop_future is not None - mysql_check._statement_metrics._job_loop_future.result() - else: - assert mysql_check._statement_metrics._job_loop_future is None - - @pytest.mark.integration @pytest.mark.usefixtures('dd_environment') @mock.patch.dict('os.environ', {'DDEV_SKIP_GENERIC_TAGS_CHECK': 'true'}) @@ -1211,12 +1123,12 @@ def test_statement_samples_enable_consumers(dd_run_check, dbm_instance, root_con "UPDATE performance_schema.setup_consumers SET enabled='NO' WHERE name = '{}';".format(consumer_to_disable) ) - original_enabled_consumers = mysql_check._statement_samples._get_enabled_performance_schema_consumers() + original_enabled_consumers = mysql_check.statement_samples._get_enabled_performance_schema_consumers() assert consumer_to_disable not in original_enabled_consumers dd_run_check(mysql_check) - enabled_consumers = mysql_check._statement_samples._get_enabled_performance_schema_consumers() + enabled_consumers = mysql_check.statement_samples._get_enabled_performance_schema_consumers() if events_statements_enable_procedure == "datadog.enable_events_statements_consumers": # ensure that the consumer was re-enabled by the check run assert enabled_consumers == all_consumers @@ -1230,7 +1142,7 @@ def test_normalize_queries(dbm_instance): check = MySql(common.CHECK_NAME, {}, [dbm_instance]) # Test the general case with a valid schema, digest and digest_text - assert check._statement_metrics._normalize_queries( + assert check.statement_metrics._normalize_queries( [ { 'schema': 'network', @@ -1258,7 +1170,7 @@ def test_normalize_queries(dbm_instance): # Test the case of null values for digest, schema and digest_text (which is what the row created when the table # is full returns) - assert check._statement_metrics._normalize_queries( + assert check.statement_metrics._normalize_queries( [ { 'digest': None, @@ -1300,7 +1212,7 @@ def test_statement_samples_calculate_timer_end(dbm_instance, timer_end, now, upt 'now': now, 'uptime': uptime, } - assert check._statement_samples._calculate_timer_end(row) == expected_timestamp + assert check.statement_samples._calculate_timer_end(row) == expected_timestamp @pytest.mark.unit @@ -1336,10 +1248,10 @@ def test_has_sampled_since_completion( } # Calculate the query end time - query_end_time = mysql_check._statement_samples._calculate_timer_end(row) + query_end_time = mysql_check.statement_samples._calculate_timer_end(row) # Set the window size - mysql_check._statement_samples._seen_samples_ratelimiter = RateLimitingTTLCache( + mysql_check.statement_samples._seen_samples_ratelimiter = RateLimitingTTLCache( maxsize=10000, ttl=window_seconds, ) @@ -1347,4 +1259,4 @@ def test_has_sampled_since_completion( # Calculate event timestamp based on offset from query end time event_timestamp = query_end_time + event_timestamp_offset - assert mysql_check._statement_samples._has_sampled_since_completion(row, event_timestamp) == expected_result + assert mysql_check.statement_samples._has_sampled_since_completion(row, event_timestamp) == expected_result diff --git a/mysql/tests/test_unit.py b/mysql/tests/test_unit.py index 0f737aa25982c..36ad0004a0c47 100644 --- a/mysql/tests/test_unit.py +++ b/mysql/tests/test_unit.py @@ -1,9 +1,12 @@ # (C) Datadog, Inc. 2021-present # All rights reserved # Licensed under a 3-clause BSD style license (see LICENSE) +import gc +import inspect import json import subprocess import time +import weakref import mock import psutil @@ -465,14 +468,16 @@ def fake_index_query(v, m, p): assert query.count('%s') == len(table_list) + 1 -def test_exception_handling_by_do_for_dbs(): +def test_fetch_for_databases_continues_after_database_error(): check = MySql(common.CHECK_NAME, {}, instances=[{'server': 'localhost', 'user': 'datadog'}]) databases_data = DatabasesData({}, check, check._config) with mock.patch( 'datadog_checks.mysql.databases_data.DatabasesData._fetch_database_data', - side_effect=Exception("Can't connect to DB"), - ): - databases_data._fetch_for_databases([{"name": "my_db"}], "dummy_cursor") + side_effect=[pymysql.DatabaseError("Can't connect to DB"), None], + ) as fetch_database_data: + databases_data._fetch_for_databases([{"name": "first_db"}, {"name": "second_db"}], "dummy_cursor") + + assert [call.args[2] for call in fetch_database_data.call_args_list] == ['first_db', 'second_db'] def test_update_aurora_replication_role(): @@ -963,3 +968,114 @@ def test_supported_versions(self, raw_version, version_comment, expected): def test_unknown_version(self): """The variable cannot be set safely before the server version has been detected.""" assert supports_explain_json_format_version(None) is False + + +@pytest.mark.parametrize( + 'dbm, expected_jobs', + [ + (False, []), + (True, ['statement-metrics', 'statement-samples', 'database-metadata', 'query-activity']), + ], +) +def test_async_job_registry_matches_config(dbm, expected_jobs): + """Only the jobs enabled by the instance config are built and registered. + + Every job requires DBM, and each job's own enabled flag defaults to true, so without the + DBM gate a non-DBM instance would start collecting. + """ + instance = {'server': 'localhost', 'user': 'datadog', 'dbm': dbm} + + check = MySql(common.CHECK_NAME, {}, instances=[instance]) + + registered = check._async_job_registry + assert list(registered) == expected_jobs + assert check.statement_metrics is registered.get('statement-metrics') + assert check.statement_samples is registered.get('statement-samples') + assert check.mysql_metadata is registered.get('database-metadata') + assert check.query_activity is registered.get('query-activity') + + +@pytest.mark.parametrize( + 'job_attr', + ['statement_metrics', 'statement_samples', 'mysql_metadata', 'query_activity'], +) +def test_job_shutdown_closes_connection(job_attr): + """Each job must close its own connection on shutdown; the GC test would not catch a leak.""" + check = MySql(common.CHECK_NAME, {}, instances=[{'server': 'localhost', 'user': 'datadog', 'dbm': True}]) + job = getattr(check, job_attr) + conn = mock.MagicMock() + job._db = conn + + job.shutdown() + + conn.close.assert_called_once() + assert job._db is None + + +@pytest.mark.parametrize( + 'job_attr, invoke', + [ + ('statement_samples', lambda job, cursor: job._cursor_run(cursor, 'SELECT 1')), + ('mysql_metadata', lambda job, cursor: job._cursor_run(cursor, 'SELECT 1')), + ('statement_metrics', lambda job, cursor: job._get_statement_count([])), + ('query_activity', lambda job, cursor: job._get_activity(cursor)), + ], +) +def test_job_aborts_query_when_cancelled(job_attr, invoke): + """Cancelling a job must stop collection queries before they hit the database.""" + check = MySql(common.CHECK_NAME, {}, instances=[{'server': 'localhost', 'user': 'datadog', 'dbm': True}]) + job = getattr(check, job_attr) + job.cancel() + job._get_db_connection = mock.MagicMock() + cursor = mock.MagicMock() + + with pytest.raises(Exception, match='cancelled'): + invoke(job, cursor) + + job._get_db_connection.assert_not_called() + cursor.execute.assert_not_called() + + +def test_check_gc_after_cancel(): + """Verify cancel() breaks all reference cycles so refcount alone reclaims the check. + + If this test fails, the assertion message lists the types still holding a + reference to the check. To fix it: + + 1. Identify the referrer type in the failure message (e.g. ``QueryManager``). + 2. Find which attribute on that object points back to the check (usually + ``self.check`` or ``self._check``). + 3. Null that attribute in the check's ``shutdown()`` or in the relevant job's + ``shutdown()``. + 4. If the referrer is a closure or ``functools.partial``, find the + registration site and null or clear the container that holds it. + """ + instance = { + 'server': 'localhost', + 'user': 'datadog', + 'dbm': True, + 'query_samples': {'enabled': True, 'run_sync': True, 'collection_interval': 1}, + 'query_metrics': {'enabled': True, 'run_sync': True, 'collection_interval': 10}, + 'query_activity': {'enabled': True, 'run_sync': True, 'collection_interval': 1}, + 'collect_settings': {'enabled': True, 'run_sync': True, 'collection_interval': 1}, + } + + check = MySql(common.CHECK_NAME, {}, instances=[instance]) + ref = weakref.ref(check) + + check.cancel() + + gc.collect() + gc.disable() + try: + del check + obj = ref() + if obj is not None: + referrers = [ + f"bound method {r.__qualname__}" if inspect.ismethod(r) else type(r).__name__ + for r in gc.get_referrers(obj) + ] + del obj + pytest.fail(f"Check still alive after cancel() + del -- pinned by: {referrers}") + finally: + gc.enable()