fix

postgres/datadog_checks/postgres/statement_samples.py

@@ -292,7 +292,7 @@ def _can_explain_statement(self, obfuscated_statement):
     def _get_db_explain_setup_state(self, dbname):
         try:
             self._get_db(dbname)
-        except psycopg2.DatabaseError as e:
+        except (psycopg2.DatabaseError, psycopg2.OperationalError) as e:
             self._log.warning(
                 "cannot collect execution plans due to failed DB connection to dbname=%s: %s", dbname, repr(e)
             )

postgres/tests/compose/resources/01_postgres.sql

@@ -7,8 +7,5 @@ CREATE DATABASE datadog_test;
 GRANT ALL PRIVILEGES ON DATABASE datadog_test TO datadog;
 CREATE DATABASE dogs;
 GRANT USAGE on SCHEMA public to bob;
-CREATE DATABASE dogs_noconn;
-REVOKE ALL PRIVILEGES ON DATABASE dogs_noconn FROM datadog;
-REVOKE CONNECT ON DATABASE dogs_noconn FROM public;
 CREATE DATABASE dogs_nofunc;
 CREATE DATABASE dogs_noschema;

postgres/tests/compose/resources/02_setup.sh

@@ -2,6 +2,9 @@
 set -e
 
 # pg_monitor is only available on 10+
+# prior to version 10 there was no `pg_read_all_stats` role so by adding a database to which the agent can't connect
+# it causes many of the tests to fail since some stats queries fail (like reading the size of the database to which you can't connect)
+# therefore we will only add this database to which you can't connect in 10+ for testing purposes
 if [[ !("$PG_MAJOR" == 9.* ) ]]; then
 psql -v ON_ERROR_STOP=1 --username "$POSTGRES_USER" datadog_test <<-'EOSQL'
     GRANT pg_monitor TO datadog;
@@ -14,7 +17,7 @@ psql -v ON_ERROR_STOP=1 --username "$POSTGRES_USER" "datadog_test" <<-'EOSQL'
 EOSQL
 
 # dogs_noschema deliberately excluded
-for DBNAME in datadog_test dogs dogs_noconn dogs_nofunc; do
+for DBNAME in datadog_test dogs dogs_nofunc; do
 
 psql -v ON_ERROR_STOP=1 --username "$POSTGRES_USER" "$DBNAME" <<-'EOSQL'
     CREATE SCHEMA datadog;

postgres/tests/compose/resources/03_load_data.sh

@@ -14,7 +14,7 @@ psql -v ON_ERROR_STOP=1 --username "$POSTGRES_USER" datadog_test <<-EOSQL
     GRANT ALL PRIVILEGES ON ALL TABLES IN SCHEMA public TO bob;
 EOSQL
 
-for DBNAME in dogs dogs_noconn dogs_noschema dogs_nofunc; do
+for DBNAME in dogs dogs_noschema dogs_nofunc; do
 
 psql -v ON_ERROR_STOP=1 --username "$POSTGRES_USER" "$DBNAME" <<-EOSQL
     CREATE TABLE breed (id SERIAL, name VARCHAR(255));
@@ -27,4 +27,3 @@ psql -v ON_ERROR_STOP=1 --username "$POSTGRES_USER" "$DBNAME" <<-EOSQL
 EOSQL
 
 done
-

postgres/tests/test_pg_integration.py

@@ -137,7 +137,7 @@ def test_schema_metrics(aggregator, integration_check, pg_instance):
         'schema:public',
     ]
     aggregator.assert_metric('postgresql.table.count', value=1, count=1, tags=expected_tags)
-    aggregator.assert_metric('postgresql.db.count', value=5, count=1)
+    aggregator.assert_metric('postgresql.db.count', value=4, count=1)
 
 
 def test_connections_metrics(aggregator, integration_check, pg_instance):
@@ -297,7 +297,6 @@ def dbm_instance(pg_instance):
     [
         ("datadog_test", DBExplainSetupState.ok),
         ("dogs", DBExplainSetupState.ok),
-        ("dogs_noconn", DBExplainSetupState.failed_connect),
         ("dogs_noschema", DBExplainSetupState.invalid_schema),
         ("dogs_nofunc", DBExplainSetupState.failed_function),
     ],
@@ -314,14 +313,6 @@ def test_get_db_explain_setup_state(integration_check, dbm_instance, dbname, exp
     [
         ("bob", "bob", "datadog_test", "SELECT city FROM persons WHERE city = %s", "hello", None),
         ("dd_admin", "dd_admin", "dogs", "SELECT * FROM breed WHERE name = %s", "Labrador", None),
-        (
-            "dd_admin",
-            "dd_admin",
-            "dogs_noconn",
-            "SELECT * FROM kennel WHERE address = %s",
-            "Fake",
-            "error:explain-{}".format(DBExplainSetupState.failed_connect),
-        ),
         (
             "dd_admin",
             "dd_admin",
@@ -373,13 +364,14 @@ def test_statement_samples_collect(
         check.check(dbm_instance)
         dbm_samples = aggregator.get_event_platform_events("dbm-samples")
 
+        expected_query = query % ('\'' + arg + '\'' if isinstance(arg, str) else arg)
+        matching = [e for e in dbm_samples if e['db']['statement'] == expected_query]
+
         if POSTGRES_VERSION.split('.')[0] == "9" and pg_stat_activity_view == "pg_stat_activity":
             # pg_monitor role exists only in version 10+
-            assert len(dbm_samples) == 0, "did not expect to catch any events"
+            assert len(matching) == 0, "did not expect to catch any events"
             return
 
-        expected_query = query % ('\'' + arg + '\'' if isinstance(arg, str) else arg)
-        matching = [e for e in dbm_samples if e['db']['statement'] == expected_query]
         assert len(matching) == 1, "missing captured event"
         event = matching[0]