diff --git a/django_postgresql_reconnect/backend/base.py b/django_postgresql_reconnect/backend/base.py index 5c6a73a..4fc8730 100644 --- a/django_postgresql_reconnect/backend/base.py +++ b/django_postgresql_reconnect/backend/base.py @@ -7,7 +7,7 @@ from django.db.backends.postgresql.base import DatabaseWrapper as PgDatabaseWrapper try: - from psycopg2 import InterfaceError + from psycopg2 import InterfaceError, OperationalError from psycopg2.extensions import STATUS_IN_TRANSACTION except ImportError as e: raise ImproperlyConfigured("Error loading psycopg2 module: %s" % e) @@ -22,9 +22,14 @@ class DatabaseWrapper(PgDatabaseWrapper): def create_cursor(self, name=None): try: return super().create_cursor(name) - except InterfaceError: + except (InterfaceError, OperationalError): if self.should_reconnect(): - self.reconnect() + try: + self.reconnect() + except OperationalError: + # reconnect() → connect() failed (e.g. transient TCP timeout). + # connection is now None; attempt one more fresh connect before giving up. + self.ensure_connection() return super().create_cursor(name) raise diff --git a/pyproject.toml b/pyproject.toml index fcf45d5..24a8497 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,6 +40,9 @@ testpaths = "tests" log_cli = true addopts = "--junitxml=tests/reports/out.xml --cov=django_postgresql_reconnect --cov-report xml:tests/reports/coverage.xml --cov-report html:tests/reports/cov_html" +[tool.coverage.run] +relative_files = true + [tool.coverage.report] omit = [ "*/migrations/*", diff --git a/tests/test_backend.py b/tests/test_backend.py index 2fea041..7ecd58c 100644 --- a/tests/test_backend.py +++ b/tests/test_backend.py @@ -2,7 +2,9 @@ import logging from copy import deepcopy +from unittest.mock import patch +import psycopg2 import pytest from django.db import InterfaceError, connections, transaction @@ -78,3 +80,59 @@ def test_should_reconnect_connection_not_initialized(): connection.connection = None assert not connection.should_reconnect() + + +@pytest.mark.django_db(transaction=True, databases=['default', 'sqlite']) +def test_reconnect_on_operational_error_from_cursor(caplog): + """OperationalError from cursor() (zombie connection) triggers reconnect.""" + from django.db.backends.postgresql.base import DatabaseWrapper as PgDatabaseWrapper + + connection = ConnectionHandler()['default'] + connection.connect() + + original_create_cursor = PgDatabaseWrapper.create_cursor + call_count = 0 + + def create_cursor_fails_once(self, name=None): + nonlocal call_count + call_count += 1 + if call_count == 1: + raise psycopg2.OperationalError('server closed the connection unexpectedly') + return original_create_cursor(self, name) + + with patch.object(PgDatabaseWrapper, 'create_cursor', create_cursor_fails_once): + cursor = connection.cursor() + + cursor.execute('SELECT 1') + assert cursor.fetchone() == (1,) + assert caplog.record_tuples == [ + ('django.db.backend', logging.WARNING, 'Reconnect to the database "default"'), + ] + + +@pytest.mark.django_db(transaction=True, databases=['default', 'sqlite']) +def test_reconnect_retries_when_connect_fails_once(caplog): + """If reconnect() → connect() fails transiently, a second attempt is made.""" + connection = ConnectionHandler()['default'] + connection.connect() + connection.connection.close() # simulate server-side disconnect (InterfaceError on cursor) + + original_connect = psycopg2.connect + attempt = 0 + + def connect_fails_once(*args, **kwargs): + nonlocal attempt + attempt += 1 + if attempt == 1: + raise psycopg2.OperationalError('could not connect to server: Connection timed out') + return original_connect(*args, **kwargs) + + with patch('psycopg2.connect', side_effect=connect_fails_once): + cursor = connection.cursor() + + cursor.execute('SELECT 1') + assert cursor.fetchone() == (1,) + assert attempt == 2 + assert caplog.record_tuples == [ + ('django.db.backend', logging.WARNING, 'Reconnect to the database "default"'), + ]