SQL Server: recover DBM<>APM correlation comment for non-procedure statements

sqlserver/changelog.d/24899.fixed

@@ -0,0 +1 @@
+Fix DBM<>APM correlation comments being dropped from Query Samples and Query Metrics for non-stored-procedure statements.

sqlserver/datadog_checks/sqlserver/activity.py

@@ -386,35 +386,33 @@ def _obfuscate_and_sanitize_row(self, row):
             comments = statement['metadata'].get('comments', [])
             row['is_proc'] = bool(row.get('procedure_name'))
             has_proc_context = row['is_proc'] or is_statement_proc(row.get('text', ''))[0]
-            if has_proc_context and row.get('text'):
+            if row.get('text') and (has_proc_context or row['text'] != row['statement_text']):
                 try:
-                    procedure_statement = obfuscate_sql_with_metadata(
+                    full_text_statement = obfuscate_sql_with_metadata(
                         row['text'], self._config.obfuscator_options, replace_null_character=True
                     )
-                    row['procedure_signature'] = compute_sql_signature(procedure_statement['query'])
-                    procedure_comments = procedure_statement['metadata'].get('comments', [])
-                    if procedure_comments:
-                        comments = list(set(comments + procedure_comments))
-                    if not row.get('procedure_name'):
-                        procedures = procedure_statement['metadata'].get('procedures')
-                        if procedures:
-                            row['procedure_name'] = procedures[0].lower()
-                            row['is_proc'] = True
+                    comments = self._merge_comments(comments, full_text_statement['metadata'])
+                    if has_proc_context:
+                        row['procedure_signature'] = compute_sql_signature(full_text_statement['query'])
+                        if not row.get('procedure_name'):
+                            procedures = full_text_statement['metadata'].get('procedures')
+                            if procedures:
+                                row['procedure_name'] = procedures[0].lower()
+                                row['is_proc'] = True
                 except Exception as e:
-                    row['procedure_signature'] = '__procedure_obfuscation_error__'
-                    # if we fail to obfuscate the procedure text,
+                    if has_proc_context:
+                        row['procedure_signature'] = '__procedure_obfuscation_error__'
+                    # if we fail to obfuscate the full text,
                     # we should not mark query statement as failed to obfuscate
                     if self._config.log_unobfuscated_queries:
-                        self.log.warning("Failed to obfuscate stored procedure=[%s] | err=[%s]", repr(row['text']), e)
+                        self.log.warning("Failed to obfuscate query text=[%s] | err=[%s]", repr(row['text']), e)
                     else:
-                        self.log.debug("Failed to obfuscate stored procedure | err=[%s]", e)
+                        self.log.debug("Failed to obfuscate query text | err=[%s]", e)
             if 'tail_text' in row:
                 tail_statement = obfuscate_sql_with_metadata(
                     row['tail_text'], self._obfuscator_options_for_tail_text, replace_null_character=True
                 )
-                appended_comments = tail_statement['metadata'].get('comments', [])
-                if appended_comments:
-                    comments = list(set(comments + appended_comments))
+                comments = self._merge_comments(comments, tail_statement['metadata'])
             obfuscated_statement = statement['query']
             metadata = statement['metadata']
             row['dd_commands'] = metadata.get('commands', None)
@@ -437,6 +435,11 @@ def _obfuscate_and_sanitize_row(self, row):
     def _remove_null_vals(row):
         return {key: val for key, val in row.items() if val is not None}
 
+    @staticmethod
+    def _merge_comments(comments, obfuscated_metadata):
+        new_comments = obfuscated_metadata.get('comments', [])
+        return list(dict.fromkeys(comments + new_comments)) if new_comments else comments
+
     @staticmethod
     def _sanitize_row(row, obfuscated_statement=None):
         # rename the statement_text field to 'text' because that

sqlserver/datadog_checks/sqlserver/statements.py

@@ -404,7 +404,6 @@ def _normalize_queries(self, rows):
             if not self._should_include_query_metrics_row(row):
                 continue
             # Attempt to obfuscate SQL statement with metadata
-            procedure_statement = None
             try:
                 statement = obfuscate_sql_with_metadata(
                     row['statement_text'], self._config.obfuscator_options, replace_null_character=True
@@ -433,39 +432,40 @@ def _normalize_queries(self, rows):
             procedure_content = None
             row['is_proc'] = bool(row.get('procedure_name'))
             has_sproc_context = row['is_proc'] or bool(row.get('sproc_object_id'))
-            if (has_sproc_context and row['text']) or self.disable_secondary_tags:
+            needs_procedure_metadata = has_sproc_context or self.disable_secondary_tags
+            if row.get('text') and (needs_procedure_metadata or row['text'] != row['statement_text']):
                 try:
-                    procedure_statement = obfuscate_sql_with_metadata(
+                    full_text_statement = obfuscate_sql_with_metadata(
                         row['text'], self._config.obfuscator_options, replace_null_character=True
                     )
-                    procedure_content = procedure_statement['query']
-                    procedure_signature = compute_sql_signature(procedure_statement['query'])
-                    procedure_comments = procedure_statement['metadata'].get('comments', [])
-                    if procedure_comments:
-                        comments = list(set(comments + procedure_comments))
-                    if not row.get('procedure_name'):
-                        procedures = procedure_statement['metadata'].get('procedures')
-                        if procedures:
-                            row['procedure_name'] = procedures[0].lower()
-                            row['is_proc'] = True
+                    full_text_comments = full_text_statement['metadata'].get('comments', [])
+                    if full_text_comments:
+                        comments = list(dict.fromkeys(comments + full_text_comments))
+                    if needs_procedure_metadata:
+                        procedure_content = full_text_statement['query']
+                        procedure_signature = compute_sql_signature(full_text_statement['query'])
+                        if not row.get('procedure_name'):
+                            procedures = full_text_statement['metadata'].get('procedures')
+                            if procedures:
+                                row['procedure_name'] = procedures[0].lower()
+                                row['is_proc'] = True
                 except Exception as e:
-                    procedure_signature = '__procedure_obfuscation_error__'
-                    procedure_content = '__procedure_obfuscation_error__'
+                    if needs_procedure_metadata:
+                        procedure_signature = '__procedure_obfuscation_error__'
+                        procedure_content = '__procedure_obfuscation_error__'
                     if self._config.log_unobfuscated_queries:
-                        self.log.warning("Failed to obfuscate stored procedure=[%s] | err=[%s]", repr(row['text']), e)
+                        self.log.warning("Failed to obfuscate query text=[%s] | err=[%s]", repr(row['text']), e)
                     else:
                         self.log.debug(
-                            "Failed to obfuscate stored procedure for query_signature=[%s] | err=[%s]",
-                            query_signature,
-                            e,
+                            "Failed to obfuscate query text for query_signature=[%s] | err=[%s]", query_signature, e
                         )
-                    self._check.count(
-                        "dd.sqlserver.statements.error",
-                        1,
-                        **self._check.debug_stats_kwargs(tags=["error:obfuscate-sproc-{}".format(type(e))]),
-                    )
-                    # If we can't obfuscate the stored procedure, we don't need to give up for this row,
-                    # we just won't have the association with the stored procedure in the metrics payload
+                    if needs_procedure_metadata:
+                        self._check.count(
+                            "dd.sqlserver.statements.error",
+                            1,
+                            **self._check.debug_stats_kwargs(tags=["error:obfuscate-sproc-{}".format(type(e))]),
+                        )
+                    # If we can't obfuscate the full text, keep the row using the obfuscated statement text.
 
             if procedure_content:
                 row['procedure_text'] = procedure_content

sqlserver/tests/test_activity.py

@@ -1056,3 +1056,31 @@ def test_sanitize_activity_row(dbm_instance, row):
     row = check.activity._obfuscate_and_sanitize_row(row)
     assert isinstance(row['query_hash'], str)
     assert isinstance(row['query_plan_hash'], str)
+
+
+@pytest.mark.unit
+def test_sanitize_activity_row_recovers_leading_comment_for_non_proc_statement(dbm_instance, datadog_agent):
+    comment = "/*dddbs='orders-service',dde='prod'*/"
+    statement_text = "SELECT * FROM orders WHERE customer_id = @P1"
+    row = {
+        # sp_executesql includes the RPC parameter declaration and leading comment in the full
+        # batch text, but SQL Server's statement offsets exclude both from statement_text.
+        'statement_text': statement_text,
+        'text': f"(@P1 int){comment} {statement_text}",
+        'procedure_name': None,
+        'query_hash': b'\xa4\xffV\x1c\xd4\x14\xbeC',
+        'query_plan_hash': b'\xfe\xba\xbf\xc6_\x9bo\x83',
+    }
+
+    def _obfuscate_sql(sql_query, options=None):
+        comments = [comment] if comment in sql_query else []
+        return json.dumps({'query': sql_query, 'metadata': {'comments': comments}})
+
+    check = SQLServer(CHECK_NAME, {}, [dbm_instance])
+    with mock.patch.object(datadog_agent, 'obfuscate_sql', passthrough=True) as mock_agent:
+        mock_agent.side_effect = _obfuscate_sql
+        row = check.activity._obfuscate_and_sanitize_row(row)
+
+    assert row['dd_comments'] == [comment]
+    assert not row.get('is_proc')
+    assert 'procedure_signature' not in row

sqlserver/tests/test_statements.py

@@ -1236,6 +1236,40 @@ def _obfuscate_sql(sql_query, options=None):
         assert not result_row.get('procedure_text')
 
 
+@pytest.mark.unit
+def test_normalize_queries_recovers_leading_comment_for_non_proc_statement(instance_docker, datadog_agent):
+    instance_docker['dbm'] = True
+    instance_docker['query_metrics'] = {'enabled': True, 'run_sync': True}
+    check = SQLServer(CHECK_NAME, {}, [instance_docker])
+
+    comment = "/*dddbs='orders-service',dde='prod'*/"
+    statement_text = "SELECT * FROM orders WHERE customer_id = @P1"
+    row = {
+        'statement_text': statement_text,
+        'text': f"(@P1 int){comment} {statement_text}",
+        'procedure_name': None,
+        'schema_name': None,
+        'sproc_object_id': None,
+        'query_hash': b'\x01\x02\x03\x04',
+        'query_plan_hash': b'\x05\x06\x07\x08',
+        'plan_handle': b'\x09\x0a\x0b\x0c',
+    }
+
+    def _obfuscate_sql(sql_query, options=None):
+        comments = [comment] if comment in sql_query else []
+        return json.dumps({'query': sql_query, 'metadata': {'comments': comments}})
+
+    with mock.patch.object(datadog_agent, 'obfuscate_sql', passthrough=True) as mock_agent:
+        mock_agent.side_effect = _obfuscate_sql
+        result = check.statement_metrics._normalize_queries([row])
+
+    assert result[0]['dd_comments'] == [comment]
+    assert result[0]['text'] == statement_text
+    assert not result[0]['is_proc']
+    assert 'procedure_signature' not in result[0]
+    assert 'procedure_text' not in result[0]
+
+
 @pytest.mark.flaky
 @pytest.mark.integration
 @pytest.mark.usefixtures('dd_environment')