Add TiDB database monitoring support

Author: takaidohigasi

mysql/changelog.d/20826.changed

@@ -0,0 +1 @@
+Implement TiDB database monitoring

mysql/datadog_checks/mysql/activity.py

@@ -183,6 +183,11 @@ def run_job(self):
                 'Waiting for events_waits_current availability to be determined by the check, skipping run.'
             )
         if self._check.events_wait_current_enabled is False:
+            # Skip warning for TiDB as it doesn't have performance schema consumers
+            if self._check._get_is_tidb(self._db):
+                self._log.debug("TiDB detected, skipping query activity collection")
+                return
+
             azure_deployment_type = self._config.cloud_metadata.get("azure", {}).get("deployment_type")
             if azure_deployment_type != "flexible_server":
                 self._check.record_warning(

mysql/datadog_checks/mysql/config.py

@@ -52,6 +52,7 @@ def __init__(self, instance, init_config):
         # Backward compatibility: check new names first, then fall back to old names
         self.schemas_config: dict = instance.get('collect_schemas', instance.get('schemas_collection', {})) or {}
         self.index_config: dict = instance.get('index_metrics', {}) or {}
+        self.slow_query_config: dict = instance.get('slow_query', {}) or {}
         self.collect_blocking_queries = is_affirmative(instance.get('collect_blocking_queries', False))
 
         self.cloud_metadata = {}

mysql/datadog_checks/mysql/mysql.py

@@ -81,10 +81,12 @@
     SQL_REPLICATION_ROLE_AWS_AURORA,
     SQL_SERVER_ID_AWS_AURORA,
     SQL_SERVER_UUID,
+    SQL_TIDB_VERSION,
     show_replica_status_query,
 )
 from .statement_samples import MySQLStatementSamples
 from .statements import MySQLStatementMetrics
+from .tidb_slow_query import TiDBSlowQueryMonitor
 from .util import DatabaseConfigurationError, connect_with_session_variables  # noqa: F401
 from .version_utils import get_version
 
@@ -121,6 +123,7 @@ def __init__(self, name, init_config, instances):
         self._agent_hostname = None
         self._database_hostname = None
         self._is_aurora = None
+        self._is_tidb = None
         self._performance_schema_enabled = None
         self._events_wait_current_enabled = None
         self._group_replication_active = None
@@ -146,6 +149,7 @@ def __init__(self, name, init_config, instances):
         self._mysql_metadata = MySQLMetadata(self, self._config, self._get_connection_args())
         self._query_activity = MySQLActivity(self, self._config, self._get_connection_args())
         self._index_metrics = MySqlIndexMetrics(self._config)
+        self._tidb_slow_query_monitor = None  # Initialized later if TiDB is detected
         # _database_instance_emitted: limit the collection and transmission of the database instance metadata
         self._database_instance_emitted = TTLCache(
             maxsize=1,
@@ -288,6 +292,7 @@ def set_resource_tags(self):
     def set_version(self, db):
         self.version = get_version(db)
         self.is_mariadb = self.version.flavor == "MariaDB"
+        self.is_tidb = self.version.flavor == "TiDB"
         self.tag_manager.set_tag("dbms_flavor", self.version.flavor.lower(), replace=True)
 
     def set_server_uuid(self, db):
@@ -314,6 +319,12 @@ def _check_database_configuration(self, db):
         self._is_group_replication_active(db)
 
     def _check_performance_schema_enabled(self, db):
+        # TiDB doesn't have performance_schema but uses information_schema.cluster_statements_summary instead
+        if self._get_is_tidb(db):
+            self._performance_schema_enabled = False
+            self.log.debug("TiDB detected, performance_schema check skipped")
+            return self._performance_schema_enabled
+
         with closing(db.cursor(CommenterCursor)) as cursor:
             cursor.execute("SHOW VARIABLES LIKE 'performance_schema'")
             results = dict(cursor.fetchall())
@@ -420,6 +431,15 @@ def check(self, _):
                     self._query_activity.run_job_loop(dbm_tags)
                     self._mysql_metadata.run_job_loop(dbm_tags)
 
+                    # Initialize and run TiDB slow query monitor if this is TiDB
+                    if self._get_is_tidb(db) and self._tidb_slow_query_monitor is None:
+                        self._tidb_slow_query_monitor = TiDBSlowQueryMonitor(
+                            self, self._config, self._get_connection_args()
+                        )
+
+                    if self._tidb_slow_query_monitor:
+                        self._tidb_slow_query_monitor.run_job_loop(dbm_tags)
+
                 # keeping track of these:
                 self._put_qcache_stats()
 
@@ -438,6 +458,8 @@ def cancel(self):
         self._statement_metrics.cancel()
         self._query_activity.cancel()
         self._mysql_metadata.cancel()
+        if self._tidb_slow_query_monitor:
+            self._tidb_slow_query_monitor.cancel()
 
     def _new_query_executor(self, queries):
         return QueryExecutor(
@@ -462,7 +484,8 @@ def _get_runtime_queries(self, db):
 
         if self.performance_schema_enabled:
             queries.extend([QUERY_USER_CONNECTIONS])
-        if self._index_metrics.include_index_metrics:
+        # TiDB doesn't have mysql.innodb_index_stats table
+        if self._index_metrics.include_index_metrics and not self._get_is_tidb(db):
             queries.extend(self._index_metrics.queries)
         self._runtime_queries_cached = self._new_query_executor(queries)
         self._runtime_queries_cached.compile_queries()
@@ -582,7 +605,11 @@ def _collect_metrics(self, db, tags):
             # Innodb metrics are not available for Aurora reader instances
             if self._is_aurora and self._replication_role == "reader":
                 self.log.debug("Skipping innodb metrics collection for reader instance")
+            # TiDB does not support InnoDB storage engine
+            elif self._get_is_tidb(db):
+                self.log.info("Skipping innodb metrics collection for TiDB instance")
             else:
+                self.log.debug("Collecting InnoDB metrics (not Aurora reader, not TiDB)")
                 with tracked_query(self, operation="innodb_metrics"):
                     results.update(self.innodb_stats.get_stats_from_innodb_status(db))
                 self.innodb_stats.process_innodb_stats(results, self._config.options, metrics)
@@ -593,23 +620,29 @@ def _collect_metrics(self, db, tags):
                 results['Binlog_space_usage_bytes'] = self._get_binary_log_stats(db)
 
         # Compute key cache utilization metric
-        key_blocks_unused = collect_scalar('Key_blocks_unused', results)
-        key_cache_block_size = collect_scalar('key_cache_block_size', results)
-        key_buffer_size = collect_scalar('key_buffer_size', results)
-        results['Key_buffer_size'] = key_buffer_size
+        # TiDB doesn't have MyISAM key cache metrics
+        if not self._get_is_tidb(db):
+            key_blocks_unused = collect_scalar('Key_blocks_unused', results)
+            key_cache_block_size = collect_scalar('key_cache_block_size', results)
+            key_buffer_size = collect_scalar('key_buffer_size', results)
+            results['Key_buffer_size'] = key_buffer_size
 
-        try:
-            # can be null if the unit is missing in the user config (4 instead of 4G for eg.)
-            if key_buffer_size != 0:
-                key_cache_utilization = 1 - ((key_blocks_unused * key_cache_block_size) / key_buffer_size)
-                results['Key_cache_utilization'] = key_cache_utilization
-
-            results['Key_buffer_bytes_used'] = collect_scalar('Key_blocks_used', results) * key_cache_block_size
-            results['Key_buffer_bytes_unflushed'] = (
-                collect_scalar('Key_blocks_not_flushed', results) * key_cache_block_size
-            )
-        except TypeError as e:
-            self.log.error("Not all Key metrics are available, unable to compute: %s", e)
+            try:
+                # can be null if the unit is missing in the user config (4 instead of 4G for eg.)
+                if key_buffer_size != 0 and key_cache_block_size is not None and key_blocks_unused is not None:
+                    key_cache_utilization = 1 - ((key_blocks_unused * key_cache_block_size) / key_buffer_size)
+                    results['Key_cache_utilization'] = key_cache_utilization
+
+                if key_cache_block_size is not None:
+                    key_blocks_used = collect_scalar('Key_blocks_used', results)
+                    key_blocks_not_flushed = collect_scalar('Key_blocks_not_flushed', results)
+
+                    if key_blocks_used is not None:
+                        results['Key_buffer_bytes_used'] = key_blocks_used * key_cache_block_size
+                    if key_blocks_not_flushed is not None:
+                        results['Key_buffer_bytes_unflushed'] = key_blocks_not_flushed * key_cache_block_size
+            except (TypeError, ZeroDivisionError) as e:
+                self.log.error("Not all Key metrics are available, unable to compute: %s", e)
 
         metrics.update(VARIABLES_VARS)
         metrics.update(INNODB_VARS)
@@ -681,7 +714,10 @@ def _collect_metrics(self, db, tags):
             metrics.update(TABLE_VARS)
 
         if self._config.replication_enabled:
-            if self.performance_schema_enabled and self._group_replication_active:
+            # TiDB does not support MySQL replication
+            if self._get_is_tidb(db):
+                self.log.debug("Skipping replication metrics collection for TiDB instance")
+            elif self.performance_schema_enabled and self._group_replication_active:
                 self.log.debug('Collecting group replication metrics.')
                 with tracked_query(self, operation="group_replication_metrics"):
                     self._collect_group_replica_metrics(db, results)
@@ -1126,6 +1162,35 @@ def _get_is_aurora(self, db):
 
         return self._is_aurora
 
+    def _get_is_tidb(self, db):
+        """
+        Tests if the instance is a TiDB database and caches the result.
+        """
+        if self._is_tidb is not None:
+            return self._is_tidb
+
+        try:
+            with closing(db.cursor(CommenterCursor)) as cursor:
+                cursor.execute(SQL_TIDB_VERSION)
+                result = cursor.fetchone()
+                self.log.debug("TiDB detection - VERSION() result: %s", result)
+                if result and 'TiDB' in result[0]:
+                    self._is_tidb = True
+                    self.log.info("Detected TiDB instance")
+                else:
+                    self._is_tidb = False
+
+        except Exception:
+            self.warning(
+                "Unable to determine if server is TiDB. If this is a TiDB database, some "
+                "information may be unavailable: %s",
+                traceback.format_exc(),
+            )
+            self._is_tidb = False
+            return False
+
+        return self._is_tidb
+
     @classmethod
     def _get_stats_from_status(cls, db):
         with closing(db.cursor(CommenterCursor)) as cursor:
@@ -1164,6 +1229,12 @@ def _check_innodb_engine_enabled(self, db):
         # table. Later is chosen because that involves no string parsing.
         if self._is_innodb_engine_enabled_cached is not None:
             return self._is_innodb_engine_enabled_cached
+
+        # TiDB doesn't have InnoDB engine
+        if self._get_is_tidb(db):
+            self._is_innodb_engine_enabled_cached = False
+            return self._is_innodb_engine_enabled_cached
+
         try:
             with closing(db.cursor(CommenterCursor)) as cursor:
                 cursor.execute(SQL_INNODB_ENGINES)
@@ -1439,13 +1510,16 @@ def _report_warnings(self):
 
     def _send_database_instance_metadata(self):
         if self.database_identifier not in self._database_instance_emitted:
+            # Keep dbms as "mysql" for all MySQL-compatible databases
+            # The specific flavor is identified through the dbms_flavor tag
+            dbms = "mysql"
             event = {
                 "host": self.reported_hostname,
                 "port": self._config.port,
                 "database_instance": self.database_identifier,
                 "database_hostname": self.database_hostname,
                 "agent_version": datadog_agent.get_version(),
-                "dbms": "mysql",
+                "dbms": dbms,
                 "kind": "database_instance",
                 "collection_interval": self._config.database_instance_collection_interval,
                 'dbms_version': self.version.version + '+' + self.version.build,
@@ -1471,6 +1545,9 @@ def set_cluster_tags(self, db):
         if self._group_replication_active:
             self.log.debug("Group replication cluster tags are not currently supported")
             return
+        if self._get_is_tidb(db):
+            self.log.debug("TiDB cluster tags are not currently supported")
+            return
 
         replica_status = self._get_replica_replication_status(db)
         if replica_status:

mysql/datadog_checks/mysql/queries.py

@@ -71,6 +71,9 @@
 FROM information_schema.replica_host_status
 WHERE server_id = @@aurora_server_id"""
 
+SQL_TIDB_VERSION = """\
+SELECT VERSION()"""
+
 SQL_GROUP_REPLICATION_MEMBER = """\
 SELECT channel_name, member_state
 FROM performance_schema.replication_group_members

mysql/datadog_checks/mysql/statement_samples.py

@@ -335,32 +335,74 @@ def _cursor_run(self, cursor, query, params=None, obfuscated_params=None, obfusc
             )
             raise
 
+    def _get_tidb_statement_samples(self):
+        """
+        Get statement samples from TiDB's cluster_statements_summary table.
+
+        TiDB doesn't have events_statements_current, but we can use cluster_statements_summary
+        with QUERY_SAMPLE_TEXT which contains a sample query for each digest.
+        """
+        with closing(self._get_db_connection().cursor(CommenterDictCursor)) as cursor:
+            # TiDB-specific query to get statement samples
+            query = """
+            SELECT
+                SCHEMA_NAME as current_schema,
+                QUERY_SAMPLE_TEXT as sql_text,
+                DIGEST as digest,
+                DIGEST_TEXT as digest_text,
+                -- TiDB doesn't have these fields, so we'll use defaults
+                0 as end_event_id,
+                0 as timer_start,
+                0 as uptime,
+                unix_timestamp() as now,
+                0 as timer_end,
+                AVG_LATENCY as timer_wait_ns,
+                0 as lock_time_ns,
+                SAMPLE_USER as processlist_user,
+                '' as processlist_host,
+                EXEC_COUNT as rows_sent,
+                AVG_AFFECTED_ROWS as rows_affected
+            FROM information_schema.cluster_statements_summary
+            WHERE QUERY_SAMPLE_TEXT IS NOT NULL
+                AND QUERY_SAMPLE_TEXT != ''
+                AND LAST_SEEN > DATE_SUB(NOW(), INTERVAL 1 MINUTE)
+            ORDER BY LAST_SEEN DESC
+            LIMIT 100
+            """
+            self._cursor_run(cursor, query)
+            return cursor.fetchall()
+
     @tracked_method(agent_check_getter=attrgetter('_check'))
     def _get_new_events_statements_current(self):
         start = time.time()
-        with closing(self._get_db_connection().cursor(CommenterDictCursor)) as cursor:
-            self._cursor_run(
-                cursor,
-                "set @uptime = {}".format(UPTIME_SUBQUERY.format(global_status_table=self._global_status_table)),
-            )
-            self._cursor_run(cursor, EVENTS_STATEMENTS_CURRENT_QUERY)
-            rows = cursor.fetchall()
-            tags = (
-                self._tags
-                + ["events_statements_table:{}".format(EVENTS_STATEMENTS_TABLE)]
-                + self._check._get_debug_tags()
-            )
-            self._check.histogram(
-                "dd.mysql.get_new_events_statements.time",
-                (time.time() - start) * 1000,
-                tags=tags,
-                hostname=self._check.reported_hostname,
-            )
-            self._check.histogram(
-                "dd.mysql.get_new_events_statements.rows", len(rows), tags=tags, hostname=self._check.reported_hostname
-            )
-            self._log.debug("Read %s rows from %s", len(rows), EVENTS_STATEMENTS_TABLE)
-            return rows
+
+        # Check if this is TiDB and use appropriate query
+        if self._check._get_is_tidb(self._db):
+            rows = self._get_tidb_statement_samples()
+        else:
+            with closing(self._get_db_connection().cursor(CommenterDictCursor)) as cursor:
+                self._cursor_run(
+                    cursor,
+                    "set @uptime = {}".format(UPTIME_SUBQUERY.format(global_status_table=self._global_status_table)),
+                )
+                self._cursor_run(cursor, EVENTS_STATEMENTS_CURRENT_QUERY)
+                rows = cursor.fetchall()
+
+        table_name = (
+            "tidb_cluster_statements_summary" if self._check._get_is_tidb(self._db) else EVENTS_STATEMENTS_TABLE
+        )
+        tags = self._tags + ["events_statements_table:{}".format(table_name)] + self._check._get_debug_tags()
+        self._check.histogram(
+            "dd.mysql.get_new_events_statements.time",
+            (time.time() - start) * 1000,
+            tags=tags,
+            hostname=self._check.reported_hostname,
+        )
+        self._check.histogram(
+            "dd.mysql.get_new_events_statements.rows", len(rows), tags=tags, hostname=self._check.reported_hostname
+        )
+        self._log.debug("Read %s rows from %s", len(rows), table_name)
+        return rows
 
     def _filter_valid_statement_rows(self, rows):
         num_sent = 0
@@ -535,6 +577,16 @@ def _get_sample_collection_strategy(self):
             self._log.debug("Using cached events_statements_strategy: %s", cached_strategy)
             return cached_strategy
 
+        # Check if this is TiDB - it doesn't have performance_schema consumers
+        if self._check._get_is_tidb(self._db):
+            self._log.debug("TiDB detected, using TiDB statement samples strategy")
+            collection_interval = self._configured_collection_interval
+            if collection_interval < 0:
+                collection_interval = 10  # Default 10 seconds for TiDB
+            strategy = ("tidb_cluster_statements_summary", collection_interval)
+            self._collection_strategy_cache["events_statements_strategy"] = strategy
+            return strategy
+
         enabled_consumers = self._get_enabled_performance_schema_consumers()
         if len(enabled_consumers) < 3:
             self._enable_events_statements_consumers()
@@ -829,6 +881,10 @@ def _run_fully_qualified_explain_procedure(self, schema, cursor, statement, obfu
             raise
 
     def _has_sampled_since_completion(self, row, event_timestamp):
+        # TiDB doesn't have end_event_id, so always return False for TiDB
+        if self._check._get_is_tidb(self._db):
+            return False
+
         # If the query has finished end_event_id will be set
         if row['end_event_id']:
             query_end_time = self._calculate_timer_end(row)