SQLServer XE add missing event fields (#20293)

* fill in extra fields

* move primary field to sqlserver struct for RQT events

* changelog

Author: azhou-datadog

sqlserver/changelog.d/20293.added

@@ -0,0 +1 @@
+Fill in missing fields for XE events

sqlserver/datadog_checks/sqlserver/xe_collection/base.py

@@ -430,6 +430,17 @@ def _create_event_payload(self, raw_event):
         if 'query_signature' in raw_event:
             normalized_event['query_signature'] = raw_event['query_signature']
 
+        # Add primary_sql_field if available
+        if 'primary_sql_field' in raw_event:
+            normalized_event['primary_sql_field'] = raw_event['primary_sql_field']
+
+        # Add metadata if available
+        normalized_event['metadata'] = {
+            'tables': raw_event.get('dd_tables'),
+            'commands': raw_event.get('dd_commands'),
+            'comments': raw_event.get('dd_comments'),
+        }
+
         return {
             "host": self._check.resolved_hostname,
             "database_instance": self._check.database_identifier,
@@ -518,7 +529,11 @@ def run_job(self):
         for event in events:
             try:
                 # Obfuscate SQL fields and get the raw statement
-                obfuscated_event, raw_sql_fields = self._obfuscate_sql_fields(event)
+                obfuscated_event, raw_sql_fields, primary_sql_field = self._obfuscate_sql_fields(event)
+
+                # Add primary SQL field to the event if available
+                if primary_sql_field:
+                    obfuscated_event['primary_sql_field'] = primary_sql_field
 
                 # Create a properly structured payload for the individual event
                 payload = self._create_event_payload(obfuscated_event)
@@ -607,6 +622,7 @@ def _obfuscate_sql_fields(self, event):
         """SQL field obfuscation and signature creation"""
         obfuscated_event = event.copy()
         raw_sql_fields = {}
+        primary_sql_field = None
 
         # Get SQL fields for this event type
         sql_fields = self.get_sql_fields(event.get('event_name', ''))
@@ -636,8 +652,9 @@ def _obfuscate_sql_fields(self, event):
                         obfuscated_event['dd_comments'].extend(result['metadata']['comments'])
 
                     # Compute query_signature and raw_query_signature from the primary field
-                    primary_field = self._get_primary_sql_field(event)
-                    if field == primary_field or 'query_signature' not in obfuscated_event:
+                    current_primary_field = self._get_primary_sql_field(event)
+                    if field == current_primary_field or 'query_signature' not in obfuscated_event:
+                        primary_sql_field = field  # Store the field used for signature
                         obfuscated_event['query_signature'] = compute_sql_signature(result['query'])
                         raw_signature = compute_sql_signature(event[field])
                         raw_sql_fields['raw_query_signature'] = raw_signature
@@ -652,7 +669,7 @@ def _obfuscate_sql_fields(self, event):
         if 'dd_comments' in obfuscated_event:
             obfuscated_event['dd_comments'] = list(set(obfuscated_event['dd_comments']))
 
-        return obfuscated_event, raw_sql_fields if raw_sql_fields else None
+        return obfuscated_event, raw_sql_fields if raw_sql_fields else None, primary_sql_field
 
     def _get_primary_sql_field(self, event):
         """
@@ -698,7 +715,7 @@ def _create_rqt_event(self, event, raw_sql_fields, query_details):
             return None
 
         # Get the primary SQL field for this event type
-        primary_field = self._get_primary_sql_field(event)
+        primary_field = event.get('primary_sql_field') or self._get_primary_sql_field(event)
         if not primary_field or primary_field not in raw_sql_fields:
             self._log.debug(
                 f"Skipping RQT event creation: Primary SQL field {primary_field} not found in raw_sql_fields"
@@ -729,6 +746,7 @@ def _create_rqt_event(self, event, raw_sql_fields, query_details):
             "session_id": event.get("session_id"),
             "xe_type": event.get("event_name"),
             "event_fire_timestamp": query_details.get("event_fire_timestamp"),
+            "primary_sql_field": primary_field,
         }
 
         # Only include duration and query_start for non-error events

sqlserver/tests/test_integration.py

@@ -1039,6 +1039,22 @@ def test_xe_collection_integration(aggregator, dd_run_check, bob_conn, instance_
             assert 'raw_query_signature' in query_details, "raw_query_signature not found in query details"
             assert query_details.get('raw_query_signature'), "raw_query_signature is empty"
 
+            # Verify primary_sql_field is present
+            assert 'primary_sql_field' in query_details, "primary_sql_field not found in query details"
+            assert query_details.get('primary_sql_field') in [
+                'batch_text',
+                'sql_text',
+                'statement',
+            ], f"Unexpected primary_sql_field value: {query_details.get('primary_sql_field')}"
+
+            # Verify metadata is present
+            assert 'metadata' in query_details, "metadata not found in query details"
+            metadata = query_details.get('metadata', {})
+            assert isinstance(metadata, dict), "metadata is not a dictionary"
+            assert 'tables' in metadata, "tables not found in metadata"
+            assert 'commands' in metadata, "commands not found in metadata"
+            assert 'comments' in metadata, "comments not found in metadata"
+
     assert found_test_query, "Could not find our specific test query in the completion events"
 
     # Verify specific error event details
@@ -1058,4 +1074,20 @@ def test_xe_collection_integration(aggregator, dd_run_check, bob_conn, instance_
             assert 'raw_query_signature' in query_details, "raw_query_signature not found in error query details"
             assert query_details.get('raw_query_signature'), "raw_query_signature is empty"
 
+            # Verify primary_sql_field is present
+            assert 'primary_sql_field' in query_details, "primary_sql_field not found in error query details"
+            assert query_details.get('primary_sql_field') in [
+                'batch_text',
+                'sql_text',
+                'statement',
+            ], f"Unexpected primary_sql_field value: {query_details.get('primary_sql_field')}"
+
+            # Verify metadata is present
+            assert 'metadata' in query_details, "metadata not found in error query details"
+            metadata = query_details.get('metadata', {})
+            assert isinstance(metadata, dict), "metadata is not a dictionary"
+            assert 'tables' in metadata, "tables not found in metadata"
+            assert 'commands' in metadata, "commands not found in metadata"
+            assert 'comments' in metadata, "comments not found in metadata"
+
     assert found_error_query, "Could not find our specific error query in the error events"

sqlserver/tests/test_xe_collection.py

@@ -733,7 +733,7 @@ def test_obfuscate_sql_fields(self, mock_compute_signature, mock_obfuscate, quer
             'sql_text': 'SELECT * FROM Customers WHERE CustomerId = 123',
         }
 
-        obfuscated_event, raw_sql_fields = query_completion_handler._obfuscate_sql_fields(event)
+        obfuscated_event, raw_sql_fields, primary_sql_field = query_completion_handler._obfuscate_sql_fields(event)
 
         # Verify obfuscated fields
         assert obfuscated_event['batch_text'] == 'SELECT * FROM Customers WHERE CustomerId = ?'
@@ -747,6 +747,9 @@ def test_obfuscate_sql_fields(self, mock_compute_signature, mock_obfuscate, quer
         assert raw_sql_fields['sql_text'] == 'SELECT * FROM Customers WHERE CustomerId = 123'
         assert raw_sql_fields['raw_query_signature'] == 'abc123'
 
+        # Verify primary SQL field
+        assert primary_sql_field == 'batch_text'
+
         # Verify raw_query_signature is added to the obfuscated event when collect_raw_query is enabled
         assert 'raw_query_signature' in obfuscated_event
         assert obfuscated_event['raw_query_signature'] == 'abc123'
@@ -836,6 +839,10 @@ def test_create_event_payload(self, mock_agent, query_completion_handler):
             'database_name': 'TestDB',
             'batch_text': 'SELECT * FROM Customers WHERE CustomerId = 123',
             'query_signature': 'abc123',
+            'primary_sql_field': 'batch_text',
+            'dd_tables': ['Customers'],
+            'dd_commands': ['SELECT'],
+            'dd_comments': [],
         }
 
         # Create payload
@@ -854,6 +861,14 @@ def test_create_event_payload(self, mock_agent, query_completion_handler):
         assert query_details['request_id'] == 456
         assert query_details['database_name'] == 'TestDB'
         assert query_details['query_signature'] == 'abc123'
+        assert query_details['primary_sql_field'] == 'batch_text'
+
+        # Verify metadata structure
+        assert 'metadata' in query_details
+        metadata = query_details['metadata']
+        assert metadata['tables'] == ['Customers']
+        assert metadata['commands'] == ['SELECT']
+        assert metadata['comments'] == []
 
     @patch('datadog_checks.sqlserver.xe_collection.base.datadog_agent')
     def test_create_rqt_event(self, mock_agent, query_completion_handler):
@@ -869,6 +884,10 @@ def test_create_rqt_event(self, mock_agent, query_completion_handler):
             'database_name': 'TestDB',
             'batch_text': 'SELECT * FROM Customers WHERE CustomerId = ?',
             'query_signature': 'abc123',
+            'primary_sql_field': 'batch_text',
+            'dd_tables': ['Customers'],
+            'dd_commands': ['SELECT'],
+            'dd_comments': [],
         }
 
         # Create raw SQL fields
@@ -892,12 +911,20 @@ def test_create_rqt_event(self, mock_agent, query_completion_handler):
         assert rqt_event['db']['raw_query_signature'] == 'def456'
         assert rqt_event['db']['statement'] == 'SELECT * FROM Customers WHERE CustomerId = 123'
 
+        # Verify metadata is present in the RQT event (RQT events already have this structure)
+        assert 'metadata' in rqt_event['db']
+        metadata = rqt_event['db']['metadata']
+        assert metadata['tables'] == ['Customers']
+        assert metadata['commands'] == ['SELECT']
+        assert metadata['comments'] == []
+
         # Verify sqlserver fields
         assert rqt_event['sqlserver']['session_id'] == 123
         assert rqt_event['sqlserver']['xe_type'] == 'sql_batch_completed'
         assert rqt_event['sqlserver']['event_fire_timestamp'] == '2023-01-01T12:00:00.123Z'
         assert rqt_event['sqlserver']['duration_ms'] == 10.0
         assert rqt_event['sqlserver']['query_start'] == '2023-01-01T11:59:50.123Z'
+        assert rqt_event['sqlserver']['primary_sql_field'] == 'batch_text'
 
     def test_create_rqt_event_disabled(self, mock_check, mock_config):
         """Test RQT event creation when disabled"""