Initial commit (Clean history)

path/to/venv/lib/python3.12/site-packages/django_db_geventpool/__init__.py

@@ -0,0 +1 @@
+version = 'v4.0.8'

path/to/venv/lib/python3.12/site-packages/django_db_geventpool/backends/base.py

@@ -0,0 +1,94 @@
+import logging
+import sys
+
+try:
+    from gevent.lock import Semaphore
+except ImportError:
+    from eventlet.semaphore import Semaphore
+
+from .creation import DatabaseCreation
+
+logger = logging.getLogger("django.geventpool")
+
+connection_pools_lock = Semaphore(value=1)
+
+
+class DatabaseWrapperMixin:
+    pool_class = None
+    creation_class = DatabaseCreation
+    INTRANS = None
+    _connection_pools = {}
+
+
+    def __init__(self, *args, **kwargs):
+        self._pool = None
+        super().__init__(*args, **kwargs)
+        self.creation = self.creation_class(self)
+
+    @property
+    def pool(self):
+        if self._pool is not None:
+            return self._pool
+        with connection_pools_lock:
+            if self.alias not in self._connection_pools:
+                self._pool = self.pool_class(**self.get_connection_params())
+                self._connection_pools[self.alias] = self._pool
+            else:
+                self._pool = self._connection_pools[self.alias]
+        return self._pool
+
+    def get_new_connection(self, conn_params: dict):
+        if self.connection is None:
+            self.connection = self.pool.get()
+            self.closed_in_transaction = False
+        return self.connection
+
+    def get_connection_params(self) -> dict:
+        conn_params = super().get_connection_params()
+        for attr in ["MAX_CONNS", "REUSE_CONNS"]:
+            if attr in self.settings_dict["OPTIONS"]:
+                conn_params[attr] = self.settings_dict["OPTIONS"][attr]
+        return conn_params
+
+    def close(self):
+        self.validate_thread_sharing()
+        if self.closed_in_transaction or self.connection is None:
+            return  # no need to close anything
+        try:
+            self._close()
+        except:
+            # In some cases (database restart, network connection lost etc...)
+            # the connection to the database is lost without giving Django a
+            # notification. If we don't set self.connection to None, the error
+            # will occur at every request.
+            self.connection = None
+            logger.warning(
+                "psycopg error while closing the connection.", exc_info=sys.exc_info()
+            )
+            raise
+        finally:
+            self.set_clean()
+
+    def close_if_unusable_or_obsolete(self):
+        # Always close the connection because it's not (usually) really being closed.
+        self.close()
+
+    def _close(self):
+        if self.connection.closed:
+            self.pool.close()
+        else:
+            if self.connection.info.transaction_status == self.INTRANS:
+                self.connection.rollback()
+                self.connection.autocommit = True
+            with self.wrap_database_errors:
+                self.pool.put(self.connection)
+        self.connection = None
+
+    def closeall(self):
+        for pool in self._connection_pools.values():
+            pool.close()
+
+    def set_clean(self):
+        if self.in_atomic_block:
+            self.closed_in_transaction = True
+            self.needs_rollback = True

path/to/venv/lib/python3.12/site-packages/django_db_geventpool/backends/creation.py

@@ -0,0 +1,16 @@
+from django.db.backends.postgresql.creation import (
+    DatabaseCreation as OriginalDatabaseCreation,
+)
+
+
+class DatabaseCreationMixin:
+    def _create_test_db(self, verbosity, autoclobber, keepdb=False):
+        return super()._create_test_db(verbosity, autoclobber, keepdb)
+
+    def _destroy_test_db(self, test_database_name, verbosity):
+        self.connection.closeall()
+        return super()._destroy_test_db(test_database_name, verbosity)
+
+
+class DatabaseCreation(DatabaseCreationMixin, OriginalDatabaseCreation):
+    pass

path/to/venv/lib/python3.12/site-packages/django_db_geventpool/backends/pool.py

@@ -0,0 +1,85 @@
+# this file is a modified version of the psycopg2 used at gevent examples
+# to be compatible with django, also checks if
+# DB connection is closed and reopen it:
+# https://github.com/surfly/gevent/blob/master/examples/psycopg2_pool.py
+import logging
+import weakref
+
+logger = logging.getLogger("django.geventpool")
+
+try:
+    from gevent import queue
+    from gevent.lock import RLock
+except ImportError:
+    from eventlet import queue
+    from ...utils import NullContextRLock as RLock
+
+
+class DatabaseConnectionPool:
+    DBERROR = None
+
+    def __init__(self, maxsize: int = 100, reuse: int = 100):
+        # Use a WeakSet here so, even if we fail to discard the connection
+        # when it is being closed, or it is closed outside of here, the item
+        # will be removed automatically
+        self._conns = weakref.WeakSet()
+        self.maxsize = maxsize
+        self.pool = queue.Queue(maxsize=max(reuse, 1))
+        self.lock = RLock()
+
+    @property
+    def size(self):
+        with self.lock:
+            return len(self._conns)
+
+    def get(self):
+        try:
+            if self.size >= self.maxsize or self.pool.qsize():
+                conn = self.pool.get()
+            else:
+                conn = self.pool.get_nowait()
+
+            try:
+                # check connection is still valid
+                self.check_usable(conn)
+                logger.debug("DB connection reused")
+            except self.DBERROR:
+                logger.debug("DB connection was closed, creating a new one")
+                conn = None
+        except queue.Empty:
+            conn = None
+            logger.debug("DB connection queue empty, creating a new one")
+
+        if conn is None:
+            try:
+                conn = self.create_connection()
+            except Exception:
+                raise
+            else:
+                self._conns.add(conn)
+
+        return conn
+
+    def put(self, item):
+        try:
+            self.pool.put_nowait(item)
+            logger.debug("DB connection returned to the pool")
+        except queue.Full:
+            item.close()
+            self._conns.discard(item)
+
+    def close(self):
+        while not self.pool.empty():
+            try:
+                conn = self.pool.get_nowait()
+            except queue.Empty:
+                continue
+            try:
+                conn.close()
+            except Exception:
+                continue
+            finally:
+                self._conns.discard(conn)
+
+        logger.debug("DB connections all closed")
+

path/to/venv/lib/python3.12/site-packages/django_db_geventpool/backends/postgis/base.py

@@ -0,0 +1,19 @@
+from django.contrib.gis.db.backends.postgis.base import (
+    DatabaseWrapper as OriginalDatabaseWrapper,
+)
+
+try:
+    # try psycopg3
+    import psycopg  # noqa
+    INTRANS = psycopg.pq.TransactionStatus.INTRANS
+    from ..postgresql_psycopg3.base import base, PostgresConnectionPool
+except ImportError:
+    # fallback to psycopg2
+    import psycopg2
+    INTRANS = psycopg2.extensions.TRANSACTION_STATUS_INTRANS
+    from ..postgresql_psycopg2.base import base, PostgresConnectionPool
+
+
+class DatabaseWrapper(base.DatabaseWrapperMixin, OriginalDatabaseWrapper):
+    pool_class = PostgresConnectionPool
+    INTRANS = INTRANS

path/to/venv/lib/python3.12/site-packages/django_db_geventpool/backends/postgresql_psycopg2/base.py

@@ -0,0 +1,41 @@
+try:
+    import psycopg2
+    import psycopg2.extras
+    import psycopg2.extensions
+except ImportError as e:
+    from django.core.exceptions import ImproperlyConfigured
+
+    raise ImproperlyConfigured("Error loading psycopg2 module: %s" % e)
+
+from django.db.backends.postgresql.base import (
+    DatabaseWrapper as OriginalDatabaseWrapper,
+)
+
+from .. import base, pool
+
+
+class PostgresConnectionPool(pool.DatabaseConnectionPool):
+    DBERROR = psycopg2.DatabaseError
+
+    def __init__(self, *args, **kwargs):
+        self.connect = kwargs.pop("connect", psycopg2.connect)
+        self.connection = None
+        maxsize = kwargs.pop("MAX_CONNS", 4)
+        reuse = kwargs.pop("REUSE_CONNS", maxsize)
+        self.args = args
+        self.kwargs = kwargs
+        self.kwargs["client_encoding"] = "UTF8"
+        super().__init__(maxsize, reuse)
+
+    def create_connection(self):
+        conn = self.connect(*self.args, **self.kwargs)
+        psycopg2.extras.register_default_jsonb(conn_or_curs=conn, loads=lambda x: x)
+        return conn
+
+    def check_usable(self, connection):
+        connection.cursor().execute("SELECT 1")
+
+
+class DatabaseWrapper(base.DatabaseWrapperMixin, OriginalDatabaseWrapper):
+    pool_class = PostgresConnectionPool
+    INTRANS = psycopg2.extensions.TRANSACTION_STATUS_INTRANS

path/to/venv/lib/python3.12/site-packages/django_db_geventpool/backends/postgresql_psycopg3/base.py

@@ -0,0 +1,41 @@
+try:
+    import psycopg
+except ImportError as e:
+    from django.core.exceptions import ImproperlyConfigured
+
+    raise ImproperlyConfigured("Error loading psycopg3 module: %s" % e)
+
+from django.db.backends.postgresql.base import (
+    DatabaseWrapper as OriginalDatabaseWrapper,
+)
+
+from .. import base, pool
+
+
+class PostgresConnectionPool(pool.DatabaseConnectionPool):
+    DBERROR = psycopg.DatabaseError
+
+    def __init__(self, *args, **kwargs):
+        self.connect = kwargs.pop("connect", psycopg.connect)
+        self.connection = None
+        maxsize = kwargs.pop("MAX_CONNS", 4)
+        reuse = kwargs.pop("REUSE_CONNS", maxsize)
+        self.args = args
+        self.kwargs = kwargs
+        self.kwargs["client_encoding"] = "UTF8"
+        super().__init__(maxsize, reuse)
+
+    def create_connection(self):
+        conn = self.connect(*self.args, **self.kwargs)
+        return conn
+
+    def check_usable(self, connection):
+        connection.cursor().execute("SELECT 1")
+
+
+class DatabaseWrapper(base.DatabaseWrapperMixin, OriginalDatabaseWrapper):
+    pool_class = PostgresConnectionPool
+    INTRANS = psycopg.pq.TransactionStatus.INTRANS
+
+
+

path/to/venv/lib/python3.12/site-packages/django_db_geventpool/utils.py

@@ -0,0 +1,25 @@
+from functools import wraps
+
+from django.core.signals import request_finished
+
+
+def close_connection(f):
+    @wraps(f)
+    def wrapper(*args, **kwargs):
+        try:
+            return f(*args, **kwargs)
+        finally:
+            request_finished.send(sender="greenlet")
+
+    return wrapper
+
+
+class NullContextRLock:
+    def __init__(self, enter_result=None):
+        self._enter_result = enter_result
+
+    def __enter__(self):
+        return self._enter_result
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        return None