diff --git a/sqlserver/changelog.d/24899.fixed b/sqlserver/changelog.d/24899.fixed new file mode 100644 index 0000000000000..ba34ea998509a --- /dev/null +++ b/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. diff --git a/sqlserver/datadog_checks/sqlserver/activity.py b/sqlserver/datadog_checks/sqlserver/activity.py index 95ee1540fe483..d16913d374777 100644 --- a/sqlserver/datadog_checks/sqlserver/activity.py +++ b/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 diff --git a/sqlserver/datadog_checks/sqlserver/statements.py b/sqlserver/datadog_checks/sqlserver/statements.py index d4a9d42cd310f..7b8e215a8ba27 100644 --- a/sqlserver/datadog_checks/sqlserver/statements.py +++ b/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 diff --git a/sqlserver/tests/test_activity.py b/sqlserver/tests/test_activity.py index f8cae9cc07757..2347fb9e4daa8 100644 --- a/sqlserver/tests/test_activity.py +++ b/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 diff --git a/sqlserver/tests/test_statements.py b/sqlserver/tests/test_statements.py index 4c9715dc6e029..73fbe9276f593 100644 --- a/sqlserver/tests/test_statements.py +++ b/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')