fix(dbm): Make pg_plan comes first before dbm_comment

ddtrace/propagation/_database_monitoring.py

@@ -1,3 +1,4 @@
+import re
 from typing import TYPE_CHECKING  # noqa:F401
 from typing import Literal  # noqa:F401
 from typing import Union  # noqa:F401
@@ -38,6 +39,9 @@
 
 log = get_logger(__name__)
 
+_LEADING_OPTIMIZER_HINT = re.compile(r"\A[^\S\r\n]*/\*\+.*?\*/\s*", re.DOTALL)
+_LEADING_OPTIMIZER_HINT_BYTES = re.compile(rb"\A[^\S\r\n]*/\*\+.*?\*/\s*", re.DOTALL)
+
 
 def _should_inject_sql_basehash():
     # type: () -> bool
@@ -49,9 +53,17 @@ def _should_inject_sql_basehash():
 def default_sql_injector(dbm_comment, sql_statement):
     # type: (str, Union[str, bytes]) -> Union[str, bytes]
     try:
+        # pg_hint_plan only reads the first block comment. Keep a hint
+        # that starts on the first line before the injected DBM comment.
         if isinstance(sql_statement, bytes):
-            return dbm_comment.encode("utf-8", errors="strict") + sql_statement
-        return dbm_comment + sql_statement
+            byte_hint_match = _LEADING_OPTIMIZER_HINT_BYTES.match(sql_statement)
+            byte_hint = byte_hint_match.group() if byte_hint_match else b""
+            byte_remaining_sql = sql_statement[len(byte_hint) :]
+            return byte_hint + dbm_comment.encode("utf-8", errors="strict") + byte_remaining_sql
+        string_hint_match = _LEADING_OPTIMIZER_HINT.match(sql_statement)
+        string_hint = string_hint_match.group() if string_hint_match else ""
+        string_remaining_sql = sql_statement[len(string_hint) :]
+        return string_hint + dbm_comment + string_remaining_sql
     except (TypeError, ValueError):
         log.warning(
             "Linking Database Monitoring profiles to spans is not supported for the following query type: %s. "

tests/internal/test_database_monitoring.py

@@ -121,6 +121,11 @@ def test_dbm_propagation_full_mode():
         # assert that args remain unchanged
         assert new_args == args
 
+        hinted_query = "/*+ Leading((a b)) */ SELECT * from table;"
+        _, hinted_kwargs = dbm_popagator.inject(dbspan, tuple(), {"procedure": hinted_query})
+        assert hinted_kwargs == {"procedure": "/*+ Leading((a b)) */ " + sqlcomment + "SELECT * from table;"}
+        assert "traceparent=" in hinted_kwargs["procedure"]
+
         # ensure that dbm tag is set (required for full mode)
         assert dbspan.get_tag(_database_monitoring.DBM_TRACE_INJECTED_TAG) == "true"
 
@@ -505,3 +510,56 @@ def test_default_sql_injector(caplog):
         assert result == non_string_object
 
     assert "Linking Database Monitoring profiles to spans is not supported for the following query type:" in caplog.text
+
+
+@pytest.mark.parametrize(
+    "query, expected",
+    [
+        (
+            "/*+ Leading((a b)) */ SELECT 1",
+            "/*+ Leading((a b)) */ /*dddbs='orders-db'*/ SELECT 1",
+        ),
+        (
+            " \t/*+ SeqScan(items) */\tSELECT * FROM items",
+            " \t/*+ SeqScan(items) */\t/*dddbs='orders-db'*/ SELECT * FROM items",
+        ),
+        (
+            b"/*+ Leading((a b)) */ SELECT 1",
+            b"/*+ Leading((a b)) */ /*dddbs='orders-db'*/ SELECT 1",
+        ),
+        (
+            b" \t/*+ SeqScan(items) */\tSELECT * FROM items",
+            b" \t/*+ SeqScan(items) */\t/*dddbs='orders-db'*/ SELECT * FROM items",
+        ),
+        (
+            "/*+ SeqScan(items) */SELECT * FROM items",
+            "/*+ SeqScan(items) *//*dddbs='orders-db'*/ SELECT * FROM items",
+        ),
+        (
+            b"/*+ SeqScan(items) */SELECT * FROM items",
+            b"/*+ SeqScan(items) *//*dddbs='orders-db'*/ SELECT * FROM items",
+        ),
+    ],
+)
+def test_default_sql_injector_preserves_leading_optimizer_hint(query, expected):
+    assert default_sql_injector("/*dddbs='orders-db'*/ ", query) == expected
+
+
+@pytest.mark.parametrize(
+    "query",
+    [
+        "SELECT '/*+ not a hint */'",
+        " \n/*+ Leading((a b)) */ SELECT 1",
+        "/* ordinary comment */ /*+ Leading((a b)) */ SELECT 1",
+        "/*+ unterminated hint",
+        b"SELECT '/*+ not a hint */'",
+        b" \n/*+ Leading((a b)) */ SELECT 1",
+        b"/* ordinary comment */ /*+ Leading((a b)) */ SELECT 1",
+        b"/*+ unterminated hint",
+    ],
+)
+def test_default_sql_injector_only_preserves_complete_leading_optimizer_hint(query):
+    dbm_comment = "/*dddbs='orders-db'*/ "
+    expected_comment = dbm_comment.encode() if isinstance(query, bytes) else dbm_comment
+
+    assert default_sql_injector(dbm_comment, query) == expected_comment + query