Initial commit

Co-authored-by: Ovidiu Sabou <[email protected]>

.gitignore

@@ -0,0 +1,29 @@
+# Ignore virtual environments
+env/
+.env/
+
+# Ignore Python byte code cache
+*.pyc
+__pycache__
+.cache
+
+# Ignore coverage reports
+.coverage
+htmlcov
+
+# Ignore build results
+*.egg-info/
+dist/
+
+# Ignore stupid .DS_Store
+.DS_Store
+
+# Ignore benchmark results
+.benchmarks/
+
+# Ignore temporary tox environments
+.tox/
+.pytest_cache/
+
+# Ignore PyCharm / IntelliJ files
+.idea/

README.md

@@ -0,0 +1,15 @@
+# Django Database Connection Retrier
+
+Automatically try to re-establish Django database connections when they fail due to DNS lookup errors.
+
+## Installation
+1. Install the package from PyPi:
+
+        $ pip install django-db-connection-retrier
+
+2. Add `dbconnectionretrier` to your `INSTALLED_APPS`:
+
+        INSTALLED_APPS = [
+            'dbconnectionretrier',
+            ...
+        ]

dbconnectionretrier/apps.py

@@ -0,0 +1,16 @@
+from django.apps import AppConfig
+
+from .patch import patch_ensure_connection
+
+
+class DBConnectionRetrierConfig(AppConfig):
+    """Django app configuration that hooks.
+
+    :see:BaseDatabaseWrapper.ensure_connection to automatically retry
+    connection failures.
+    """
+
+    name = "dbconnectionretrier"
+
+    def ready(self):
+        return patch_ensure_connection()

dbconnectionretrier/ensure_connection.py

@@ -0,0 +1,39 @@
+import logging
+
+from time import sleep
+
+import aspectlib
+
+from django.db import OperationalError
+
+LOGGER = logging.getLogger(__name__)
+
+
+@aspectlib.Aspect
+def ensure_connection(instance):
+    """Aspect that tries to ensure a DB connection by retrying.
+
+    Catches name resolution errors, by filtering on OperationalError
+    exceptions that contain name resolution error messages
+
+    Useful in case the DNS resolution is shaky, as in the case
+    of the Heroku environment
+    """
+    max_tries = 3
+    for trial in range(0, max_tries):
+        try:
+            result = yield aspectlib.Proceed
+            yield aspectlib.Return(result)
+        except OperationalError as error:
+            message = str(error)
+            if "could not translate host name" not in message:
+                raise
+            if trial == max_tries - 1:
+                raise
+            sleep(2 ** trial)
+
+            LOGGER.warning(
+                "Database connection lost, retrying trial %d: %s",
+                trial,
+                message,
+            )

dbconnectionretrier/patch.py

@@ -0,0 +1,35 @@
+from contextlib import contextmanager
+
+import aspectlib
+
+from django.db.backends.base.base import BaseDatabaseWrapper
+
+from .ensure_connection import ensure_connection
+
+
+def patch_ensure_connection():
+    """Monkey patch BaseDatabaseWrapper.ensure_connection.
+
+    See the doc of the patch function for details about what it does.
+
+    Rturns:
+        An object representing the patch. see:
+        https://python-aspectlib.readthedocs.io/en/latest/testing.html?highlight=rollback#spy-mock-toolkit-record-mock-decorators
+    """
+
+    return aspectlib.weave(
+        BaseDatabaseWrapper.ensure_connection, ensure_connection
+    )
+
+
+@contextmanager
+def patch_ensure_connection_contextual():
+    """Monkey patch BaseDatabaseWrapper.ensure_connection for the duration of
+    the context."""
+
+    patch = patch_ensure_connection()
+
+    try:
+        yield patch
+    finally:
+        patch.rollback()

pyproject.toml

@@ -0,0 +1,2 @@
+[tool.black]
+line-length = 80

requirements.txt

@@ -0,0 +1,16 @@
+-e .
+
+psycopg2==2.8.2
+coverage==4.5.3
+pytest==4.5.0
+pytest-django==3.4.8
+pytest-cov==2.7.1
+tox==3.11.1
+sl-docformatter==1.2
+black==19.3b0
+flake8==3.7.7
+pycodestyle==2.5.0
+autoflake==1.3
+autopep8==1.4.4
+isort==4.3.20
+dj-database-url==0.5.0

setup.cfg

@@ -0,0 +1,15 @@
+[flake8]
+ignore = E252,E501
+exclude = env,.tox,.git,config/settings,*/migrations/*
+
+[pycodestyle]
+ignore = E501
+exclude=env,.tox,.git
+
+[isort]
+line_length=80
+multi_line_output=3
+lines_between_types=1
+include_trailing_comma=True
+not_skip=__init__.py
+known_standard_library=dataclasses

setup.py

@@ -0,0 +1,133 @@
+import distutils.cmd
+import subprocess
+
+from setuptools import find_packages, setup
+
+
+class BaseCommand(distutils.cmd.Command):
+    user_options = []
+
+    def initialize_options(self):
+        pass
+
+    def finalize_options(self):
+        pass
+
+
+def create_command(text, commands):
+    """Creates a custom setup.py command."""
+
+    class CustomCommand(BaseCommand):
+        description = text
+
+        def run(self):
+            for cmd in commands:
+                subprocess.check_call(cmd)
+
+    return CustomCommand
+
+
+setup(
+    name="django-db-connection-retrier",
+    version="1.0",
+    packages=find_packages(),
+    include_package_data=True,
+    license="MIT License",
+    description="Automatically ty re-establishing the Django database connection when it gets lost.",
+    url="https://github.com/SectorLabs/django-db-connection-retrier",
+    author="Sector Labs",
+    author_email="[email protected]",
+    keywords=["django", "postgres", "extra", "hstore", "ltree"],
+    install_requires=["aspectlib==1.4.2", "Django==2.2.5"],
+    classifiers=[
+        "Environment :: Web Environment",
+        "Framework :: Django",
+        "Intended Audience :: Developers",
+        "License :: OSI Approved :: MIT License",
+        "Operating System :: OS Independent",
+        "Programming Language :: Python",
+        "Programming Language :: Python :: 3.5",
+    ],
+    cmdclass={
+        "lint": create_command(
+            "Lints the code",
+            [
+                ["flake8", "setup.py", "dbconnectionretrier", "tests"],
+                ["pycodestyle", "setup.py", "dbconnectionretrier", "tests"],
+            ],
+        ),
+        "lint_fix": create_command(
+            "Lints the code",
+            [
+                [
+                    "autoflake",
+                    "--remove-all-unused-imports",
+                    "-i",
+                    "-r",
+                    "setup.py",
+                    "dbconnectionretrier",
+                    "tests",
+                ],
+                [
+                    "autopep8",
+                    "-i",
+                    "-r",
+                    "setup.py",
+                    "dbconnectionretrier",
+                    "tests",
+                ],
+            ],
+        ),
+        "format": create_command(
+            "Formats the code",
+            [["black", "setup.py", "dbconnectionretrier", "tests"]],
+        ),
+        "format_verify": create_command(
+            "Checks if the code is auto-formatted",
+            [["black", "--check", "setup.py", "dbconnectionretrier", "tests"]],
+        ),
+        "format_docstrings": create_command(
+            "Auto-formats doc strings", [["docformatter", "-r", "-i", "."]]
+        ),
+        "format_docstrings_verify": create_command(
+            "Verifies that doc strings are properly formatted",
+            [["docformatter", "-r", "-c", "."]],
+        ),
+        "sort_imports": create_command(
+            "Automatically sorts imports",
+            [
+                ["isort", "setup.py"],
+                ["isort", "-rc", "dbconnectionretrier"],
+                ["isort", "-rc", "tests"],
+            ],
+        ),
+        "sort_imports_verify": create_command(
+            "Verifies all imports are properly sorted.",
+            [
+                ["isort", "-c", "setup.py"],
+                ["isort", "-c", "-rc", "dbconnectionretrier"],
+                ["isort", "-c", "-rc", "tests"],
+            ],
+        ),
+        "fix": create_command(
+            "Automatically format code and fix linting errors",
+            [
+                ["python", "setup.py", "format"],
+                ["python", "setup.py", "format_docstrings"],
+                ["python", "setup.py", "sort_imports"],
+                ["python", "setup.py", "lint_fix"],
+            ],
+        ),
+        "verify": create_command(
+            "Verifies whether the code is auto-formatted and has no linting errors",
+            [
+                [
+                    ["python", "setup.py", "format_verify"],
+                    ["python", "setup.py", "format_docstrings_verify"],
+                    ["python", "setup.py", "sort_imports_verify"],
+                    ["python", "setup.py", "lint"],
+                ]
+            ],
+        ),
+    },
+)

tests/conftest.py

@@ -0,0 +1,28 @@
+import dj_database_url
+import pytest
+
+
+def set_defaults(db):
+    """Set the mandatory settings on the connection."""
+    db.setdefault("TIME_ZONE", None)
+    db.setdefault("CONN_MAX_AGE", None)
+    db.setdefault("OPTIONS", {})
+    db.setdefault("ATOMIC_REQUESTS", True)
+
+
+@pytest.fixture()
+def unknown_host():
+    """DB connection with a broken host name."""
+    db = dj_database_url.config(
+        default="postgres://this_domain_should_not_exist/test_strat"
+    )
+    set_defaults(db)
+    return db
+
+
+@pytest.fixture()
+def unknown_db():
+    """DB connection with a broken database name."""
+    db = dj_database_url.config(default="postgres:///this_db_should_not_exist")
+    set_defaults(db)
+    return db

tests/test_app_config.py

@@ -0,0 +1,30 @@
+from io import StringIO
+from logging import StreamHandler
+
+import pytest
+
+from django.db import OperationalError
+from django.db.utils import load_backend
+
+from dbconnectionretrier.apps import DBConnectionRetrierConfig
+from dbconnectionretrier.ensure_connection import LOGGER
+
+
+def test_app_config_install_patch(unknown_host):
+    """Tests whether the Django app config properly installs the retrier and
+    retries connection failures."""
+
+    try:
+        patch = DBConnectionRetrierConfig.ready(None)
+
+        io = StringIO()
+        LOGGER.addHandler(StreamHandler(io))
+        with pytest.raises(OperationalError):
+            backend = load_backend(unknown_host["ENGINE"])
+            conn = backend.DatabaseWrapper(unknown_host, "unknown_host")
+            conn.ensure_connection()
+
+        # test that retrying has taken place (DNS errors might have been fixed)
+        assert "trial 0" in io.getvalue()
+    finally:
+        patch.rollback()