Add support for Django Connection pool (celery#9953)

* Add support for Django Connection pool

https://docs.djangoproject.com/en/dev/ref/databases/#postgresql-pool

* tests for close_pool

* close_pool called only if DB pool is enabled in Django settings

* conn pool docs

---------

Co-authored-by: Asif Saif Uddin {"Auvi":"অভি"} <[email protected]>

Author: petrprikryl

celery/fixups/django.py

@@ -1,4 +1,5 @@
 """Django-specific customization."""
+import contextlib
 import os
 import sys
 import warnings
@@ -201,6 +202,10 @@ def _close_database(self) -> None:
         for conn in self._db.connections.all():
             try:
                 conn.close()
+                pool_enabled = self._settings.DATABASES.get(conn.alias, {}).get("OPTIONS", {}).get("pool")
+                if pool_enabled and hasattr(conn, "close_pool"):
+                    with contextlib.suppress(KeyError):
+                        conn.close_pool()
             except self.interface_errors:
                 pass
             except self.DatabaseError as exc:

docs/django/first-steps-with-django.rst

@@ -216,6 +216,16 @@ However, if your app :ref:`uses a custom task base class <task-custom-classes>`,
 you'll need inherit from :class:`~celery.contrib.django.task.DjangoTask` instead of
 :class:`~celery.app.task.Task` to get this behaviour.
 
+Django Connection pool
+----------------------
+From Django 5.1+ there is built-in support for database connection pooling.
+If you enable it in Django ``DATABASES`` settings Celery will automatically
+handle connection pool closing in worker processes via ``close_pool``
+database backend method as
+`sharing connections across processes is not possible. <https://github.com/psycopg/psycopg/issues/544#issuecomment-1500886864>`_
+
+You can find more about Connection pool at `Django docs. <https://docs.djangoproject.com/en/dev/ref/databases/#connection-pool>`_
+
 Extensions
 ==========
 

t/unit/fixups/test_django.py

@@ -284,6 +284,58 @@ def test_close_database_always_closes_connections(self):
             # it to optimize connection handling.
             conn.close_if_unusable_or_obsolete.assert_not_called()
 
+    def test_close_database_skip_conn_pool(self):
+        class Connection:
+            """Mock connection without `close_pool` method."""
+            alias = 'default'
+
+            def close(self):
+                pass
+
+        with self.fixup_context(self.app) as (f, _, _):
+            conn = Mock(spec=Connection)
+            f._db.connections.all = Mock(return_value=[conn])
+            f.close_database()
+            assert not hasattr(conn, "close_pool")
+            conn.close.assert_called_once_with()
+
+    def test_close_database_suppresses_close_pool_keyerror(self):
+        with self.fixup_context(self.app) as (f, _, _):
+            conn = Mock()
+            conn.close_pool = Mock(side_effect=KeyError("pool already closed"))
+            f._db.connections.all = Mock(return_value=[conn])
+            f.close_database()  # should not raise
+            conn.close.assert_called_once_with()
+            conn.close_pool.assert_called_once_with()
+
+    def test_close_database_conn_pool_based_on_settings(self):
+        class DJSettings:
+            DATABASES = {}
+
+        with self.fixup_context(self.app) as (f, _, _):
+            conn = Mock()
+            conn.alias = "default"
+            conn.close_pool = Mock()
+            f._db.connections.all = Mock(return_value=[conn])
+            f._settings = DJSettings
+
+            f._settings.DATABASES["default"] = {"OPTIONS": {}}
+            f.close_database()
+            conn.close.assert_called_once_with()
+            conn.close_pool.assert_not_called()
+
+            conn.reset_mock()
+            f._settings.DATABASES["default"] = {"OPTIONS": {"pool": True}}
+            f.close_database()
+            conn.close.assert_called_once_with()
+            conn.close_pool.assert_called_once_with()
+
+            conn.reset_mock()
+            f._settings.DATABASES["default"] = {"OPTIONS": {"pool": False}}
+            f.close_database()
+            conn.close.assert_called_once_with()
+            conn.close_pool.assert_not_called()
+
     def test_close_cache_raises_error(self):
         with self.fixup_context(self.app) as (f, _, _):
             f._cache.close_caches.side_effect = AttributeError