├── tests ├── __init__.py ├── backends │ ├── __init__.py │ └── django_postgrespool2 │ │ ├── __init__.py │ │ ├── test_creation.py │ │ └── tests.py ├── test_base.py ├── models.py ├── test_settings.py └── test_django.py ├── django_postgrespool2 ├── postgis │ ├── __init__.py │ └── base.py ├── psycopg2 │ ├── __init__.py │ └── base.py ├── __init__.py └── base.py ├── setup.cfg ├── MANIFEST.in ├── docker-compose.dev.yml ├── .gitignore ├── .github └── workflows │ ├── deploy.yml │ └── test.yml ├── LICENSE ├── manage.py ├── tox.ini ├── setup.py ├── CHANGELOG.md └── README.md /tests/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /tests/backends/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /django_postgrespool2/postgis/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /django_postgrespool2/psycopg2/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /setup.cfg: -------------------------------------------------------------------------------- 1 | [bdist_wheel] 2 | universal=1 3 | -------------------------------------------------------------------------------- /tests/backends/django_postgrespool2/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /MANIFEST.in: -------------------------------------------------------------------------------- 1 | include LICENSE 2 | include README.rst -------------------------------------------------------------------------------- /django_postgrespool2/psycopg2/base.py: -------------------------------------------------------------------------------- 1 | from ..base import * 2 | -------------------------------------------------------------------------------- /django_postgrespool2/__init__.py: -------------------------------------------------------------------------------- 1 | __version__ = "2.0.5" 2 | __author__ = 'lcd1232' 3 | -------------------------------------------------------------------------------- /tests/test_base.py: -------------------------------------------------------------------------------- 1 | import unittest 2 | 3 | 4 | class Test(unittest.TestCase): 5 | pass 6 | 7 | 8 | if __name__ == '__main__': 9 | unittest.main() 10 | -------------------------------------------------------------------------------- /docker-compose.dev.yml: -------------------------------------------------------------------------------- 1 | services: 2 | postgres_backend: 3 | image: postgres:alpine 4 | ports: 5 | - 5432:5432 6 | environment: 7 | POSTGRES_DB: pool 8 | POSTGRES_USER: pool 9 | POSTGRES_PASSWORD: pool 10 | -------------------------------------------------------------------------------- /django_postgrespool2/postgis/base.py: -------------------------------------------------------------------------------- 1 | from django.contrib.gis.db.backends.postgis.base import * 2 | from django.contrib.gis.db.backends.postgis.base import DatabaseWrapper as Psycopg2DatabaseWrapper 3 | from django.contrib.gis.db.backends.postgis.creation import DatabaseCreation as Psycopg2DatabaseCreation 4 | from ..base import * 5 | -------------------------------------------------------------------------------- /tests/models.py: -------------------------------------------------------------------------------- 1 | from django.db import models 2 | from django.utils import timezone 3 | 4 | try: 5 | from django.db.models import JSONField 6 | except ImportError: 7 | from django.contrib.postgres.fields import JSONField 8 | 9 | 10 | class DogModel(models.Model): 11 | 12 | name = models.CharField(max_length=200) 13 | created = models.DateTimeField(default=timezone.now) 14 | attrs = JSONField(default=dict) 15 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | *.py[co] 2 | 3 | # Packages 4 | *.egg 5 | *.egg-info 6 | dist 7 | build 8 | eggs 9 | parts 10 | bin 11 | var 12 | sdist 13 | develop-eggs 14 | .installed.cfg 15 | 16 | # Installer logs 17 | pip-log.txt 18 | 19 | # Unit test / coverage reports 20 | .coverage 21 | .tox 22 | 23 | #Translations 24 | *.mo 25 | 26 | #Mr Developer 27 | .mr.developer.cfg 28 | 29 | # Sublime Codeintel 30 | .codeintel/ 31 | 32 | # virtualenv 33 | venv/ 34 | -------------------------------------------------------------------------------- /tests/test_settings.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | os.environ.setdefault("DJANGO_DB_HOST", "localhost") 4 | 5 | SITE_ID = 1 6 | 7 | DATABASES = { 8 | "default": { 9 | "ENGINE": "django_postgrespool2", 10 | "NAME": "pool", 11 | "USER": "pool", 12 | "PASSWORD": "pool", 13 | "HOST": os.environ["DJANGO_DB_HOST"], 14 | } 15 | } 16 | 17 | DATABASE_POOL_ARGS = { 18 | "echo": True 19 | } 20 | 21 | INSTALLED_APPS = [ 22 | "django.contrib.auth", 23 | "django.contrib.contenttypes", 24 | "django.contrib.sessions", 25 | "tests", 26 | ] 27 | 28 | SECRET_KEY = "local" 29 | -------------------------------------------------------------------------------- /.github/workflows/deploy.yml: -------------------------------------------------------------------------------- 1 | name: Deploy package 2 | 3 | on: 4 | - push 5 | 6 | jobs: 7 | build: 8 | runs-on: ubuntu-latest 9 | steps: 10 | - uses: actions/checkout@v1 11 | - name: Set up Python 3.7 12 | uses: actions/setup-python@v2 13 | with: 14 | python-version: 3.7 15 | - name: Install dependencies 16 | run: | 17 | python -m pip install --upgrade pip 18 | - name: Create dist 19 | run: | 20 | python setup.py sdist 21 | - name: Publish package 22 | if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags') 23 | uses: pypa/gh-action-pypi-publish@release/v1 24 | with: 25 | user: __token__ 26 | password: ${{ secrets.PYPI_API_TOKEN }} 27 | -------------------------------------------------------------------------------- /tests/test_django.py: -------------------------------------------------------------------------------- 1 | from django.test import TestCase, override_settings 2 | from django.utils import timezone 3 | from tests.models import DogModel 4 | 5 | 6 | class TestPool(TestCase): 7 | 8 | def test_simple_request(self): 9 | DogModel.objects.create(name='wow') 10 | DogModel.objects.create(name='gog') 11 | 12 | @override_settings(USE_TZ=True, TIME_ZONE='Asia/Hong_Kong') 13 | def test_datetime_timezone(self): 14 | dog = DogModel.objects.create(name="wow", created=timezone.now()) 15 | self.assertEqual(dog.created.tzname(), "UTC") 16 | dog = DogModel.objects.get(name="wow") 17 | self.assertEqual(dog.created.tzname(), "UTC") 18 | 19 | def test_with_json(self): 20 | DogModel.objects.create(name="wow", attrs={"color": "pink"}) 21 | dog = DogModel.objects.get(name="wow") 22 | self.assertEqual(dog.attrs, {"color": "pink"}) 23 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Copyright (c) 2012 Kenneth Reitz 2 | 3 | Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: 4 | 5 | The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. 6 | 7 | THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. -------------------------------------------------------------------------------- /manage.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | import os 3 | import sys 4 | 5 | if __name__ == "__main__": 6 | os.environ.setdefault("DJANGO_SETTINGS_MODULE", "tests.test_settings") 7 | 8 | try: 9 | from django.core.management import execute_from_command_line 10 | except ImportError: 11 | # The above import may fail for some other reason. Ensure that the 12 | # issue is really that Django is missing to avoid masking other 13 | # exceptions on Python 2. 14 | try: 15 | import django # noqa 16 | except ImportError: 17 | raise ImportError( 18 | "Couldn't import Django. Are you sure it's installed and " 19 | "available on your PYTHONPATH environment variable? Did you " 20 | "forget to activate a virtual environment?" 21 | ) 22 | 23 | raise 24 | 25 | # This allows easy placement of apps within the interior 26 | # website directory. 27 | current_path = os.path.dirname(os.path.abspath(__file__)) 28 | sys.path.append(os.path.join(current_path, "tests")) 29 | 30 | execute_from_command_line(sys.argv) 31 | -------------------------------------------------------------------------------- /tests/backends/django_postgrespool2/test_creation.py: -------------------------------------------------------------------------------- 1 | """ 2 | Tests for django test runner 3 | """ 4 | import unittest 5 | try: 6 | from unittest import mock 7 | except ImportError: 8 | import mock 9 | 10 | 11 | from django import db 12 | from django.conf import settings 13 | from django.test.runner import DiscoverRunner 14 | 15 | 16 | class SetupDatabasesTests(unittest.TestCase): 17 | 18 | def setUp(self): 19 | self.runner_instance = DiscoverRunner(verbosity=0) 20 | 21 | def test_destroy_test_db_restores_db_name(self): 22 | tested_connections = db.ConnectionHandler({ 23 | 'default': { 24 | 'ENGINE': settings.DATABASES[db.DEFAULT_DB_ALIAS]["ENGINE"], 25 | 'NAME': 'xxx_test_database', 26 | }, 27 | }) 28 | # Using the real current name as old_name to not mess with the test suite. 29 | old_name = settings.DATABASES[db.DEFAULT_DB_ALIAS]["NAME"] 30 | with mock.patch('django.db.connections', new=tested_connections): 31 | tested_connections['default'].creation.destroy_test_db(old_name, verbosity=0, keepdb=True) 32 | self.assertEqual(tested_connections['default'].settings_dict["NAME"], old_name) 33 | -------------------------------------------------------------------------------- /tox.ini: -------------------------------------------------------------------------------- 1 | [tox] 2 | skipsdist = True 3 | usedevelop = True 4 | minversion = 1.8 5 | envlist = 6 | py{3.5,3.6,3.7,3.8,3.9}-dj2.2 7 | py{3.6,3.7,3.8,3.9}-dj3.0 8 | py{3.6,3.7,3.8,3.9}-dj3.1 9 | py{3.6,3.7,3.8,3.9}-djmaster 10 | 11 | [testenv] 12 | passenv = DJANGO_DB_HOST 13 | basepython = 14 | py3.5: python3 15 | py3.6: python3 16 | py3.7: python3 17 | py3.8: python3 18 | py3.9: python3 19 | pypy: pypy 20 | usedevelop = true 21 | setenv = 22 | PYTHONPATH = {toxinidir} 23 | DJANGO_SETTINGS_MODULE=tests.test_settings 24 | deps = 25 | py{3.5,3.6,3.7,3.8,3.9,pypy}: coverage 26 | psycopg2-binary 27 | dj2.2: https://github.com/django/django/archive/stable/2.2.x.tar.gz#egg=django 28 | dj3.0: https://github.com/django/django/archive/stable/3.0.x.tar.gz#egg=django 29 | dj3.1: https://github.com/django/django/archive/stable/3.1.x.tar.gz#egg=django 30 | dj3.2: https://github.com/django/django/archive/stable/3.2.x.tar.gz#egg=django 31 | dj4.0: https://github.com/django/django/archive/stable/4.0.x.tar.gz#egg=django 32 | djmaster: https://github.com/django/django/archive/master.tar.gz#egg=django 33 | 34 | commands = 35 | {envbindir}/python -Wd manage.py test -v2 {posargs:tests} 36 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup, find_packages 2 | from django_postgrespool2 import __version__, __author__ 3 | 4 | required = [ 5 | "sqlalchemy>=1.1", 6 | "django>=2.2", 7 | ] 8 | 9 | setup( 10 | name="django-postgrespool2", 11 | version=__version__, 12 | description="PostgreSQL connection pooling for Django.", 13 | long_description=open("README.md", encoding="utf-8").read(), 14 | long_description_content_type="text/markdown", 15 | author=__author__, 16 | author_email="malexey1984@gmail.com", 17 | url="https://github.com/lcd1232/django-postgrespool2", 18 | packages=find_packages(), 19 | python_requires=">=3.6", 20 | install_requires=required, 21 | license="MIT", 22 | classifiers=( 23 | "Development Status :: 5 - Production/Stable", 24 | "Intended Audience :: Developers", 25 | "Natural Language :: English", 26 | "License :: OSI Approved :: MIT License", 27 | "Programming Language :: Python", 28 | "Programming Language :: Python :: 3 :: Only", 29 | "Programming Language :: Python :: 3.6", 30 | "Programming Language :: Python :: 3.7", 31 | "Programming Language :: Python :: 3.8", 32 | "Programming Language :: Python :: 3.9", 33 | "Framework :: Django", 34 | "Framework :: Django :: 2.2", 35 | "Framework :: Django :: 3.0", 36 | "Framework :: Django :: 3.1", 37 | "Topic :: Database", 38 | ), 39 | keywords=["postgresql", "django", "pool", "pgbouncer",] 40 | ) 41 | -------------------------------------------------------------------------------- /.github/workflows/test.yml: -------------------------------------------------------------------------------- 1 | name: Test package 2 | 3 | on: 4 | - push 5 | 6 | jobs: 7 | build: 8 | runs-on: ubuntu-latest 9 | services: 10 | # Label used to access the service container 11 | postgres: 12 | # Docker Hub image 13 | image: postgres:12 14 | # Provide the password for postgres 15 | env: 16 | POSTGRES_USER: pool 17 | POSTGRES_DB: pool 18 | POSTGRES_PASSWORD: pool 19 | # Set health checks to wait until postgres has started 20 | ports: 21 | - 5432:5432 22 | options: >- 23 | --health-cmd pg_isready 24 | --health-interval 10s 25 | --health-timeout 5s 26 | --health-retries 5 27 | strategy: 28 | matrix: 29 | python-version: 30 | - "3.6" 31 | - "3.7" 32 | - "3.8" 33 | - "3.9" 34 | django-version: 35 | - "2.2" 36 | - "3.0" 37 | - "3.1" 38 | - "3.2" 39 | - "4.0" 40 | exclude: 41 | # django 4.0 supports python >= 3.8 42 | - python-version: "3.6" 43 | django-version: "4.0" 44 | - python-version: "3.7" 45 | django-version: "4.0" 46 | steps: 47 | - uses: actions/checkout@v1 48 | - name: Set up Python ${{ matrix.python-version }} 49 | uses: actions/setup-python@v2 50 | with: 51 | python-version: ${{ matrix.python-version }} 52 | - name: Install dependencies 53 | run: | 54 | python -m pip install --upgrade pip 55 | pip install tox tox-gh-actions 56 | - name: Test with tox 57 | run: tox -e py${{ matrix.python-version }}-dj${{ matrix.django-version }} 58 | -------------------------------------------------------------------------------- /CHANGELOG.md: -------------------------------------------------------------------------------- 1 | # Changelog 2 | All notable changes to this project will be documented in this file. 3 | ## 2.0.5 - 2022-10-03 4 | ### Fixed 5 | - Disable psycopg2 parsing JSONField value in Django 3.1+ ([#29](https://github.com/lcd1232/django-postgrespool2/pull/29)) 6 | 7 | ## 2.0.3 - 2022-07-17 8 | ### Fixed 9 | - Fixed timezone aware ([#27](https://github.com/lcd1232/django-postgrespool2/pull/27)) 10 | 11 | ## 2.0.1 - 2021-02-28 12 | ### Added 13 | - Support django 2.2, 3.0 and 3.1 14 | - Support python 3.6, 3.7, 3.8 and 3.9 15 | ### Removed 16 | - Remove support django 1.8, 1.9, 2.0, 2.1 17 | - Remove support python 3.4 and 3.5 18 | 19 | ## 1.0.1 - 2020-03-15 20 | ### Fixed 21 | - Fixed installation on python 3.7, thanks [chickahoona](https://github.com/lcd1232/django-postgrespool2/pull/16) 22 | 23 | ## 1.0.0 - 2019-09-09 24 | ### Added 25 | - New setting `DATABASE_POOL_CLASS`, thanks [mozartilize](https://github.com/mozartilize) 26 | ## Changed 27 | - Rewrite internal logic of library, thanks [mozartilize](https://github.com/mozartilize) 28 | ### Removed 29 | - Remove support python 3.3 30 | - Remove support django 1.7 31 | ### Fixed 32 | - Add missed backend `django_postgrespool2.postgis` 33 | 34 | ## 0.2.0 - 2018-02-13 35 | ### Added 36 | - Now you can specify backend for engine. Example: 37 | ```python 38 | DATABASES = { 39 | 'default': { 40 | 'ENGINE': 'django_postgrespool2.psycopg2', 41 | 'HOST': 'localhost', 42 | 'PORT': '5432', 43 | 'USER': 'test', 44 | 'PASSWORD': 'test', 45 | } 46 | } 47 | ``` 48 | Available backends: **django_postgrespool2.psycopg2**, **django_postgrespool2.postgis**. By default it using psycopg2 as backend engine. 49 | ### Changed 50 | - Update to Django 2.0 51 | - Fix error with pre_ping option 52 | - Fix error with __version__ 53 | 54 | ## 0.1.0 - 2017-09-14 55 | - Initial release 56 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | ![PyPI - Django Version](https://img.shields.io/pypi/djversions/django-postgrespool2) 2 | [![PyPI - License](https://img.shields.io/pypi/l/django-postgrespool2)](https://github.com/lcd1232/django-postgrespool2/blob/master/LICENSE) 3 | [![PyPI](https://img.shields.io/pypi/v/django-postgrespool2)](https://pypi.org/project/django-postgrespool2/) 4 | 5 | # Django-PostgresPool2 6 | This is simple PostgreSQL connection pooling for Django. You can use it as an alternative for [PgBouncer](https://www.pgbouncer.org/). 7 | This is a fork of the original [django-postgrespool](https://github.com/kennethreitz/django-postgrespool). 8 | 9 | ## Installation 10 | 11 | Installing Django-PostgresPool2 is simple, with pip: 12 | ```bash 13 | $ pip install django-postgrespool2 14 | ``` 15 | 16 | ## Usage 17 | 18 | Using Django-PostgresPool2 is simple, just set `django_postgrespool2` as your connection engine: 19 | ```python 20 | DATABASES = { 21 | "default": { 22 | "ENGINE": "django_postgrespool2", 23 | "NAME": "yourdb", 24 | "USER": "user", 25 | "PASSWORD": "some_password", 26 | "HOST": "localhost", 27 | } 28 | } 29 | ``` 30 | If you're using the [environ](https://github.com/joke2k/django-environ) module: 31 | ```python 32 | import environ 33 | 34 | env = environ.Env() 35 | 36 | DATABASES = {"default": env.db("DATABASE_URL", engine="django_postgrespool2")} 37 | ``` 38 | Everything should work as expected. 39 | 40 | Configuration 41 | ------------- 42 | 43 | Optionally, you can provide pool class to construct the pool (default `sqlalchemy.pool.QueuePool`) or additional options to pass to SQLAlchemy's pool creation. 44 | List of possible values `DATABASE_POOL_CLASS` is [here](https://docs.sqlalchemy.org/en/14/core/pooling.html#api-documentation-available-pool-implementations) 45 | ```python 46 | DATABASE_POOL_CLASS = 'sqlalchemy.pool.QueuePool' 47 | 48 | DATABASE_POOL_ARGS = { 49 | 'max_overflow': 10, 50 | 'pool_size': 5, 51 | 'recycle': 300, 52 | } 53 | ``` 54 | Here's a basic explanation of two of these options: 55 | 56 | - **pool_size** – The *minimum* number of connections to maintain in the pool. 57 | - **max_overflow** – The maximum *overflow* size of the pool. This is not the maximum size of the pool. 58 | - **recycle** - Number of seconds between connection recycling, which means upon checkout, if this timeout is surpassed the connection will be closed and replaced with a newly opened connection. 59 | 60 | The total number of "sleeping" connections the pool will allow is `pool_size`. The total simultaneous connections the pool will allow is `pool_size + max_overflow`. 61 | 62 | As an example, databases in the [Heroku Postgres](https://www.heroku.com/postgres) starter tier have a maximum connection limit of 20. In that case your `pool_size` and `max_overflow`, when combined, should not exceed 20. 63 | 64 | Check out the official [SQLAlchemy Connection Pooling](http://docs.sqlalchemy.org/en/latest/core/pooling.html#sqlalchemy.pool.QueuePool.__init__) docs to learn more about the optoins that can be defined in `DATABASE_POOL_ARGS`. 65 | -------------------------------------------------------------------------------- /tests/backends/django_postgrespool2/tests.py: -------------------------------------------------------------------------------- 1 | import unittest 2 | 3 | try: 4 | from unittest import mock 5 | except ImportError: 6 | import mock 7 | import warnings 8 | 9 | from django.db import DatabaseError, connection 10 | from django.test import TestCase 11 | from django.utils import version 12 | 13 | 14 | @unittest.skipUnless(connection.vendor == 'postgresql', 'PostgreSQL tests') 15 | class Tests(TestCase): 16 | 17 | def test_nodb_connection(self): 18 | """ 19 | Test that the _nodb_connection property fallbacks to the default connection 20 | database when access to the 'postgres' database is not granted. 21 | """ 22 | def mocked_connect(self): 23 | if self.settings_dict['NAME'] is None: 24 | raise DatabaseError() 25 | return '' 26 | 27 | if version.get_version_tuple(version.get_version()) >= (3, 1, 0): 28 | # I don't know what to test here 29 | pass 30 | else: 31 | nodb_conn = connection._nodb_connection 32 | self.assertIsNone(nodb_conn.settings_dict['NAME']) 33 | # Now assume the 'postgres' db isn't available 34 | with warnings.catch_warnings(record=True) as w: 35 | with mock.patch('django.db.backends.base.base.BaseDatabaseWrapper.connect', 36 | side_effect=mocked_connect, autospec=True): 37 | warnings.simplefilter('always', RuntimeWarning) 38 | nodb_conn = connection._nodb_connection 39 | self.assertIsNotNone(nodb_conn.settings_dict['NAME']) 40 | self.assertEqual(nodb_conn.settings_dict['NAME'], connection.settings_dict['NAME']) 41 | # Check a RuntimeWarning nas been emitted 42 | self.assertEqual(len(w), 1) 43 | self.assertEqual(w[0].message.__class__, RuntimeWarning) 44 | 45 | def test_connect_and_rollback(self): 46 | """ 47 | PostgreSQL shouldn't roll back SET TIME ZONE, even if the first 48 | transaction is rolled back (#17062). 49 | """ 50 | new_connection = connection.copy() 51 | try: 52 | # Ensure the database default time zone is different than 53 | # the time zone in new_connection.settings_dict. We can 54 | # get the default time zone by reset & show. 55 | with new_connection.cursor() as cursor: 56 | cursor.execute("RESET TIMEZONE") 57 | cursor.execute("SHOW TIMEZONE") 58 | db_default_tz = cursor.fetchone()[0] 59 | new_tz = 'Europe/Paris' if db_default_tz == 'UTC' else 'UTC' 60 | new_connection.close() 61 | 62 | if hasattr(new_connection, 'timezone_name'): # django 1.8 63 | # Invalidate timezone name cache, because the setting_changed 64 | # handler cannot know about new_connection. 65 | del new_connection.timezone_name 66 | 67 | # Fetch a new connection with the new_tz as default 68 | # time zone, run a query and rollback. 69 | with self.settings(TIME_ZONE=new_tz): 70 | new_connection.set_autocommit(False) 71 | new_connection.rollback() 72 | 73 | # Now let's see if the rollback rolled back the SET TIME ZONE. 74 | with new_connection.cursor() as cursor: 75 | cursor.execute("SHOW TIMEZONE") 76 | tz = cursor.fetchone()[0] 77 | self.assertEqual(new_tz, tz) 78 | 79 | finally: 80 | new_connection.dispose() 81 | 82 | def test_connect_non_autocommit(self): 83 | """ 84 | The connection wrapper shouldn't believe that autocommit is enabled 85 | after setting the time zone when AUTOCOMMIT is False (#21452). 86 | """ 87 | new_connection = connection.copy() 88 | new_connection.settings_dict['AUTOCOMMIT'] = False 89 | 90 | try: 91 | # Open a database connection. 92 | new_connection.cursor() 93 | self.assertFalse(new_connection.get_autocommit()) 94 | finally: 95 | new_connection.dispose() 96 | 97 | def test_connect_isolation_level(self): 98 | """ 99 | The transaction level can be configured with 100 | DATABASES ['OPTIONS']['isolation_level']. 101 | """ 102 | import psycopg2 103 | from psycopg2.extensions import ( 104 | ISOLATION_LEVEL_READ_COMMITTED as read_committed, 105 | ISOLATION_LEVEL_SERIALIZABLE as serializable, 106 | ) 107 | # Since this is a django.test.TestCase, a transaction is in progress 108 | # and the isolation level isn't reported as 0. This test assumes that 109 | # PostgreSQL is configured with the default isolation level. 110 | 111 | # Check the level on the psycopg2 connection, not the Django wrapper. 112 | default_level = read_committed if psycopg2.__version__ < '2.7' else None 113 | self.assertEqual(connection.connection.isolation_level, default_level) 114 | 115 | new_connection = connection.copy() 116 | new_connection.settings_dict['OPTIONS']['isolation_level'] = serializable 117 | try: 118 | # Start a transaction so the isolation level isn't reported as 0. 119 | new_connection.set_autocommit(False) 120 | # Check the level on the psycopg2 connection, not the Django wrapper. 121 | self.assertEqual(new_connection.connection.isolation_level, serializable) 122 | finally: 123 | new_connection.dispose() 124 | 125 | def test_connect_no_is_usable_checks(self): 126 | new_connection = connection.copy() 127 | with mock.patch.object(new_connection, 'is_usable') as is_usable: 128 | new_connection.connect() 129 | is_usable.assert_not_called() 130 | new_connection.dispose() 131 | -------------------------------------------------------------------------------- /django_postgrespool2/base.py: -------------------------------------------------------------------------------- 1 | import datetime 2 | import logging 3 | from functools import partial 4 | from importlib import import_module 5 | 6 | from django.conf import settings 7 | from django.db.backends.postgresql.base import ( 8 | psycopg2, 9 | Database, 10 | DatabaseWrapper as Psycopg2DatabaseWrapper, 11 | ) 12 | from django.db.backends.postgresql.creation import DatabaseCreation as Psycopg2DatabaseCreation 13 | from django.dispatch import Signal 14 | from django.utils import version 15 | try: 16 | from django.utils.asyncio import async_unsafe 17 | except ImportError: 18 | # dummy decorator 19 | def async_unsafe(func): 20 | return func 21 | try: 22 | # django 2.2 23 | from django.db.backends.postgresql.utils import utc_tzinfo_factory 24 | from django.utils.timezone import utc 25 | 26 | def utc_tzinfo_factory(offset): 27 | zero = 0 28 | # psycopg>=2.9 sends offset as timedelta 29 | if isinstance(offset, datetime.timedelta): 30 | zero = datetime.timedelta() 31 | if offset != zero: 32 | raise AssertionError("database connection isn't set to UTC") 33 | return utc 34 | 35 | except ImportError: 36 | utc_tzinfo_factory = None 37 | from sqlalchemy import event 38 | from sqlalchemy.dialects import postgresql 39 | from sqlalchemy.pool import manage 40 | 41 | # DATABASE_POOL_ARGS should be something like: 42 | # {'max_overflow':10, 'pool_size':5, 'recycle':300} 43 | pool_args = {'max_overflow': 10, 'pool_size': 5, 'recycle': 300} 44 | pool_args.update(getattr(settings, 'DATABASE_POOL_ARGS', {})) 45 | dialect = postgresql.dialect(dbapi=psycopg2) 46 | pool_args['dialect'] = dialect 47 | 48 | POOL_CLS = getattr(settings, 'DATABASE_POOL_CLASS', 'sqlalchemy.pool.QueuePool') 49 | pool_module_name, pool_cls_name = POOL_CLS.rsplit('.', 1) 50 | pool_cls = getattr(import_module(pool_module_name), pool_cls_name) 51 | pool_args['poolclass'] = pool_cls 52 | 53 | db_pool = manage(Database, **pool_args) 54 | pool_disposed = Signal() 55 | 56 | log = logging.getLogger('z.pool') 57 | 58 | django_version = version.get_version_tuple(version.get_version()) 59 | 60 | def _log(message, *args): 61 | log.debug(message) 62 | 63 | 64 | # Only hook up the listeners if we are in debug mode. 65 | if settings.DEBUG: 66 | event.listen(pool_cls, 'checkout', partial(_log, 'retrieved from pool')) 67 | event.listen(pool_cls, 'checkin', partial(_log, 'returned to pool')) 68 | event.listen(pool_cls, 'connect', partial(_log, 'new connection')) 69 | 70 | 71 | class DatabaseCreation(Psycopg2DatabaseCreation): 72 | def _clone_test_db(self, *args, **kw): 73 | self.connection.dispose() 74 | super(DatabaseCreation, self)._clone_test_db(*args, **kw) 75 | 76 | def create_test_db(self, *args, **kw): 77 | self.connection.dispose() 78 | super(DatabaseCreation, self).create_test_db(*args, **kw) 79 | 80 | def destroy_test_db(self, *args, **kw): 81 | """Ensure connection pool is disposed before trying to drop database. 82 | """ 83 | self.connection.dispose() 84 | super(DatabaseCreation, self).destroy_test_db(*args, **kw) 85 | 86 | 87 | class DatabaseWrapper(Psycopg2DatabaseWrapper): 88 | """SQLAlchemy FTW.""" 89 | 90 | def __init__(self, *args, **kwargs): 91 | super(DatabaseWrapper, self).__init__(*args, **kwargs) 92 | self._pool = None 93 | self._pool_connection = None 94 | self.creation = DatabaseCreation(self) 95 | 96 | @property 97 | def pool(self): 98 | return self._pool 99 | 100 | def _close(self): 101 | if self._pool_connection is not None: 102 | if not self.is_usable(): 103 | self._pool_connection.invalidate() 104 | with self.wrap_database_errors: 105 | return self._pool_connection.close() 106 | 107 | @async_unsafe 108 | def create_cursor(self, name=None): 109 | if name: 110 | # In autocommit mode, the cursor will be used outside of a 111 | # transaction, hence use a holdable cursor. 112 | cursor = self._pool_connection.cursor( 113 | name, scrollable=False, withhold=self.connection.autocommit) 114 | else: 115 | cursor = self._pool_connection.cursor() 116 | cursor.tzinfo_factory = self.tzinfo_factory if settings.USE_TZ else None 117 | return cursor 118 | 119 | def tzinfo_factory(self, offset): 120 | if utc_tzinfo_factory: 121 | # for Django 2.2 122 | return utc_tzinfo_factory(offset) 123 | return self.timezone 124 | 125 | def dispose(self): 126 | """ 127 | Dispose of the pool for this instance, closing all connections. 128 | """ 129 | self.close() 130 | self._pool_connection = None 131 | # _DBProxy.dispose doesn't actually call dispose on the pool 132 | if self.pool: 133 | self.pool.dispose() 134 | self._pool = None 135 | conn_params = self.get_connection_params() 136 | db_pool.dispose(**conn_params) 137 | pool_disposed.send(sender=self.__class__, connection=self) 138 | 139 | @async_unsafe 140 | def get_new_connection(self, conn_params): 141 | if not self._pool: 142 | self._pool = db_pool.get_pool(**conn_params) 143 | # get new connection through pool, not creating a new one outside. 144 | self._pool_connection = self.pool.connect() 145 | c = self._pool_connection.connection # dbapi connection 146 | 147 | options = self.settings_dict['OPTIONS'] 148 | try: 149 | self.isolation_level = options['isolation_level'] 150 | except KeyError: 151 | self.isolation_level = c.isolation_level 152 | else: 153 | # Set the isolation level to the value from OPTIONS. 154 | if self.isolation_level != c.isolation_level: 155 | c.set_session(isolation_level=self.isolation_level) 156 | 157 | if django_version >= (3, 1, 1): 158 | psycopg2.extras.register_default_jsonb(conn_or_curs=c, loads=lambda x: x) 159 | return c 160 | 161 | def is_usable(self): 162 | if not self.connection: 163 | return False 164 | return self.connection.closed == 0 165 | --------------------------------------------------------------------------------