Fix pgerror reference (DataDog#20848)

* Fix pgerror reference

* Changelog

* Create real error for pg_stat_statements

* Restore extension

Author: sethsamuel

postgres/datadog_checks/postgres/statements.py

@@ -426,9 +426,7 @@ def _load_pg_stat_statements(self):
         except psycopg.Error as e:
             error_tag = "error:database-{}".format(type(e).__name__)
 
-            if (
-                isinstance(e, psycopg.errors.ObjectNotInPrerequisiteState)
-            ) and 'pg_stat_statements must be loaded' in str(e.pgerror):
+            if (isinstance(e, psycopg.errors.ObjectNotInPrerequisiteState)) and 'pg_stat_statements' in str(e):
                 error_tag = "error:database-{}-pg_stat_statements_not_loaded".format(type(e).__name__)
                 self._check.record_warning(
                     DatabaseConfigurationError.pg_stat_statements_not_loaded,
@@ -444,7 +442,7 @@ def _load_pg_stat_statements(self):
                         code=DatabaseConfigurationError.pg_stat_statements_not_loaded.value,
                     ),
                 )
-            elif isinstance(e, psycopg.errors.UndefinedTable) and 'pg_stat_statements' in str(e.pgerror):
+            elif isinstance(e, psycopg.errors.UndefinedTable) and 'pg_stat_statements' in str(e):
                 error_tag = "error:database-{}-pg_stat_statements_not_created".format(type(e).__name__)
                 self._check.record_warning(
                     DatabaseConfigurationError.pg_stat_statements_not_created,

postgres/tests/test_statements.py

@@ -1888,79 +1888,9 @@ def test_statement_samples_config_invalid_number(integration_check, pg_instance,
         integration_check(pg_instance)
 
 
-class ObjectNotInPrerequisiteState(psycopg.errors.ObjectNotInPrerequisiteState):
-    """
-    A fake ObjectNotInPrerequisiteState that allows setting pg_error on construction since ObjectNotInPrerequisiteState
-    has it as read-only and not settable at construction-time
-    """
-
-    def __init__(self, pg_error):
-        self.pg_error = pg_error
-
-    def __getattribute__(self, attr):
-        if attr == 'pgerror':
-            return self.pg_error
-        else:
-            return super(ObjectNotInPrerequisiteState, self).__getattribute__(attr)
-
-    def __str__(self):
-        return self.pg_error
-
-
-class UndefinedTable(psycopg.errors.UndefinedTable):
-    """
-    A fake UndefinedTable that allows setting pg_error on construction since UndefinedTable
-    has it as read-only and not settable at construction-time
-    """
-
-    def __init__(self, pg_error):
-        self.pg_error = pg_error
-
-    def __getattribute__(self, attr):
-        if attr == 'pgerror':
-            return self.pg_error
-        else:
-            return super(UndefinedTable, self).__getattribute__(attr)
-
-    def __str__(self):
-        return self.pg_error
-
-
 @pytest.mark.parametrize(
     "error,metric_columns,expected_error_tag,expected_warnings",
     [
-        (
-            ObjectNotInPrerequisiteState('pg_stat_statements must be loaded via shared_preload_libraries'),
-            [],
-            'error:database-ObjectNotInPrerequisiteState-pg_stat_statements_not_loaded',
-            [
-                'Unable to collect statement metrics because pg_stat_statements extension is '
-                "not loaded in database 'datadog_test'. See https://docs.datadoghq.com/database_monitoring/"
-                'setup_postgres/troubleshooting#pg-stat-statements-not-loaded'
-                ' for more details\ncode=pg-stat-statements-not-loaded dbname=datadog_test host=stubbed.hostname',
-            ],
-        ),
-        (
-            UndefinedTable('ERROR:  relation "pg_stat_statements" does not exist'),
-            [],
-            'error:database-UndefinedTable-pg_stat_statements_not_created',
-            [
-                'Unable to collect statement metrics because pg_stat_statements is not '
-                "created in database 'datadog_test'. See https://docs.datadoghq.com/database_monitoring/"
-                'setup_postgres/troubleshooting#pg-stat-statements-not-created'
-                ' for more details\ncode=pg-stat-statements-not-created dbname=datadog_test host=stubbed.hostname',
-            ],
-        ),
-        (
-            ObjectNotInPrerequisiteState('cannot insert into view'),
-            [],
-            'error:database-ObjectNotInPrerequisiteState',
-            [
-                "Unable to collect statement metrics because of an error running queries in database 'datadog_test'. "
-                "See https://docs.datadoghq.com/database_monitoring/troubleshooting for help: cannot insert into view\n"
-                "dbname=datadog_test host=stubbed.hostname"
-            ],
-        ),
         (
             psycopg.errors.DatabaseError('connection reset'),
             [],
@@ -2008,6 +1938,52 @@ def test_statement_metrics_database_errors(
     assert check.warnings == expected_warnings
 
 
+@pytest.mark.parametrize(
+    "error,metric_columns,expected_error_tag,expected_warnings",
+    [
+        (
+            'not-created',
+            [],
+            'error:database-UndefinedTable-pg_stat_statements_not_created',
+            [
+                'Unable to collect statement metrics because pg_stat_statements is not '
+                "created in database 'datadog_test'. See https://docs.datadoghq.com/database_monitoring/"
+                'setup_postgres/troubleshooting#pg-stat-statements-not-created'
+                ' for more details\ncode=pg-stat-statements-not-created dbname=datadog_test host=stubbed.hostname',
+            ],
+        ),
+    ],
+)
+def test_statement_metrics_database_extension_errors(
+    aggregator, integration_check, dbm_instance, error, metric_columns, expected_error_tag, expected_warnings
+):
+    # don't need samples for this test
+    dbm_instance['query_samples']['enabled'] = False
+    dbm_instance['query_activity']['enabled'] = False
+    check = integration_check(dbm_instance)
+
+    # Break databased on purpose to simulate the extension not being loaded or created
+    if error == 'not-created':
+        superconn = _get_superconn(dbm_instance)
+        with superconn.cursor() as cur:
+            cur.execute("DROP EXTENSION pg_stat_statements CASCADE;")
+
+            run_one_check(check)
+
+            # Restore extension for next test
+            cur.execute("CREATE EXTENSION IF NOT EXISTS pg_stat_statements SCHEMA public;")
+
+    expected_tags = _get_expected_tags(
+        check, dbm_instance, with_host=False, with_db=True, agent_hostname='stubbed.hostname'
+    ) + [expected_error_tag]
+
+    aggregator.assert_metric(
+        'dd.postgres.statement_metrics.error', value=1.0, count=1, tags=expected_tags, hostname='stubbed.hostname'
+    )
+
+    assert check.warnings == expected_warnings
+
+
 @pytest.mark.parametrize(
     "pg_stat_statements_max_threshold,expected_warnings",
     [