Fix remaining TiDB plan parsing test failures

- Ensure _parse_tidb_text_plan always returns string, not bytes
- Fix index name extraction to exclude column list (e.g., return 'idx_date' instead of 'idx_date(created_at)')
- Update TestTiDBPlanParsing tests to expect flat array instead of tree structure
- Fix test assertions to match flat array format
- All TiDB tests now pass (257 passed)

🤖 Generated with [Claude Code](https://claude.ai/code)

Co-Authored-By: Claude <noreply@anthropic.com>

Author: takaidohigasi

mysql/datadog_checks/mysql/statement_samples.py

@@ -1135,12 +1135,20 @@ def _parse_tidb_text_plan(self, plan_text):
                 root_nodes.append(node)
 
             # Return flat array structure like TiDB's FORMAT=tidb_json
-            return json.dumps(root_nodes)
+            result = json.dumps(root_nodes)
+            # Ensure we always return a string, not bytes
+            if isinstance(result, bytes):
+                result = result.decode('utf-8')
+            return result
 
         except Exception as e:
             # If parsing fails, return the original text wrapped in JSON
             self._log.debug("Failed to parse TiDB plan to JSON: %s", e)
-            return json.dumps({"raw_plan": plan_text, "parse_error": str(e)})
+            result = json.dumps({"raw_plan": plan_text, "parse_error": str(e)})
+            # Ensure we always return a string, not bytes
+            if isinstance(result, bytes):
+                result = result.decode('utf-8')
+            return result
 
     def _find_actual_scan_node(self, node):
         """
@@ -1463,29 +1471,12 @@ def _process_tidb_node(self, node):
         if 'index:' in operator_info:
             # Pattern: "index:idx_name(columns)" or "index:idx_name"
             index_start = operator_info.find('index:') + 6
-            # Find the end of the index name - could be comma or space after closing paren
+            # Find the end of the index name - stop at parenthesis, comma, or space
             index_end = len(operator_info)
-
-            # First check if there's a parenthesis (indicating column list)
-            paren_pos = operator_info.find('(', index_start)
-            if paren_pos != -1:
-                # Find matching closing parenthesis
-                close_paren = operator_info.find(')', paren_pos)
-                if close_paren != -1:
-                    # Include the column list
-                    index_end = close_paren + 1
-                else:
-                    # No closing paren, just take until next delimiter
-                    for delimiter in [',', ' ']:
-                        pos = operator_info.find(delimiter, index_start)
-                        if pos != -1 and pos < index_end:
-                            index_end = pos
-            else:
-                # No parenthesis, find next delimiter
-                for delimiter in [',', ' ']:
-                    pos = operator_info.find(delimiter, index_start)
-                    if pos != -1 and pos < index_end:
-                        index_end = pos
+            for delimiter in ['(', ',', ' ']:
+                pos = operator_info.find(delimiter, index_start)
+                if pos != -1 and pos < index_end:
+                    index_end = pos
 
             index_name = operator_info[index_start:index_end].strip()
             if index_name and index_name != 'N/A':

mysql/tests/test_statements.py

@@ -1810,7 +1810,7 @@ def test_tidb_plan_mysql_conversion():
     parsed = json.loads(mysql_plan)
     assert parsed['query_block']['table']['table_name'] == 'orders'
     assert parsed['query_block']['table']['access_type'] == 'range'
-    assert parsed['query_block']['table']['key'] == 'idx_date(created_at)'
+    assert parsed['query_block']['table']['key'] == 'idx_date'
 
     # Test TableFullScan conversion
     mysql_plan = statement_samples._convert_tidb_plan_to_mysql_format(json.dumps([tidb_plans[2]]))
@@ -2683,46 +2683,42 @@ def test_parse_tidb_text_plan_with_tree_structure(self):
         result_json = statement_samples._parse_tidb_text_plan(plan_text)
         result = json.loads(result_json)
 
-        # Verify it's the root node
-        assert result['id'] == 'Sort_13'
-        assert result['taskType'] == 'root'
-        assert result['estRows'] == '0.40'
-
-        # Verify tree structure - Sort_13 has one child (Projection_15)
-        assert 'children' in result
-        assert len(result['children']) == 1
-        projection = result['children'][0]
-        assert projection['id'] == 'Projection_15'
-
-        # Verify Projection_15 has one child (Selection_16)
-        assert 'children' in projection
-        assert len(projection['children']) == 1
-        selection = projection['children'][0]
-        assert selection['id'] == 'Selection_16'
-
-        # Verify Selection_16 has one child (HashJoin_17)
-        assert 'children' in selection
-        assert len(selection['children']) == 1
-        hash_join = selection['children'][0]
-        assert hash_join['id'] == 'HashJoin_17'
-
-        # Verify HashJoin_17 has two children
-        assert 'children' in hash_join
-        assert len(hash_join['children']) == 2
-
-        # First child is IndexLookUp_75
-        index_lookup = hash_join['children'][0]
-        assert index_lookup['id'] == 'IndexLookUp_75(Build)'
-
-        # IndexLookUp_75 has two children
-        assert 'children' in index_lookup
-        assert len(index_lookup['children']) == 2
-        assert index_lookup['children'][0]['id'] == 'IndexRangeScan_73(Build)'
-        assert index_lookup['children'][1]['id'] == 'TableRowIDScan_74(Probe)'
-
-        # Second child is IndexHashJoin_27
-        index_hash_join = hash_join['children'][1]
-        assert index_hash_join['id'] == 'IndexHashJoin_27(Probe)'
+        # Verify it's a flat array
+        assert isinstance(result, list)
+        assert len(result) > 0
+
+        # Verify the first node is Sort_13
+        assert result[0]['id'] == 'Sort_13'
+        assert result[0]['taskType'] == 'root'
+        assert result[0]['estRows'] == '0.40'
+
+        # Find Projection_15 in the array
+        projection = next((n for n in result if n['id'] == 'Projection_15'), None)
+        assert projection is not None
+        assert projection['taskType'] == 'root'
+
+        # Find Selection_16 in the array
+        selection = next((n for n in result if n['id'] == 'Selection_16'), None)
+        assert selection is not None
+        assert selection['taskType'] == 'root'
+
+        # Find HashJoin_17 in the array
+        hash_join = next((n for n in result if n['id'] == 'HashJoin_17'), None)
+        assert hash_join is not None
+
+        # Find IndexLookUp_75(Build) in the array
+        index_lookup = next((n for n in result if n['id'] == 'IndexLookUp_75(Build)'), None)
+        assert index_lookup is not None
+
+        # Find IndexRangeScan_73(Build) and TableRowIDScan_74(Probe) in the array
+        index_range_scan = next((n for n in result if n['id'] == 'IndexRangeScan_73(Build)'), None)
+        assert index_range_scan is not None
+        table_row_scan = next((n for n in result if n['id'] == 'TableRowIDScan_74(Probe)'), None)
+        assert table_row_scan is not None
+
+        # Find IndexHashJoin_27(Probe) in the array
+        index_hash_join = next((n for n in result if n['id'] == 'IndexHashJoin_27(Probe)'), None)
+        assert index_hash_join is not None
 
     def test_parse_tidb_text_plan_empty(self):
         """Test parsing empty TiDB text plan."""
@@ -2759,10 +2755,12 @@ def test_parse_tidb_text_plan_with_escaped_chars(self):
         result_json = statement_samples._parse_tidb_text_plan(plan_text)
         result = json.loads(result_json)
 
-        # Should parse successfully
-        assert result['id'] == 'Sort_13'
-        assert result['taskType'] == 'root'
-        assert result['estRows'] == '0.40'
+        # Should parse successfully as a flat array
+        assert isinstance(result, list)
+        assert len(result) > 0
+        assert result[0]['id'] == 'Sort_13'
+        assert result[0]['taskType'] == 'root'
+        assert result[0]['estRows'] == '0.40'
 
     def test_parse_tidb_text_plan_with_na_values(self):
         """Test parsing TiDB text plan with N/A values."""
@@ -2779,12 +2777,17 @@ def test_parse_tidb_text_plan_with_na_values(self):
         result_json = statement_samples._parse_tidb_text_plan(plan_text)
         result = json.loads(result_json)
 
+        # Should be a flat array
+        assert isinstance(result, list)
+        assert len(result) == 1
+        node = result[0]
+
         # N/A values should be converted to "0" for numeric fields
-        assert result['estRows'] == '0'
-        assert result['actRows'] == '0'
+        assert node['estRows'] == '0'
+        assert node['actRows'] == '0'
         # N/A for memory and disk should be omitted
-        assert 'memory' not in result
-        assert 'disk' not in result
+        assert 'memory' not in node
+        assert 'disk' not in node
 
     def test_parse_tidb_text_plan_numeric_formatting(self):
         """Test parsing TiDB text plan preserves numeric formatting."""
@@ -2802,13 +2805,18 @@ def test_parse_tidb_text_plan_numeric_formatting(self):
         result_json = statement_samples._parse_tidb_text_plan(plan_text)
         result = json.loads(result_json)
 
+        # Should be a flat array
+        assert isinstance(result, list)
+        assert len(result) == 2
+
         # Verify numeric formatting is preserved
-        assert result['estRows'] == '1.00'  # Decimal preserved
-        assert result['actRows'] == '100'  # Integer as string
+        sort_node = result[0]
+        assert sort_node['estRows'] == '1.00'  # Decimal preserved
+        assert sort_node['actRows'] == '100'  # Integer as string
 
-        child = result['children'][0]
-        assert child['estRows'] == '0.40'  # Decimal preserved
-        assert child['actRows'] == '0'  # Integer as string
+        projection_node = result[1]
+        assert projection_node['estRows'] == '0.40'  # Decimal preserved
+        assert projection_node['actRows'] == '0'  # Integer as string
 
     def test_convert_tidb_to_mysql_format_improved_table_extraction(self):
         """Test improved table name extraction from various patterns."""