Repository: DataDog/integrations-core
Pull request #20319: Fix set explain regex
Author: login: jasonmp85
Labels: name: integration/postgres
name: team/agent-integrations
name: team/database-monitoring-agent
name: ecosystems/review-requested
name: product/review-requested
name: agent/review-requested
Created: 2025-05-16T22:28:44Z | Merged: 2025-05-28T06:47:37Z

### What does this PR do?
Reapplies the feature added in #20106, addressing the exponential runtime that certain SQL payloads could cause in the original implementation.

Notably, it removes repetition operators from inner groups in the regex and refactors it so such repetition appears in the outer groupings of inner alternations.

### Testing

To test this, I wired up `sqlsmith` to a Python script which could generate the whole `SET` syntax. Between zero and ten `SET` statements were prepended to every `sqlsmith`-produced statement, as well as inline comments, at random. Over 100,000 statements were processed in this manner, and across many such runs, no exponential behavior was observed with the new regex (the old regex hung regularly, and quickly).

### Performance

One such run of over 200,000 statements is illustrated with the below histogram. The buckets are runtime, in nanoseconds.

```
                      ┌                                        ┐ 
   [    0.0,  5000.0) ┤▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇ 173761   
   [ 5000.0, 10000.0) ┤▇▇▇▇▇▇▇▇ 42857                            
   [10000.0, 15000.0) ┤▇ 6345                                    
   [15000.0, 20000.0) ┤ 1244                                     
   [20000.0, 25000.0) ┤ 341                                      
   [25000.0, 30000.0) ┤ 162                                      
   [30000.0, 35000.0) ┤ 52                                       
   [35000.0, 40000.0) ┤ 22                                       
   [40000.0, 45000.0) ┤ 12                                       
   [45000.0, 50000.0) ┤ 4                                        
   [50000.0, 55000.0) ┤ 4                                        
   [55000.0, 60000.0) ┤ 3                                        
   [60000.0, 65000.0) ┤ 2                                        
   [65000.0, 70000.0) ┤ 0                                        
   [70000.0, 75000.0) ┤ 0                                        
   [75000.0, 80000.0) ┤ 0                                        
   [80000.0, 85000.0) ┤ 0                                        
   [85000.0, 90000.0) ┤ 1                                        
                      └                                        ┘ 
                                      Frequency
```

### Motivation
I'd like to ship this feature and had to pull it before because of it pegging the CPU to 100%.

### Review checklist (to be filled by reviewers)

- [ ] Feature or bugfix MUST have appropriate tests (unit, integration, e2e)
- [ ] Add the `qa/skip-qa` label if the PR doesn't need to be tested during QA.
- [ ] If you need to backport this PR to another branch, you can add the `backport/<branch-name>` label to the PR and it will automatically open a backport PR once this one is merged



Diff:
diff --git a/postgres/changelog.d/20319.fixed b/postgres/changelog.d/20319.fixed
new file mode 100644
index 0000000000000..2db9917b8ed0a
--- /dev/null
+++ b/postgres/changelog.d/20319.fixed
@@ -0,0 +1 @@
+Allow EXPLAIN on multi-statement SQL where one or more SET commands appear before another supported statement type
diff --git a/postgres/datadog_checks/postgres/explain_parameterized_queries.py b/postgres/datadog_checks/postgres/explain_parameterized_queries.py
index 1b6ca39cc86ab..0eea44d762aca 100644
--- a/postgres/datadog_checks/postgres/explain_parameterized_queries.py
+++ b/postgres/datadog_checks/postgres/explain_parameterized_queries.py
@@ -7,7 +7,6 @@
 
 import psycopg2
 
-from datadog_checks.base.utils.db.sql import compute_sql_signature
 from datadog_checks.base.utils.tracking import tracked_method
 from datadog_checks.postgres.cursor import CommenterDictCursor
 
@@ -73,7 +72,7 @@ def __init__(self, check, config, explain_function):
         self._explain_function = explain_function
 
     @tracked_method(agent_check_getter=agent_check_getter)
-    def explain_statement(self, dbname, statement, obfuscated_statement):
+    def explain_statement(self, dbname, statement, obfuscated_statement, query_signature):
         if self._check.version < V12:
             # if pg version < 12, skip explaining parameterized queries because
             # plan_cache_mode is not supported
@@ -85,7 +84,6 @@ def explain_statement(self, dbname, statement, obfuscated_statement):
             return None, DBExplainError.parameterized_query, '{}'.format(type(e))
         self._set_plan_cache_mode(dbname)
 
-        query_signature = compute_sql_signature(obfuscated_statement)
         try:
             self._create_prepared_statement(dbname, statement, obfuscated_statement, query_signature)
         except psycopg2.errors.IndeterminateDatatype as e:
diff --git a/postgres/datadog_checks/postgres/statement_samples.py b/postgres/datadog_checks/postgres/statement_samples.py
index c72bfe524e2ec..f80371fd8409a 100644
--- a/postgres/datadog_checks/postgres/statement_samples.py
+++ b/postgres/datadog_checks/postgres/statement_samples.py
@@ -32,7 +32,7 @@
 from datadog_checks.base.utils.tracking import tracked_method
 from datadog_checks.postgres.explain_parameterized_queries import ExplainParameterizedQueries
 
-from .util import DatabaseConfigurationError, DBExplainError, warning_with_tags
+from .util import DatabaseConfigurationError, DBExplainError, trim_leading_set_stmts, warning_with_tags
 from .version_utils import V9_6, V10
 
 # according to https://unicodebook.readthedocs.io/unicode_encodings.html, the max supported size of a UTF-8 encoded
@@ -749,15 +749,22 @@ def _run_and_track_explain(self, dbname, statement, obfuscated_statement, query_
     @tracked_method(agent_check_getter=agent_check_getter)
     def _run_explain_safe(self, dbname, statement, obfuscated_statement, query_signature):
         # type: (str, str, str, str) -> Tuple[Optional[Dict], Optional[DBExplainError], Optional[str]]
+
+        orig_statement = statement
+
+        # remove leading SET statements from our SQL
+        if obfuscated_statement[:3].lower() == "set":
+            statement = trim_leading_set_stmts(statement)
+            obfuscated_statement = trim_leading_set_stmts(obfuscated_statement)
+
         if not self._can_explain_statement(obfuscated_statement):
             return None, DBExplainError.no_plans_possible, None
 
         track_activity_query_size = self._get_track_activity_query_size()
 
-        if (
-            self._get_truncation_state(track_activity_query_size, statement, query_signature)
-            == StatementTruncationState.truncated
-        ):
+        # truncation check is on the original query, not the trimmed version
+        stmt_trunc = self._get_truncation_state(track_activity_query_size, orig_statement, query_signature)
+        if stmt_trunc == StatementTruncationState.truncated:
             return (
                 None,
                 DBExplainError.query_truncated,
@@ -779,7 +786,7 @@ def _run_explain_safe(self, dbname, statement, obfuscated_statement, query_signa
             if self._explain_parameterized_queries._is_parameterized_query(statement):
                 if is_affirmative(self._config.statement_samples_config.get('explain_parameterized_queries', True)):
                     return self._explain_parameterized_queries.explain_statement(
-                        dbname, statement, obfuscated_statement
+                        dbname, statement, obfuscated_statement, query_signature
                     )
                 e = psycopg2.errors.UndefinedParameter("Unable to explain parameterized query")
                 self._log.debug(
diff --git a/postgres/datadog_checks/postgres/util.py b/postgres/datadog_checks/postgres/util.py
index be6bda6a41f66..42364db3bae4b 100644
--- a/postgres/datadog_checks/postgres/util.py
+++ b/postgres/datadog_checks/postgres/util.py
@@ -1,6 +1,7 @@
 # (C) Datadog, Inc. 2019-present
 # All rights reserved
 # Licensed under Simplified BSD License (see LICENSE)
+import re
 import string
 from enum import Enum
 from typing import Any, List, Tuple  # noqa: F401
@@ -130,6 +131,40 @@ def get_list_chunks(lst, n):
         yield lst[i : i + n]
 
 
+SET_TRIM_PATTERN = re.compile(
+    r"""
+    ^
+    (?:
+        \s*
+        # match one leading comment
+        (?:
+            /\*
+            .*
+            \*/
+            \s*
+        )?
+
+        # match leading SET commands
+        SET\b
+        (?:
+            [^';] | # keywords, integer literals, etc.
+            '[^']*' # single-quoted strings
+        )+
+        ;
+    )+
+    """,
+    flags=(re.I | re.X),
+)
+
+
+# Expects one or more SQL statements in a string. If the string
+# begins with any SET statements, they are removed and the rest
+# of the string is returned. Otherwise, the string is returned
+# as it was received.
+def trim_leading_set_stmts(sql):
+    return SET_TRIM_PATTERN.sub('', sql, 1).lstrip()
+
+
 fmt = PartialFormatter()
 
 AWS_RDS_HOSTNAME_SUFFIX = ".rds.amazonaws.com"
diff --git a/postgres/tests/test_explain_parameterized_queries.py b/postgres/tests/test_explain_parameterized_queries.py
index 977ce53b8b94b..8dec9bd613ecd 100644
--- a/postgres/tests/test_explain_parameterized_queries.py
+++ b/postgres/tests/test_explain_parameterized_queries.py
@@ -51,7 +51,9 @@ def test_explain_parameterized_queries(integration_check, dbm_instance, query, e
     if check.version < V12:
         return
 
-    plan_dict, explain_err_code, err = check.statement_samples._run_and_track_explain(DB_NAME, query, query, query)
+    plan_dict, explain_err_code, err = check.statement_samples._run_and_track_explain(
+        DB_NAME, query, query, "7231596c8b5536d1"
+    )
     assert plan_dict is not None
     assert explain_err_code == expected_explain_err_code
     assert err is None
@@ -111,7 +113,10 @@ def test_explain_parameterized_queries_version_below_12(integration_check, dbm_i
         return
 
     plan_dict, explain_err_code, err = check.statement_samples._run_and_track_explain(
-        DB_NAME, "SELECT * FROM pg_settings WHERE name = $1", "SELECT * FROM pg_settings WHERE name = $1", ""
+        DB_NAME,
+        "SELECT * FROM pg_settings WHERE name = $1",
+        "SELECT * FROM pg_settings WHERE name = $1",
+        "7231596c8b5536d1",
     )
     assert plan_dict is None
     assert explain_err_code == DBExplainError.parameterized_query
@@ -133,7 +138,10 @@ def test_explain_parameterized_queries_create_prepared_statement_exception(integ
         side_effect=psycopg2.errors.DatabaseError("unexpected exception"),
     ):
         plan_dict, explain_err_code, err = check.statement_samples._run_and_track_explain(
-            DB_NAME, "SELECT * FROM pg_settings WHERE name = $1", "SELECT * FROM pg_settings WHERE name = $1", ""
+            DB_NAME,
+            "SELECT * FROM pg_settings WHERE name = $1",
+            "SELECT * FROM pg_settings WHERE name = $1",
+            "7231596c8b5536d1",
         )
         assert plan_dict is None
         assert explain_err_code == DBExplainError.failed_to_explain_with_prepared_statement
@@ -155,7 +163,9 @@ def test_explain_parameterized_queries_explain_prepared_statement_exception(inte
         side_effect=psycopg2.errors.DatabaseError("unexpected exception"),
     ):
         query = "SELECT * FROM pg_settings WHERE name = $1"
-        plan_dict, explain_err_code, err = check.statement_samples._run_and_track_explain(DB_NAME, query, query, "")
+        plan_dict, explain_err_code, err = check.statement_samples._run_and_track_explain(
+            DB_NAME, query, query, "7231596c8b5536d1"
+        )
         assert plan_dict is None
         assert explain_err_code == DBExplainError.failed_to_explain_with_prepared_statement
         assert err is not None
@@ -184,7 +194,10 @@ def test_explain_parameterized_queries_explain_prepared_statement_no_plan_return
         return_value=None,
     ):
         plan_dict, explain_err_code, err = check.statement_samples._run_and_track_explain(
-            DB_NAME, "SELECT * FROM pg_settings WHERE name = $1", "SELECT * FROM pg_settings WHERE name = $1", ""
+            DB_NAME,
+            "SELECT * FROM pg_settings WHERE name = $1",
+            "SELECT * FROM pg_settings WHERE name = $1",
+            "7231596c8b5536d1",
         )
         assert plan_dict is None
         assert explain_err_code == DBExplainError.no_plan_returned_with_prepared_statement
diff --git a/postgres/tests/test_statements.py b/postgres/tests/test_statements.py
index f422deb8f9141..756847edb1a3c 100644
--- a/postgres/tests/test_statements.py
+++ b/postgres/tests/test_statements.py
@@ -583,6 +583,44 @@ def test_get_db_explain_setup_state(integration_check, dbm_instance, dbname, exp
 failed_explain_test_repeat_count = 5
 
 
+@pytest.mark.parametrize(
+    "query",
+    [
+        "SELECT * FROM pg_class",
+        "SET LOCAL datestyle TO postgres; SELECT * FROM pg_class",
+    ],
+)
+def test_successful_explain(
+    integration_check,
+    dbm_instance,
+    aggregator,
+    query,
+):
+    dbname = "datadog_test"
+    # Don't need metrics for this one
+    dbm_instance['query_metrics']['enabled'] = False
+    dbm_instance['query_samples']['explain_parameterized_queries'] = False
+    check = integration_check(dbm_instance)
+    check._connect()
+
+    # run check so all internal state is correctly initialized
+    run_one_check(check)
+
+    # clear out contents of aggregator so we measure only the metrics generated during this specific part of the test
+    aggregator.reset()
+
+    db_explain_error, err = check.statement_samples._get_db_explain_setup_state(dbname)
+    assert db_explain_error is None
+    assert err is None
+
+    plan, *rest = check.statement_samples._run_and_track_explain(dbname, query, query, "7231596c8b5536d1")
+    assert plan is not None
+
+    plan = plan['Plan']
+    assert plan['Node Type'] == 'Seq Scan'
+    assert plan['Relation Name'] == 'pg_class'
+
+
 @pytest.mark.parametrize(
     "query,expected_error_tag,explain_function_override,expected_fail_count,skip_on_versions",
     [
@@ -663,7 +701,7 @@ def test_failed_explain_handling(
     assert err is None
 
     for _ in range(failed_explain_test_repeat_count):
-        check.statement_samples._run_and_track_explain(dbname, query, query, query)
+        check.statement_samples._run_and_track_explain(dbname, query, query, "7231596c8b5536d1")
 
     expected_tags = _get_expected_tags(
         check, dbm_instance, with_host=False, with_db=True, agent_hostname='stubbed.hostname'
@@ -1480,7 +1518,9 @@ def test_statement_run_explain_errors(
     check._connect()
 
     run_one_check(check)
-    _, explain_err_code, err = check.statement_samples._run_and_track_explain("datadog_test", query, query, query)
+    _, explain_err_code, err = check.statement_samples._run_and_track_explain(
+        "datadog_test", query, query, "7231596c8b5536d1"
+    )
     run_one_check(check)
 
     assert explain_err_code == expected_explain_err_code
@@ -1534,7 +1574,9 @@ def test_statement_run_explain_parameterized_queries(
         return
 
     run_one_check(check)
-    _, explain_err_code, err = check.statement_samples._run_and_track_explain("datadog_test", query, query, query)
+    _, explain_err_code, err = check.statement_samples._run_and_track_explain(
+        "datadog_test", query, query, "7231596c8b5536d1"
+    )
     run_one_check(check)
 
     assert explain_err_code == expected_explain_err_code
diff --git a/postgres/tests/test_unit.py b/postgres/tests/test_unit.py
index 51df63a0f74f7..80f4c59020d81 100644
--- a/postgres/tests/test_unit.py
+++ b/postgres/tests/test_unit.py
@@ -201,3 +201,37 @@ def test_database_identifier(pg_instance, template, expected, tags):
     pg_instance['tags'] = tags
     check = PostgreSql('postgres', {}, [pg_instance])
     assert check.database_identifier == expected
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize(
+    "query,expected_trimmed_query",
+    [
+        ("SELECT * FROM pg_settings WHERE name = $1", "SELECT * FROM pg_settings WHERE name = $1"),
+        ("SELECT * FROM pg_settings; DELETE FROM pg_settings;", "SELECT * FROM pg_settings; DELETE FROM pg_settings;"),
+        ("SET search_path TO 'my_schema', public; SELECT * FROM pg_settings", "SELECT * FROM pg_settings"),
+        ("SET TIME ZONE 'Europe/Rome'; SELECT * FROM pg_settings", "SELECT * FROM pg_settings"),
+        (
+            "SET LOCAL request_id = 1234; SET LOCAL hostname TO 'Bob''s Laptop'; SELECT * FROM pg_settings",
+            "SELECT * FROM pg_settings",
+        ),
+        ("SET LONG;" * 1024 + "SELECT *;", "SELECT *;"),
+        ("SET " + "'quotable'" * 1024 + "; SELECT *;", "SELECT *;"),
+        ("SET 'l" + "o" * 1024 + "ng'; SELECT *;", "SELECT *;"),
+        (" /** pl/pgsql **/ SET 'comment'; SELECT *;", "SELECT *;"),
+        ("this isn't SQL", "this isn't SQL"),
+        (
+            "SET SESSION min_wal_size = 14400; "
+            + "SET LOCAL wal_buffers TO 2048; "
+            + "/* testing id 1234 */ set send_abort_for_kill TO 'stderr'; "
+            + "set id = case when (false) and ((((cast(null as box) ~= cast(null as box)) "
+            + "or (cast(null as point) <@ cast(null as line))) or (public.my table",
+            "set id = case when (false) and ((((cast(null as box) ~= cast(null as box)) "
+            + "or (cast(null as point) <@ cast(null as line))) or (public.my table",
+        ),
+        ("", ""),
+    ],
+)
+def test_trim_set_stmts(query, expected_trimmed_query):
+    trimmed_query = util.trim_leading_set_stmts(query)
+    assert trimmed_query == expected_trimmed_query


Reviews:
user: login: nenadnoveljic
state: CHANGES_REQUESTED
author_association: MEMBER
submitted_at: 2025-05-19T08:31:45Z
user: login: jasonmp85
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T16:17:46Z
user: login: jasonmp85
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T16:18:34Z
user: login: jasonmp85
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T21:59:45Z
user: login: nenadnoveljic
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-20T07:03:32Z
user: login: jasonmp85
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-21T18:37:04Z
user: login: nenadnoveljic
state: DISMISSED
author_association: MEMBER
submitted_at: 2025-05-21T20:02:52Z
user: login: nenadnoveljic
state: APPROVED
author_association: MEMBER
submitted_at: 2025-05-28T06:46:46Z

Review comments:
diff_hunk: @@ -749,15 +749,22 @@ def _run_and_track_explain(self, dbname, statement, obfuscated_statement, query_
     @tracked_method(agent_check_getter=agent_check_getter)
     def _run_explain_safe(self, dbname, statement, obfuscated_statement, query_signature):
         # type: (str, str, str, str) -> Tuple[Optional[Dict], Optional[DBExplainError], Optional[str]]
+
+        orig_statement = statement
+
+        # remove leading SET statements from our SQL
+        if obfuscated_statement[:3].lower() == "set":
+            statement = trim_leading_set_stmts(statement)
+            obfuscated_statement = trim_leading_set_stmts(obfuscated_statement)
path: postgres/datadog_checks/postgres/statement_samples.py
user: login: nenadnoveljic
body: We compute the query signature from a trimmed statement [here](https://github.com/DataDog/integrations-core/blob/master/postgres/datadog_checks/postgres/explain_parameterized_queries.py#L88), and send that with the plan event. But could this cause issues matching it with query samples and metrics? Since those use the untrimmed statement, the query won't match.
created_at: 2025-05-19T08:29:38Z
author_association: MEMBER
diff_hunk: @@ -130,6 +131,40 @@ def get_list_chunks(lst, n):
         yield lst[i : i + n]
 
 
+SET_TRIM_PATTERN = re.compile(
+    r"""
+    ^
+    (?:
+        \s*
+        # match one leading comment
+        (?:
+            /\*
+            .*
+            \*/
+            \s*
+        )?
+
+        # match leading SET commands
+        SET\b
+        (?:
+            [^';] | # keywords, integer literals, etc.
+            '[^']*' # single-quoted strings
+        )+
+        ;
+    )+
+    """,
+    flags=(re.I | re.X),
+)
+
+
+# Expects one or more SQL statements in a string. If the string
+# begins with any SET statements, they are removed and the rest
+# of the string is returned. Otherwise, the string is returned
+# as it was received.
+def trim_leading_set_stmts(sql):
+    return SET_TRIM_PATTERN.sub('', sql, 1).lstrip()
path: postgres/datadog_checks/postgres/util.py
user: login: nenadnoveljic
body: Should we try to remove the regex entirely?
created_at: 2025-05-19T08:31:13Z
author_association: MEMBER
diff_hunk: @@ -130,6 +131,40 @@ def get_list_chunks(lst, n):
         yield lst[i : i + n]
 
 
+SET_TRIM_PATTERN = re.compile(
+    r"""
+    ^
+    (?:
+        \s*
+        # match one leading comment
+        (?:
+            /\*
+            .*
+            \*/
+            \s*
+        )?
+
+        # match leading SET commands
+        SET\b
+        (?:
+            [^';] | # keywords, integer literals, etc.
+            '[^']*' # single-quoted strings
+        )+
+        ;
+    )+
+    """,
+    flags=(re.I | re.X),
+)
+
+
+# Expects one or more SQL statements in a string. If the string
+# begins with any SET statements, they are removed and the rest
+# of the string is returned. Otherwise, the string is returned
+# as it was received.
+def trim_leading_set_stmts(sql):
+    return SET_TRIM_PATTERN.sub('', sql, 1).lstrip()
path: postgres/datadog_checks/postgres/util.py
user: login: jasonmp85
body: I'd end up writing an iterative lexer/parser and in that case would require several mini-regexes to match lexemes, in additional several state variables to keep track of whether I'm in a comment, in a quote, etc., etc.…

I think it would _drastically_ hurt readability and I don't think there's any benefit at this point.
created_at: 2025-05-19T16:17:46Z
author_association: CONTRIBUTOR
diff_hunk: @@ -749,15 +749,22 @@ def _run_and_track_explain(self, dbname, statement, obfuscated_statement, query_
     @tracked_method(agent_check_getter=agent_check_getter)
     def _run_explain_safe(self, dbname, statement, obfuscated_statement, query_signature):
         # type: (str, str, str, str) -> Tuple[Optional[Dict], Optional[DBExplainError], Optional[str]]
+
+        orig_statement = statement
+
+        # remove leading SET statements from our SQL
+        if obfuscated_statement[:3].lower() == "set":
+            statement = trim_leading_set_stmts(statement)
+            obfuscated_statement = trim_leading_set_stmts(obfuscated_statement)
path: postgres/datadog_checks/postgres/statement_samples.py
user: login: jasonmp85
body: Looking into this, I thought I vetted this logic on the first PR, but it's possible I missed something…
created_at: 2025-05-19T16:18:34Z
author_association: CONTRIBUTOR
diff_hunk: @@ -749,15 +749,22 @@ def _run_and_track_explain(self, dbname, statement, obfuscated_statement, query_
     @tracked_method(agent_check_getter=agent_check_getter)
     def _run_explain_safe(self, dbname, statement, obfuscated_statement, query_signature):
         # type: (str, str, str, str) -> Tuple[Optional[Dict], Optional[DBExplainError], Optional[str]]
+
+        orig_statement = statement
+
+        # remove leading SET statements from our SQL
+        if obfuscated_statement[:3].lower() == "set":
+            statement = trim_leading_set_stmts(statement)
+            obfuscated_statement = trim_leading_set_stmts(obfuscated_statement)
path: postgres/datadog_checks/postgres/statement_samples.py
user: login: jasonmp85
body: Ok, looking back into this, I don't see a problem with this code… the query signature generated within `explain_statement` is used solely to give a unique identifier to the prepared statement created for the purposes of `EXPLAIN`. It isn't sent on anywhere, or otherwise used in a metric calculation. As you mention, the original query has a signature calculated and included in its row quite a bit further up the stack, which is unchanged.

It's unclear and tortuous, but my trimming the SQL doesn't particularly add to that. All that happens due to this change is certain statements that previously had no plan passed up the stack will now have a plan passed up from `_run_and_track_explain`. The caller (`_collect_plan_for_statement`) still uses the query signature calculated much earlier, pre-trim. The trimming only occurs at the lower part of the call stack (in `_run_explain_safe` and where it calls `_explain_parameterized_queries` and `_run_explain`). As you point out, they may call `explain_statement`, which will calculate a different query signature, but as far as I can see, it's only used in the name of the prepared statement created (and removed) during the course of an `EXPLAIN`.
created_at: 2025-05-19T21:59:45Z
author_association: CONTRIBUTOR
diff_hunk: @@ -749,15 +749,22 @@ def _run_and_track_explain(self, dbname, statement, obfuscated_statement, query_
     @tracked_method(agent_check_getter=agent_check_getter)
     def _run_explain_safe(self, dbname, statement, obfuscated_statement, query_signature):
         # type: (str, str, str, str) -> Tuple[Optional[Dict], Optional[DBExplainError], Optional[str]]
+
+        orig_statement = statement
+
+        # remove leading SET statements from our SQL
+        if obfuscated_statement[:3].lower() == "set":
+            statement = trim_leading_set_stmts(statement)
+            obfuscated_statement = trim_leading_set_stmts(obfuscated_statement)
path: postgres/datadog_checks/postgres/statement_samples.py
user: login: nenadnoveljic
body: Since this trimmed `query_signature` does show up in logs, for example [here](https://github.com/DataDog/integrations-core/blob/master/postgres/datadog_checks/postgres/explain_parameterized_queries.py#L132-L133), but isn’t present in any emitted telemetry (metrics, samples, plan events), could that be misleading during troubleshooting? For example, someone debugging a plan issue might search for that signature elsewhere and find nothing, since it was generated from the trimmed version.
created_at: 2025-05-20T07:03:32Z
author_association: MEMBER
diff_hunk: @@ -749,15 +749,22 @@ def _run_and_track_explain(self, dbname, statement, obfuscated_statement, query_
     @tracked_method(agent_check_getter=agent_check_getter)
     def _run_explain_safe(self, dbname, statement, obfuscated_statement, query_signature):
         # type: (str, str, str, str) -> Tuple[Optional[Dict], Optional[DBExplainError], Optional[str]]
+
+        orig_statement = statement
+
+        # remove leading SET statements from our SQL
+        if obfuscated_statement[:3].lower() == "set":
+            statement = trim_leading_set_stmts(statement)
+            obfuscated_statement = trim_leading_set_stmts(obfuscated_statement)
path: postgres/datadog_checks/postgres/statement_samples.py
user: login: jasonmp85
body: That's a good point, I'm still getting a handle on Datadog the user-facing product, and didn't think about that difference to the end user… but I guess since the agent runs on user infra, they will see logs from the agent?

I've changed the code to pass down the `query_signature`. It isn't clear to my why `_run_explain` ever recalculated it, given it's available to the caller and is solely used to name this prepared statement.

Does that address this?
created_at: 2025-05-21T18:37:03Z
author_association: CONTRIBUTOR

PR comments:
user: login: jasonmp85
created_at: 2025-05-16T22:30:54Z
body: The performance benchmark was around the call to trim the `SET` statements from the SQL using the new regex. The vast majority of calls run in 5–10 microseconds on my laptop.

Additionally, after determining a snippet of affected SQL using the fuzzer, I added it to our unit tests. Without the fix, it hangs indefinitely. With the fix the whole suite runs imperceptibly fast.
author_association: CONTRIBUTOR
user: login: codecov[bot]
created_at: 2025-05-16T23:27:44Z
body: ## [Codecov](https://app.codecov.io/gh/DataDog/integrations-core/pull/20319?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=DataDog) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 90.33%. Comparing base [(`0245920`)](https://app.codecov.io/gh/DataDog/integrations-core/commit/024592040086a7901ed3614a93c0a9355b313d6f?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=DataDog) to head [(`43636d8`)](https://app.codecov.io/gh/DataDog/integrations-core/commit/43636d885e37f96946f81b001ca7bb2a13972330?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=DataDog).
> Report is 1 commits behind head on master.

<details><summary>Additional details and impacted files</summary>


| [Flag](https://app.codecov.io/gh/DataDog/integrations-core/pull/20319/flags?src=pr&el=flags&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=DataDog) | Coverage Δ | |
|---|---|---|
| [activemq](https://app.codecov.io/gh/DataDog/integrations-core/pull/20319/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=DataDog) | `?` | |
| [cassandra](https://app.codecov.io/gh/DataDog/integrations-core/pull/20319/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=DataDog) | `?` | |
| [confluent_platform](https://app.codecov.io/gh/DataDog/integrations-core/pull/20319/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=DataDog) | `?` | |
| [hive](https://app.codecov.io/gh/DataDog/integrations-core/pull/20319/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=DataDog) | `?` | |
| [hivemq](https://app.codecov.io/gh/DataDog/integrations-core/pull/20319/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=DataDog) | `?` | |
| [hudi](https://app.codecov.io/gh/DataDog/integrations-core/pull/20319/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=DataDog) | `?` | |
| [ignite](https://app.codecov.io/gh/DataDog/integrations-core/pull/20319/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=DataDog) | `?` | |
| [jboss_wildfly](https://app.codecov.io/gh/DataDog/integrations-core/pull/20319/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=DataDog) | `?` | |
| [kafka](https://app.codecov.io/gh/DataDog/integrations-core/pull/20319/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=DataDog) | `?` | |
| [postgres](https://app.codecov.io/gh/DataDog/integrations-core/pull/20319/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=DataDog) | `93.22% <100.00%> (+3.62%)` | :arrow_up: |
| [presto](https://app.codecov.io/gh/DataDog/integrations-core/pull/20319/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=DataDog) | `?` | |
| [solr](https://app.codecov.io/gh/DataDog/integrations-core/pull/20319/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=DataDog) | `?` | |
| [tomcat](https://app.codecov.io/gh/DataDog/integrations-core/pull/20319/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=DataDog) | `?` | |
| [weblogic](https://app.codecov.io/gh/DataDog/integrations-core/pull/20319/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=DataDog) | `?` | |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=DataDog#carryforward-flags-in-the-pull-request-comment) to find out more.

</details>
<details><summary> :rocket: New features to boost your workflow: </summary>

- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
- :package: [JS Bundle Analysis](https://docs.codecov.com/docs/javascript-bundle-analysis): Save yourself from yourself by tracking and limiting bundle sizes in JS merges.
</details>
author_association: NONE
user: login: jasonmp85
created_at: 2025-05-27T22:09:07Z
body: @nenadnoveljic I thought this auto-merged but something happened with a test. Between being sick and the US holiday, your review became stale… can you re-post approval when you have the chance?
author_association: CONTRIBUTOR
user: login: jasonmp85
created_at: 2025-05-28T16:19:16Z
body: 🙇🏻 
author_association: CONTRIBUTOR