Revert "fix: allow EXPLAIN on multi-statement SQL beginning with SET (#20106)" (#20282)

This reverts commit 9b43811.

Author: jasonmp85

postgres/changelog.d/20106.fixed

@@ -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, trim_leading_set_stmts, warning_with_tags
+from .util import DatabaseConfigurationError, DBExplainError, 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,22 +749,15 @@ 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()
 
-        # 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:
+        if (
+            self._get_truncation_state(track_activity_query_size, statement, query_signature)
+            == StatementTruncationState.truncated
+        ):
             return (
                 None,
                 DBExplainError.query_truncated,

postgres/datadog_checks/postgres/statement_samples.py

@@ -1,7 +1,6 @@
 # (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
@@ -131,43 +130,6 @@ def get_list_chunks(lst, n):
         yield lst[i : i + n]
 
 
-SET_TRIM_PATTERN = re.compile(
-    r"""
-    ^(?:
-        # match one leading comment
-        (?:
-            \s*
-            /\*
-            .*?
-            \*/
-        )?
-        # match leading SET commands
-        \s*SET\b
-        (?:
-            [^';]*? | # keywords, integer literals, etc.
-            (?:'[^']*?')* # single-quoted strings
-        )+
-        ;
-    )+
-    \s*(.+?)$ # actual non-SET cmds
-    """,
-    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):
-    match = SET_TRIM_PATTERN.match(sql)
-
-    if match:
-        return match.group(1)
-    else:
-        return sql
-
-
 fmt = PartialFormatter()
 
 AWS_RDS_HOSTNAME_SUFFIX = ".rds.amazonaws.com"

postgres/datadog_checks/postgres/util.py

@@ -583,44 +583,6 @@ 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, query)
-    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",
     [

postgres/tests/test_statements.py

@@ -201,28 +201,3 @@ 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"),
-        ("", ""),
-    ],
-)
-def test_trim_set_stmts(query, expected_trimmed_query):
-    trimmed_query = util.trim_leading_set_stmts(query)
-    assert trimmed_query == expected_trimmed_query