Fixed #30398 -- Added CONN_HEALTH_CHECKS database setting.

The CONN_HEALTH_CHECKS setting can be used to enable database
connection health checks for Django's persistent DB connections.

Thanks Florian Apolloner for reviews.

1 files changed, 155 insertions, 2 deletions

diff --git a/tests/backends/base/test_base.py b/tests/backends/base/test_base.py
index b07d6d3ceb..b2ebea5829 100644
--- a/tests/backends/base/test_base.py
+++ b/tests/backends/base/test_base.py
@@ -1,8 +1,8 @@
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, patch
 
 from django.db import DEFAULT_DB_ALIAS, connection, connections
 from django.db.backends.base.base import BaseDatabaseWrapper
-from django.test import SimpleTestCase, TestCase
+from django.test import SimpleTestCase, TestCase, skipUnlessDBFeature
 
 from ..models import Square
 
@@ -134,3 +134,156 @@ class ExecuteWrapperTests(TestCase):
         self.assertFalse(wrapper.called)
         self.assertEqual(connection.execute_wrappers, [])
         self.assertEqual(connections['other'].execute_wrappers, [])
+
+
+class ConnectionHealthChecksTests(SimpleTestCase):
+    databases = {'default'}
+
+    def setUp(self):
+        # All test cases here need newly configured and created connections.
+        # Use the default db connection for convenience.
+        connection.close()
+        self.addCleanup(connection.close)
+
+    def patch_settings_dict(self, conn_health_checks):
+        self.settings_dict_patcher = patch.dict(connection.settings_dict, {
+            **connection.settings_dict,
+            'CONN_MAX_AGE': None,
+            'CONN_HEALTH_CHECKS': conn_health_checks,
+        })
+        self.settings_dict_patcher.start()
+        self.addCleanup(self.settings_dict_patcher.stop)
+
+    def run_query(self):
+        with connection.cursor() as cursor:
+            cursor.execute('SELECT 42' + connection.features.bare_select_suffix)
+
+    @skipUnlessDBFeature('test_db_allows_multiple_connections')
+    def test_health_checks_enabled(self):
+        self.patch_settings_dict(conn_health_checks=True)
+        self.assertIsNone(connection.connection)
+        # Newly created connections are considered healthy without performing
+        # the health check.
+        with patch.object(connection, 'is_usable', side_effect=AssertionError):
+            self.run_query()
+
+        old_connection = connection.connection
+        # Simulate request_finished.
+        connection.close_if_unusable_or_obsolete()
+        self.assertIs(old_connection, connection.connection)
+
+        # Simulate connection health check failing.
+        with patch.object(connection, 'is_usable', return_value=False) as mocked_is_usable:
+            self.run_query()
+            new_connection = connection.connection
+            # A new connection is established.
+            self.assertIsNot(new_connection, old_connection)
+            # Only one health check per "request" is performed, so the next
+            # query will carry on even if the health check fails. Next query
+            # succeeds because the real connection is healthy and only the
+            # health check failure is mocked.
+            self.run_query()
+            self.assertIs(new_connection, connection.connection)
+        self.assertEqual(mocked_is_usable.call_count, 1)
+
+        # Simulate request_finished.
+        connection.close_if_unusable_or_obsolete()
+        # The underlying connection is being reused further with health checks
+        # succeeding.
+        self.run_query()
+        self.run_query()
+        self.assertIs(new_connection, connection.connection)
+
+    @skipUnlessDBFeature('test_db_allows_multiple_connections')
+    def test_health_checks_enabled_errors_occurred(self):
+        self.patch_settings_dict(conn_health_checks=True)
+        self.assertIsNone(connection.connection)
+        # Newly created connections are considered healthy without performing
+        # the health check.
+        with patch.object(connection, 'is_usable', side_effect=AssertionError):
+            self.run_query()
+
+        old_connection = connection.connection
+        # Simulate errors_occurred.
+        connection.errors_occurred = True
+        # Simulate request_started (the connection is healthy).
+        connection.close_if_unusable_or_obsolete()
+        # Persistent connections are enabled.
+        self.assertIs(old_connection, connection.connection)
+        # No additional health checks after the one in
+        # close_if_unusable_or_obsolete() are executed during this "request"
+        # when running queries.
+        with patch.object(connection, 'is_usable', side_effect=AssertionError):
+            self.run_query()
+
+    @skipUnlessDBFeature('test_db_allows_multiple_connections')
+    def test_health_checks_disabled(self):
+        self.patch_settings_dict(conn_health_checks=False)
+        self.assertIsNone(connection.connection)
+        # Newly created connections are considered healthy without performing
+        # the health check.
+        with patch.object(connection, 'is_usable', side_effect=AssertionError):
+            self.run_query()
+
+        old_connection = connection.connection
+        # Simulate request_finished.
+        connection.close_if_unusable_or_obsolete()
+        # Persistent connections are enabled (connection is not).
+        self.assertIs(old_connection, connection.connection)
+        # Health checks are not performed.
+        with patch.object(connection, 'is_usable', side_effect=AssertionError):
+            self.run_query()
+            # Health check wasn't performed and the connection is unchanged.
+            self.assertIs(old_connection, connection.connection)
+            self.run_query()
+            # The connection is unchanged after the next query either during
+            # the current "request".
+            self.assertIs(old_connection, connection.connection)
+
+    @skipUnlessDBFeature('test_db_allows_multiple_connections')
+    def test_set_autocommit_health_checks_enabled(self):
+        self.patch_settings_dict(conn_health_checks=True)
+        self.assertIsNone(connection.connection)
+        # Newly created connections are considered healthy without performing
+        # the health check.
+        with patch.object(connection, 'is_usable', side_effect=AssertionError):
+            # Simulate outermost atomic block: changing autocommit for
+            # a connection.
+            connection.set_autocommit(False)
+            self.run_query()
+            connection.commit()
+            connection.set_autocommit(True)
+
+        old_connection = connection.connection
+        # Simulate request_finished.
+        connection.close_if_unusable_or_obsolete()
+        # Persistent connections are enabled.
+        self.assertIs(old_connection, connection.connection)
+
+        # Simulate connection health check failing.
+        with patch.object(connection, 'is_usable', return_value=False) as mocked_is_usable:
+            # Simulate outermost atomic block: changing autocommit for
+            # a connection.
+            connection.set_autocommit(False)
+            new_connection = connection.connection
+            self.assertIsNot(new_connection, old_connection)
+            # Only one health check per "request" is performed, so a query will
+            # carry on even if the health check fails. This query succeeds
+            # because the real connection is healthy and only the health check
+            # failure is mocked.
+            self.run_query()
+            connection.commit()
+            connection.set_autocommit(True)
+            # The connection is unchanged.
+            self.assertIs(new_connection, connection.connection)
+        self.assertEqual(mocked_is_usable.call_count, 1)
+
+        # Simulate request_finished.
+        connection.close_if_unusable_or_obsolete()
+        # The underlying connection is being reused further with health checks
+        # succeeding.
+        connection.set_autocommit(False)
+        self.run_query()
+        connection.commit()
+        connection.set_autocommit(True)
+        self.assertIs(new_connection, connection.connection)