instrumentation: make query scanning for dollar quotes a bit more correct

elasticapm/instrumentation/packages/dbapi2.py

@@ -34,6 +34,7 @@
 """
 
 import re
+import string
 
 import wrapt
 
@@ -85,6 +86,7 @@ def scan(tokens):
     literal_started = None
     prev_was_escape = False
     lexeme = []
+    digits = set(string.digits)
 
     i = 0
     while i < len(tokens):
@@ -114,6 +116,11 @@ def scan(tokens):
                 literal_start_idx = i
                 literal_started = token
             elif token == "$":
+                # exclude query parameters that have a digit following the dollar
+                if True and len(tokens) > i + 1 and tokens[i + 1] in digits:
+                    yield i, token
+                    i += 1
+                    continue
                 # Postgres can use arbitrary characters between two $'s as a
                 # literal separation token, e.g.: $fish$ literal $fish$
                 # This part will detect that and skip over the literal.

tests/instrumentation/dbapi2_tests.py

@@ -122,6 +122,20 @@ def test_extract_signature_bytes():
     assert actual == expected
 
 
+def test_extract_signature_pathological():
+    # tune for performance testing
+    multiplier = 10
+    values = []
+    for chunk in range(multiplier):
+        i = chunk * 3
+        values.append(f" (${1+i}::varchar, ${2+i}::varchar, ${3+i}::varchar), ")
+
+    sql = f"SELECT * FROM (VALUES {''.join(values)})\n"
+    actual = extract_signature(sql)
+    expected = "SELECT FROM"
+    assert actual == expected
+
+
 @pytest.mark.parametrize(
     ["sql", "expected"],
     [