style

postgres/datadog_checks/postgres/statement_samples.py

@@ -1,4 +1,3 @@
-import itertools
 import logging
 import os
 import re
@@ -138,7 +137,7 @@ def run_sampler(self, tags):
         self._tags_no_db = [t for t in tags if not t.startswith('db:')]
         for t in self._tags:
             if t.startswith('service:'):
-                self._service = t[len('service:'):]
+                self._service = t[len('service:') :]
 
         self._last_check_run = time.time()
         if self._run_sync or is_affirmative(os.environ.get('DBM_STATEMENT_SAMPLER_RUN_SYNC', "false")):
@@ -294,8 +293,9 @@ def _get_db_explain_setup_state(self, dbname):
         try:
             self._get_db(dbname)
         except psycopg2.DatabaseError as e:
-            self._log.warning("cannot collect execution plans due to failed DB connection to dbname=%s: %s", dbname,
-                              repr(e))
+            self._log.warning(
+                "cannot collect execution plans due to failed DB connection to dbname=%s: %s", dbname, repr(e)
+            )
             return DBExplainSetupState.failed_connect
 
         try:
@@ -330,8 +330,7 @@ def _get_db_explain_setup_state_cached(self, dbname):
     def _run_explain(self, dbname, statement, obfuscated_statement):
         start_time = time.time()
         with self._get_db(dbname).cursor() as cursor:
-            self._log.debug("Running query on dbname=%s: %s(%s)", dbname, self._explain_function,
-                            obfuscated_statement)
+            self._log.debug("Running query on dbname=%s: %s(%s)", dbname, self._explain_function, obfuscated_statement)
             cursor.execute(
                 """SELECT {explain_function}($stmt${statement}$stmt$)""".format(
                     explain_function=self._explain_function, statement=statement
@@ -352,8 +351,9 @@ def _run_explain_safe(self, dbname, statement, obfuscated_statement):
         explain_setup_state = self._get_db_explain_setup_state_cached(dbname)
         if explain_setup_state != DBExplainSetupState.ok:
             self._check.count(
-                "dd.postgres.statement_samples.error", 1,
-                tags=self._dbtags(dbname) + ["error:explain-{}".format(explain_setup_state)]
+                "dd.postgres.statement_samples.error",
+                1,
+                tags=self._dbtags(dbname) + ["error:explain-{}".format(explain_setup_state)],
             )
             return None
 
@@ -362,8 +362,9 @@ def _run_explain_safe(self, dbname, statement, obfuscated_statement):
         except psycopg2.errors.DatabaseError as e:
             self._log.warning("Failed to collect execution plan: %s", repr(e))
             self._check.count(
-                "dd.postgres.statement_samples.error", 1,
-                tags=self._dbtags(dbname) + ["error:explain-{}".format(type(e))]
+                "dd.postgres.statement_samples.error",
+                1,
+                tags=self._dbtags(dbname) + ["error:explain-{}".format(type(e))],
             )
             return None
 
@@ -372,8 +373,9 @@ def _collect_plan_for_statement(self, row):
             obfuscated_statement = datadog_agent.obfuscate_sql(row['query'])
         except Exception as e:
             self._log.debug("Failed to obfuscate statement: %s", e)
-            self._check.count("dd.postgres.statement_samples.error", 1,
-                              tags=self._dbtags(row['datname']) + ["error:sql-obfuscate"])
+            self._check.count(
+                "dd.postgres.statement_samples.error", 1, tags=self._dbtags(row['datname']) + ["error:sql-obfuscate"]
+            )
             return None
 
         # limit the rate of explains done to the database
@@ -446,8 +448,9 @@ def _explain_pg_stat_activity(self, rows):
                 if event:
                     yield event
             except Exception:
-                self._log.exception("Crashed trying to collect execution plan for statement in dbname=%s",
-                                    row['datname'])
+                self._log.exception(
+                    "Crashed trying to collect execution plan for statement in dbname=%s", row['datname']
+                )
                 self._check.count(
                     "dd.postgres.statement_samples.error",
                     1,

postgres/tests/test_pg_integration.py

@@ -1,7 +1,6 @@
 # (C) Datadog, Inc. 2010-present
 # All rights reserved
 # Licensed under Simplified BSD License (see LICENSE)
-import pdb
 import socket
 import time
 from concurrent.futures.thread import ThreadPoolExecutor
@@ -13,7 +12,7 @@
 
 from datadog_checks.base.utils.serialization import json
 from datadog_checks.postgres import PostgreSql
-from datadog_checks.postgres.statement_samples import PostgresStatementSamples, DBExplainSetupState
+from datadog_checks.postgres.statement_samples import DBExplainSetupState, PostgresStatementSamples
 from datadog_checks.postgres.util import PartialFormatter, fmt
 
 from .common import DB_NAME, HOST, PORT, POSTGRES_VERSION, check_bgw_metrics, check_common_metrics
@@ -315,22 +314,50 @@ 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", "dogs_noschema", "SELECT * FROM kennel WHERE id = %s", 123,
-         "error:explain-{}".format(DBExplainSetupState.invalid_schema)),
-        ("dd_admin", "dd_admin", "dogs_nofunc", "SELECT * FROM kennel WHERE id = %s", 123,
-         "error:explain-{}".format(DBExplainSetupState.failed_function)),
+        (
+            "dd_admin",
+            "dd_admin",
+            "dogs_noconn",
+            "SELECT * FROM kennel WHERE address = %s",
+            "Fake",
+            "error:explain-{}".format(DBExplainSetupState.failed_connect),
+        ),
+        (
+            "dd_admin",
+            "dd_admin",
+            "dogs_noschema",
+            "SELECT * FROM kennel WHERE id = %s",
+            123,
+            "error:explain-{}".format(DBExplainSetupState.invalid_schema),
+        ),
+        (
+            "dd_admin",
+            "dd_admin",
+            "dogs_nofunc",
+            "SELECT * FROM kennel WHERE id = %s",
+            123,
+            "error:explain-{}".format(DBExplainSetupState.failed_function),
+        ),
     ],
 )
-def test_statement_samples_collect(aggregator, integration_check, dbm_instance, pg_stat_activity_view, user, password,
-                                   dbname, query, arg, expected_error_tag):
+def test_statement_samples_collect(
+    aggregator,
+    integration_check,
+    dbm_instance,
+    pg_stat_activity_view,
+    user,
+    password,
+    dbname,
+    query,
+    arg,
+    expected_error_tag,
+):
     dbm_instance['pg_stat_activity_view'] = pg_stat_activity_view
     check = integration_check(dbm_instance)
     check._connect()
 
     tags = dbm_instance['tags'] + [
-        'server:{}'.format(HOST) ,
+        'server:{}'.format(HOST),
         'port:{}'.format(PORT),
         'db:{}'.format(dbname),
     ]