Django PostgreSQL backend with connection logging

base.py

import logging
import os
import time

from django.contrib.gis.db.backends.postgis.base import (
    DatabaseWrapper as PostGISDatabaseWrapper,
)

logger = logging.getLogger(__name__)

TRUE_VALUES = ["True", "true", "1"]


class DatabaseWrapper(PostGISDatabaseWrapper):
    """Workarounds to debug database "connection already closed" errors."""

    def ensure_connection(self):
        """Drop the connection if it is closed."""
        logger.debug("Ensuring connection")
        drop_closed_connection = (
            os.getenv("DATABASE_DROP_CLOSED_CONNECTION", "False") in TRUE_VALUES
        )
        if drop_closed_connection:
            if (
                self.connection is not None
                and hasattr(self.connection, "closed")
                and self.connection.closed
            ):
                logger.error("Attempted to use a closed connection. Reconnecting.")
                self.connection = None

        super().ensure_connection()

    def close_if_unusable_or_obsolete(self):
        """Log every attempt to close the connection."""
        if self.connection is not None:
            connection_pid = self.connection.get_backend_pid()
            logger.info("Checking if connection %s is usable", connection_pid)
            self.health_check_done = False
            # If the application didn't restore the original autocommit setting,
            # don't take chances, drop the connection.
            if self.get_autocommit() != self.settings_dict["AUTOCOMMIT"]:
                logger.info("Closing connection due to autocommit")
                self.close()
                return

            # If an exception other than DataError or IntegrityError occurred
            # since the last commit / rollback, check if the connection works.
            if self.errors_occurred:
                if self.is_usable():
                    logger.info("Connection has errors but is usable")
                    self.errors_occurred = False
                    self.health_check_done = True
                else:
                    logger.info("Closing connection due to the occurred errors")
                    self.close()
                    return

            if self.close_at is not None and time.monotonic() >= self.close_at:
                logger.info("Closing connection due to age")
                self.close()
                return

    def _close(self):
        """Log connection PID when closing."""
        if self.connection is not None:
            connection_pid = self.connection.get_backend_pid()
            logger.info("Closed connection %s", connection_pid)
            super()._close()

    def get_new_connection(self, conn_params):
        """Log connection PID when opening."""
        connection = super().get_new_connection(conn_params)
        connection_pid = connection.get_backend_pid()
        logger.info("Opened connection %s", connection_pid)
        return connection