working on psycopg3 (#81)

* refactor some code

* ruff formatting

* split common code for psycopg2 and psycopg3

* Removed python2 code

* split code for pool

* added psycopg3 pool

* split tests for psycopg

* updated readme

---------

Co-authored-by: Javier Cordero <[email protected]>

.github/workflows/ci.yml → .github/workflows/ci_psycopg2.yml

@@ -1,4 +1,4 @@
-name: CI
+name: CI-psycopg2
 'on':
   push:
     branches:
@@ -60,6 +60,7 @@ jobs:
       - uses: actions/checkout@v4
       - run: pip install django==${{ matrix.django-version}}
       - run: pip install psycopg2
+      - run: pip install psycogreen
       - run: pip install gevent
       - run: python setup.py -q install
-      - run: python runtests.py
+      - run: python runtests_psycopg2.py

.github/workflows/ci_psycopg3.yml

@@ -0,0 +1,56 @@
+name: CI-psycopg3
+'on':
+  push:
+    branches:
+      - master
+  pull_request:
+    branches:
+      - master
+jobs:
+  build:
+    env:
+      POSTGRES_USER: postgres
+      PGPASSWORD: postgres
+    runs-on: '${{ matrix.os }}'
+    strategy:
+      matrix:
+        include:
+          - os: ubuntu-latest
+            python-version: '3.8'
+            django-version: '4.2.9'
+          - os: ubuntu-latest
+            python-version: '3.9'
+            django-version: '4.2.9'
+          - os: ubuntu-latest
+            python-version: '3.10'
+            django-version: '4.2.9'
+          - os: ubuntu-latest
+            python-version: '3.11'
+            django-version: '4.2.9'
+          - os: ubuntu-latest
+            python-version: '3.12'
+            django-version: '4.2.9'
+    services:
+      postgres:
+        image: postgres
+        env:
+          POSTGRES_USER: postgres
+          POSTGRES_PASSWORD: postgres
+        options: >-
+          --health-cmd pg_isready
+          --health-interval 10s
+          --health-timeout 5s
+          --health-retries 5
+        ports:
+          - 5432:5432
+    steps:
+      - name: 'Set up Python ${{ matrix.python-version }}'
+        uses: actions/setup-python@v5
+        with:
+          python-version: '${{ matrix.python-version }}'
+      - uses: actions/checkout@v4
+      - run: pip install django==${{ matrix.django-version}}
+      - run: pip install psycopg[binary]
+      - run: pip install gevent
+      - run: python setup.py -q install
+      - run: python runtests_psycopg3.py

README.md

@@ -9,46 +9,54 @@ django-db-geventpool
 
 Another DB pool using gevent for PostgreSQL DB.
 
-If **gevent** is not installed, the pool will use **eventlet** as fallback.
+
+psycopg3
+---------
+
+Django, since 4.2, supports psycopg3. One of the advantages is that gevent is supported without needing extra patches, just install the package
+
+```
+$ pip install psycopg[binary]
+```
+
 
 psycopg2
 --------
 
-django-db-geventpool requires psycopg2:
+If **gevent** is not installed, the pool will use **eventlet** as fallback.
 
 -   `psycopg2>=2.5.1` for CPython 2 and 3 (or
     [psycopg2-binary](https://pypi.org/project/psycopg2-binary/)---see
     [notes in the psycopg2 2.7.4
     release](http://initd.org/psycopg/articles/2018/02/08/psycopg-274-released/))
 -   `psycopg2cffi>=2.7` for PyPy
 
-Patch psycopg2
---------------
+    Patch psycopg2
+    --------------
 
-Before using the pool, psycopg2 must be patched with psycogreen, if you
-are using [gunicorn webserver](http://www.gunicorn.org/), a good place
-is the
-[post\_fork()](http://docs.gunicorn.org/en/latest/settings.html#post-fork)
-function at the config file:
+    Before using the pool, psycopg2 must be patched with psycogreen, if you
+    are using [gunicorn webserver](http://www.gunicorn.org/), a good place
+    is the
+    [post\_fork()](http://docs.gunicorn.org/en/latest/settings.html#post-fork)
+    function at the config file:
 
-``` {.python}
-from psycogreen.gevent import patch_psycopg     # use this if you use gevent workers
-from psycogreen.eventlet import patch_psycopg   # use this if you use eventlet workers
+    ``` {.python}
+    from psycogreen.gevent import patch_psycopg     # use this if you use gevent workers
+    from psycogreen.eventlet import patch_psycopg   # use this if you use eventlet workers
 
-def post_fork(server, worker):
-    patch_psycopg()
-    worker.log.info("Made Psycopg2 Green")
-```
+    def post_fork(server, worker):
+        patch_psycopg()
+        worker.log.info("Made Psycopg2 Green")
+    ```
 
 Settings
 --------
 
-> -
->
->     Set *ENGINE* in your database settings to:
+> - Set *ENGINE* in your database settings to:
 >
->     :   -   *\'django\_db\_geventpool.backends.postgresql\_psycopg2\'*
->         -   For postgis: *\'django\_db\_geventpool.backends.postgis\'*
+>         - For psycopg3: 'django_db_geventpool.backends.postgresql_psycopg3'
+>         - For psycopg2: 'django_db_geventpool.backends.postgresql_psycopg2'
+>         - For postgis: 'django_db_geventpool.backends.postgis'
 >
 > -   Add *MAX\_CONNS* to *OPTIONS* to set the maximun number of
 >     connections allowed to database (default=4)
@@ -64,7 +72,7 @@ Settings
 ``` {.python}
 DATABASES = {
     'default': {
-        'ENGINE': 'django_db_geventpool.backends.postgresql_psycopg2',
+        'ENGINE': 'django_db_geventpool.backends.postgresql_psycopg',
         'NAME': 'db',
         'USER': 'postgres',
         'PASSWORD': 'postgres',

django_db_geventpool/backends/base.py

@@ -0,0 +1,93 @@
+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 = {}
+connection_pools_lock = Semaphore(value=1)
+
+
+class DatabaseWrapperMixin(object):
+    pool_class = None
+    creation_class = DatabaseCreation
+    INTRANS = None
+
+    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 connection_pools:
+                self._pool = self.pool_class(**self.get_connection_params())
+                connection_pools[self.alias] = self._pool
+            else:
+                self._pool = 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(
+                "psycopg2 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.closeall()
+        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 connection_pools.values():
+            pool.closeall()
+
+    def set_clean(self):
+        if self.in_atomic_block:
+            self.closed_in_transaction = True
+            self.needs_rollback = True

django_db_geventpool/backends/postgresql_psycopg2/creation.py → django_db_geventpool/backends/creation.py

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

django_db_geventpool/backends/postgresql_psycopg2/psycopg2_pool.py → django_db_geventpool/backends/pool.py

@@ -1,13 +1,11 @@
-# coding=utf-8
-
 # 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 sys
 import weakref
-logger = logging.getLogger('django.geventpool')
+
+logger = logging.getLogger("django.geventpool")
 
 try:
     from gevent import queue
@@ -16,27 +14,11 @@
     from eventlet import queue
     from ...utils import NullContextRLock as RLock
 
-try:
-    from psycopg2 import connect, DatabaseError
-    import psycopg2.extras
-except ImportError as e:
-    from django.core.exceptions import ImproperlyConfigured
-    raise ImproperlyConfigured("Error loading psycopg2 module: %s" % e)
-
-if sys.version_info[0] >= 3:
-    integer_types = int,
-else:
-    import __builtin__
-    integer_types = int, __builtin__.long
-
 
-class DatabaseConnectionPool(object):
-    def __init__(self, maxsize=100, reuse=100):
-        if not isinstance(maxsize, integer_types):
-            raise TypeError('Expected integer, got %r' % (maxsize,))
-        if not isinstance(reuse, integer_types):
-            raise TypeError('Expected integer, got %r' % (reuse,))
+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
@@ -61,7 +43,7 @@ def get(self):
                 # check connection is still valid
                 self.check_usable(conn)
                 logger.debug("DB connection reused")
-            except DatabaseError:
+            except self.DBERROR:
                 logger.debug("DB connection was closed, creating a new one")
                 conn = None
         except queue.Empty:
@@ -100,24 +82,3 @@ def closeall(self):
                 self._conns.discard(conn)
 
         logger.debug("DB connections all closed")
-
-
-class PostgresConnectionPool(DatabaseConnectionPool):
-    def __init__(self, *args, **kwargs):
-        self.connect = kwargs.pop('connect', connect)
-        self.connection = None
-        maxsize = kwargs.pop('MAX_CONNS', 4)
-        reuse = kwargs.pop('REUSE_CONNS', maxsize)
-        self.args = args
-        self.kwargs = kwargs
-        super(PostgresConnectionPool, self).__init__(maxsize, reuse)
-
-    def create_connection(self):
-        conn = self.connect(*self.args, **self.kwargs)
-        # set correct encoding
-        conn.set_client_encoding('UTF8')
-        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')

django_db_geventpool/backends/postgis/base.py

@@ -1,8 +1,13 @@
-# coding=utf-8
-
-from django.contrib.gis.db.backends.postgis.base import DatabaseWrapper as OriginalDatabaseWrapper
-
-from django_db_geventpool.backends.postgresql_psycopg2.base import DatabaseWrapperMixin
+from django.contrib.gis.db.backends.postgis.base import (
+    DatabaseWrapper as OriginalDatabaseWrapper,
+)
+
+try:
+    import psycopg  # noqa
+    from ..postgresql_psycopg3.base import DatabaseWrapperMixin
+except ImportError:
+    # fallback to psycopg3
+    from ..postgresql_psycopg2.base import DatabaseWrapperMixin
 
 
 class DatabaseWrapper(DatabaseWrapperMixin, OriginalDatabaseWrapper):