From 375fe2e3e9d0a1c570fa1818ed473de406da3ce4 Mon Sep 17 00:00:00 2001 From: Jo-Ann Meunier Date: Sun, 6 Aug 2023 14:46:16 -0400 Subject: [PATCH] cleanup --- postgres/datadog_checks/postgres/connections.py | 9 ++++----- postgres/tests/test_statements.py | 4 ++-- 2 files changed, 6 insertions(+), 7 deletions(-) diff --git a/postgres/datadog_checks/postgres/connections.py b/postgres/datadog_checks/postgres/connections.py index 28194a4ce34e8..2570f65860c0e 100644 --- a/postgres/datadog_checks/postgres/connections.py +++ b/postgres/datadog_checks/postgres/connections.py @@ -217,7 +217,7 @@ def _terminate_connection_unsafe(self, dbname: str): return False return True - def get_main_db(self): + def _get_main_db(self): """ Returns a memoized, persistent psycopg connection to `self.dbname`. Utilizes the db connection pool, and is meant to be shared across multiple threads. @@ -241,18 +241,17 @@ def execute_main_db_safe(self, query, params=None, row_format=None): with self._query_lock: self._log.debug("Running query [{}] {}".format(query, params)) if row_format: - with self.get_main_db().cursor(row_factory=row_format) as cursor: + with self._get_main_db().cursor(row_factory=row_format) as cursor: cursor.execute(query, params) return cursor.fetchall() else: - with self.get_main_db().cursor() as cursor: + with self._get_main_db().cursor() as cursor: cursor.execute(query, params) return cursor.fetchall() def get_main_db_columns_safe(self, query, params): with self._query_lock: - with self.get_main_db().cursor() as cursor: + with self._get_main_db().cursor() as cursor: self._log.debug("Running query [%s] %s", query, params) cursor.execute(query, params) return [desc[0] for desc in cursor.description] if cursor.description else [] - diff --git a/postgres/tests/test_statements.py b/postgres/tests/test_statements.py index 5850d6f4e2c1c..92148969c0742 100644 --- a/postgres/tests/test_statements.py +++ b/postgres/tests/test_statements.py @@ -1359,10 +1359,10 @@ def test_pg_settings_caching(aggregator, integration_check, dbm_instance): check = integration_check(dbm_instance) assert not check.pg_settings, "pg_settings should not have been initialized yet" check._connect() - check.db_pool.get_main_db() + check.db_pool._get_main_db() assert "track_activity_query_size" in check.pg_settings check.pg_settings["test_key"] = True - check.db_pool.get_main_db() + check.db_pool._get_main_db() assert ( "test_key" in check.pg_settings ), "key should not have been blown away. If it was then pg_settings was not cached correctly"