[DBMON-6882] Make Postgres check teardown idempotent

postgres/changelog.d/24852.fixed

@@ -0,0 +1 @@
+Only tear the check down once when it is cancelled, and keep logging usable through the end of teardown.

postgres/datadog_checks/postgres/postgres.py

@@ -199,6 +199,7 @@ def __init__(self, name, init_config, instances):
         self._cancel_lock = threading.Lock()
         self._is_running = False
         self._cancelled = False
+        self._finalized = False
 
     def database_monitoring_column_statistics(self, raw_event: str):
         self.event_platform_event(raw_event, "dbm-column-statistics")
@@ -537,19 +538,26 @@ def _register_async_jobs(self):
             self.data_observability = self.register_async_job(PostgresDataObservability(self, self._config))
 
     def _finalize(self):
-        """Tear down check state. Must not run while check() is executing."""
+        """Tear down check state. Runs at most once, and never while check() is executing."""
+        with self._cancel_lock:
+            if self._finalized:
+                self.log.debug("Check already finalized, nothing to tear down")
+                return
+            self._finalized = True
         self.log.debug("Finalizing check: closing connections and clearing state")
         self.shutdown_async_jobs()
         self._clean_state()
         self.check_initializations.clear()
         # TODO: move diagnosis cleanup into AgentCheck.cancel() in the base class
         self._diagnosis = None
-        self.log.check = None
         self._query_manager = None
         self.health = None
         self._close_db()
         self._close_db_pool()
         self.log.debug("Check cleanup complete")
+        # Must come last: the logging adapter reads back through this attribute for checks whose
+        # check_id was never resolved, so anything logged after this would fail.
+        self.log.check = None
 
     def _clean_state(self):
         self.log.debug("Cleaning state")

postgres/tests/test_unit.py

@@ -519,6 +519,21 @@ def test_run_after_cancel_returns_immediately(pg_instance):
     assert result == ''
 
 
+def test_finalize_runs_once_across_repeated_cancels(pg_instance):
+    """Verify that teardown is idempotent."""
+    check = PostgreSql('postgres', {}, [pg_instance])
+    conn = mock.MagicMock()
+    check._db = conn
+
+    with mock.patch.object(check.db_pool, 'close_all', wraps=check.db_pool.close_all) as close_all:
+        check.cancel()
+        check.cancel()
+        check._finalize()
+
+    conn.close.assert_called_once()
+    close_all.assert_called_once()
+
+
 @pytest.mark.parametrize(
     'dbm, data_observability_enabled, expected_jobs',
     [