Recover appended DBM comments from SQL Server activity tail text

sqlserver/changelog.d/25884.fixed

@@ -0,0 +1 @@
+Fix Query Samples missing the calling APM service for queries longer than `stored_procedure_characters_limit` with an appended SQL comment.

sqlserver/datadog_checks/sqlserver/activity.py

@@ -8,7 +8,6 @@
 import time
 
 from datadog_checks.base import is_affirmative
-from datadog_checks.base.utils.common import to_native_string
 from datadog_checks.base.utils.db.sql import compute_sql_signature
 from datadog_checks.base.utils.db.utils import (
     DBMAsyncJob,
@@ -20,7 +19,12 @@
 from datadog_checks.base.utils.tracking import tracked_method
 from datadog_checks.sqlserver.config import SQLServerConfig
 from datadog_checks.sqlserver.const import STATIC_INFO_ENGINE_EDITION, STATIC_INFO_VERSION
-from datadog_checks.sqlserver.utils import is_statement_proc, needs_comment_recovery, raise_if_cancelled
+from datadog_checks.sqlserver.utils import (
+    extract_trailing_dbm_comments,
+    is_statement_proc,
+    needs_comment_recovery,
+    raise_if_cancelled,
+)
 
 DEFAULT_COLLECTION_INTERVAL = 10
 MAX_PAYLOAD_BYTES = 19e6
@@ -175,14 +179,9 @@ class SqlserverActivity(DBMAsyncJob):
     def __init__(self, check, config: SQLServerConfig):
         self.log = check.log
         self._config = config
-        self._obfuscator_options_for_tail_text = to_native_string(
-            json.dumps(
-                {
-                    'dbms': 'mssql',
-                    'obfuscation_mode': 'normalize_only',  # only need the trailing comments from the tail text
-                    'collect_procedures': True,
-                }
-            )
+        obfuscator_options = json.loads(self._config.obfuscator_options)
+        self._collect_comments = bool(
+            obfuscator_options.get('return_json_metadata') and obfuscator_options.get('collect_comments')
         )
         collection_interval = float(
             self._config.activity_config.get('collection_interval', DEFAULT_COLLECTION_INTERVAL)
@@ -407,11 +406,10 @@ def _obfuscate_and_sanitize_row(self, row):
                         self.log.warning("Failed to obfuscate query text=[%s] | err=[%s]", repr(row['text']), e)
                     else:
                         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
-                )
-                comments = self._merge_comments(comments, tail_statement['metadata'])
+            if self._collect_comments and row.get('tail_text'):
+                # tail_text is cut at a fixed length and can start inside a string literal, which hides
+                # the trailing comments from the obfuscator, so read them directly from the end instead.
+                comments = self._merge_comments(comments, {'comments': extract_trailing_dbm_comments(row['tail_text'])})
             obfuscated_statement = statement['query']
             metadata = statement['metadata']
             row['dd_commands'] = metadata.get('commands', None)
@@ -450,8 +448,8 @@ def _sanitize_row(row, obfuscated_statement=None):
         if 'query_plan_hash' in row:
             row['query_plan_hash'] = _hash_to_hex(row['query_plan_hash'])
         # remove deobfuscated sql text from event
-        if 'statement_text' in row:
-            del row['statement_text']
+        row.pop('statement_text', None)
+        row.pop('tail_text', None)
         # convert `context_info` to string, since the default decoding format
         # for json is utf-8, and `context_info` need not comply to utf-8
         if 'context_info' in row:

sqlserver/datadog_checks/sqlserver/utils.py

@@ -164,6 +164,35 @@ def needs_comment_recovery(full_text: Optional[str], statement_comments: Optiona
     )
 
 
+_SQLCOMMENTER_COMMENT = re.compile(r"/\*\s*[\w.-]+='[^']*'(?:\s*,\s*[\w.-]+='[^']*')*\s*\*/")
+
+
+def extract_trailing_dbm_comments(text: Optional[str]) -> list[str]:
+    """Return the DBM propagation comments at the end of `text`, in order.
+
+    Scans backward over trailing whitespace, semicolons and `/* ... */` blocks, so it is safe on text cut
+    at an arbitrary position, such as inside a string literal. A backward scan cannot tell whether a
+    delimiter belongs to a string literal or a line comment, so only blocks in the `key='value'` format
+    the tracers emit, with a DBM marker, are accepted, and the scan stops at the first block that is not.
+    """
+    comments = []
+    end = len(text) if text else 0
+    while True:
+        while end > 0 and (text[end - 1].isspace() or text[end - 1] == ';'):
+            end -= 1
+        if end < 4 or text[end - 2 : end] != '*/':
+            break
+        start = text.rfind('/*', 0, end - 2)
+        if start < 0:
+            break
+        comment = text[start:end]
+        if not _SQLCOMMENTER_COMMENT.fullmatch(comment) or not has_dbm_comment_marker(comment):
+            break
+        comments.append(comment)
+        end = start
+    return comments[::-1]
+
+
 def _get_procedure_name(t, idx):
     if idx >= 0 and idx + 1 < len(t):
         return t[idx + 1].lower()

sqlserver/tests/test_activity.py

@@ -1072,6 +1072,41 @@ def _obfuscate_sql(sql_query, options=None):
     assert mock_agent.call_count == 2
 
 
+@pytest.mark.unit
+@pytest.mark.parametrize(
+    'obfuscator_options,expect_comments',
+    [
+        pytest.param({}, True, id='defaults'),
+        pytest.param({'collect_comments': False}, False, id='collect_comments_disabled'),
+        pytest.param({'collect_metadata': False}, False, id='collect_metadata_disabled'),
+    ],
+)
+def test_sanitize_activity_row_recovers_appended_comment_from_tail_text(
+    dbm_instance, datadog_agent, obfuscator_options, expect_comments
+):
+    case_comment = "/* orders-lookup */"
+    dbm_comment = "/*ddps='orders-service',dde='prod',ddpv='1.0.0'*/"
+    statement_text = "SELECT id, '" + 'x' * 600 + "' AS pad FROM dbo.orders"
+    batch = f"(@P0 nvarchar(4000)){statement_text} {case_comment} {dbm_comment}"
+    row = {
+        'statement_text': statement_text,
+        'text': batch[:500],
+        # RIGHT(text, 200) can start inside a string literal
+        'tail_text': batch[-200:],
+        'procedure_name': None,
+        'query_hash': b'\xa4\xffV\x1c\xd4\x14\xbeC',
+        'query_plan_hash': b'\xfe\xba\xbf\xc6_\x9bo\x83',
+    }
+    assert row['tail_text'].startswith('x')
+
+    check = SQLServer(CHECK_NAME, {}, [dict(dbm_instance, obfuscator_options=obfuscator_options)])
+    row = check.activity._obfuscate_and_sanitize_row(row)
+
+    assert row['dd_comments'] == ([dbm_comment] if expect_comments else [])
+    assert 'tail_text' not in row
+    assert 'statement_text' not in row
+
+
 @pytest.mark.unit
 @pytest.mark.parametrize(
     'text,statement_text,expected_comments',

sqlserver/tests/test_unit.py

@@ -47,6 +47,7 @@
     Database,
     construct_use_statement,
     extract_sql_comments_and_procedure_name,
+    extract_trailing_dbm_comments,
     get_unixodbc_sysconfig,
     needs_comment_recovery,
     parse_sqlserver_major_version,
@@ -1569,6 +1570,35 @@ def test_needs_comment_recovery_requires_missing_dbm_comment(full_text, statemen
     assert needs_comment_recovery(full_text, statement_comments) is expected
 
 
+@pytest.mark.parametrize(
+    'text,expected',
+    [
+        pytest.param(None, [], id='none'),
+        pytest.param('', [], id='empty'),
+        pytest.param('SELECT 1 FROM t', [], id='no_comment'),
+        pytest.param("SELECT 1 /*ddps='svc'*/", ["/*ddps='svc'*/"], id='single'),
+        pytest.param(
+            "SELECT 1 /*ddps='svc',dde='prod'*/ /*traceparent='00-abc-def-01'*/ ;\n",
+            ["/*ddps='svc',dde='prod'*/", "/*traceparent='00-abc-def-01'*/"],
+            id='multiple_with_trailing_semicolon_and_whitespace',
+        ),
+        pytest.param(
+            "xxxx' AS pad FROM dbo.t /* case */ /*ddps='svc',ddh='db%2C1433'*/",
+            ["/*ddps='svc',ddh='db%2C1433'*/"],
+            id='starts_inside_string_literal_stops_at_non_dbm_comment',
+        ),
+        pytest.param("SELECT 1 /*ddps='svc'*/ /* case */", [], id='non_dbm_comment_last'),
+        pytest.param("/* leading */ SELECT 1 FROM t", [], id='leading_comment_only'),
+        pytest.param("SELECT '/* not a comment */' AS c", [], id='comment_inside_trailing_literal'),
+        pytest.param("SELECT '/*customer-secret' -- */", [], id='delimiters_from_literal_and_line_comment'),
+        pytest.param("SELECT '/*ddps=''x' -- */", [], id='marker_inside_literal_and_line_comment'),
+        pytest.param("ps='svc',dde='prod'*/", [], id='starts_inside_comment'),
+    ],
+)
+def test_extract_trailing_dbm_comments(text, expected):
+    assert extract_trailing_dbm_comments(text) == expected
+
+
 @pytest.mark.parametrize(
     "instance_host,split_host,split_port",
     [