Inspect base database back-end instance to figure out base classes

The PostGIS back-end doesn't set `DatabaseWrapper.ops_class` properly
and instead initializes its own `Operations` class in the constructor.

This results in django-postgres-extra subclassing the wrong base class.

As a solution, we create an instance and inspect what class is actually
used for `DatabaseWrapper.ops_class` and `DatabaseWrapper.introspection_class`
and subclass from that.

[#181799346]

Author: Photonios

psqlextra/backend/base.py

@@ -24,18 +24,34 @@ class DatabaseWrapper(base_impl.backend()):
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        if not isinstance(self.ops, PostgresOperations):
-            # PostGis replaces the ops object instead of setting the ops_class attribute
-            if self.ops.compiler_module != 'django.db.models.sql.compiler':
-                raise NotImplementedError(
-                    f'''The Django ops object has been replaced by {self.ops} and a custom compiler module {self.ops.compiler_module} has been set.
-Replacing both at the same time is incompatible with psqlextra. '''
-                )
-            self.ops._compiler_cache = None
-            self.ops.compiler_module = 'psqlextra.compiler'
 
+        # Some base back-ends such as the PostGIS back-end don't properly
+        # set `ops_class` and `introspection_class` and initialize these
+        # classes themselves.
+        #
+        # This can lead to broken functionality. We fix this automatically.
+
+        if not isinstance(self.ops, self.introspection_class):
+            self.introspection = self.introspection_class(self)
+
+        if not isinstance(self.ops, self.ops_class):
+            self.ops = self.ops_class(self)
+
+        for expected_compiler_class in self.ops.compiler_classes:
+            compiler_class = self.ops.compiler(expected_compiler_class.__name__)
+
+            if not issubclass(compiler_class, expected_compiler_class):
+                logger.warning(
+                    "Compiler '%s.%s' is not properly deriving from '%s.%s'."
+                    % (
+                        compiler_class.__module__,
+                        compiler_class.__name__,
+                        expected_compiler_class.__module__,
+                        expected_compiler_class.__name__,
+                    )
+                )
 
-def prepare_database(self):
+    def prepare_database(self):
         """Ran to prepare the configured database.
 
         This is where we enable the `hstore` extension if it wasn't

psqlextra/backend/base_impl.py

@@ -1,16 +1,16 @@
 import importlib
-import os
 
 from django.conf import settings
 from django.core.exceptions import ImproperlyConfigured
+from django.db import DEFAULT_DB_ALIAS, connections
 
 from django.db.backends.postgresql.base import (  # isort:skip
     DatabaseWrapper as Psycopg2DatabaseWrapper,
 )
 
 
-def backend():
-    """Gets the base class for the custom database back-end.
+def base_backend_instance():
+    """Gets an instance of the base class for the custom database back-end.
 
     This should be the Django PostgreSQL back-end. However,
     some people are already using a custom back-end from
@@ -20,11 +20,15 @@ def backend():
     As long as the specified base eventually also has
     the PostgreSQL back-end as a base, then everything should
     work as intended.
+
+    We create an instance to inspect what classes to subclass
+    because not all back-ends set properties such as `ops_class`
+    properly. The PostGIS back-end is a good example.
     """
     base_class_name = getattr(
         settings,
         "POSTGRES_EXTRA_DB_BACKEND_BASE",
-        os.environ.get("POSTGRES_EXTRA_DB_BACKEND_BASE") or "django.db.backends.postgresql",
+        "django.db.backends.postgresql",
     )
 
     base_class_module = importlib.import_module(base_class_name + ".base")
@@ -50,7 +54,24 @@ def backend():
             % base_class_name
         )
 
-    return base_class
+    base_instance = base_class(connections.databases[DEFAULT_DB_ALIAS])
+    if base_instance.connection:
+        raise ImproperlyConfigured(
+            (
+                "'%s' establishes a connection during initialization."
+                " This is not expected and can lead to more connections"
+                " being established than neccesarry."
+            )
+            % base_class_name
+        )
+
+    return base_instance
+
+
+def backend():
+    """Gets the base class for the database back-end."""
+
+    return base_backend_instance().__class__
 
 
 def schema_editor():
@@ -60,7 +81,7 @@ def schema_editor():
     this.
     """
 
-    return backend().SchemaEditorClass
+    return base_backend_instance().SchemaEditorClass
 
 
 def introspection():
@@ -70,7 +91,7 @@ def introspection():
     for this.
     """
 
-    return backend().introspection_class
+    return base_backend_instance().introspection.__class__
 
 
 def operations():
@@ -80,4 +101,4 @@ def operations():
     this.
     """
 
-    return backend().ops_class
+    return base_backend_instance().ops.__class__

psqlextra/backend/operations.py

@@ -1,5 +1,3 @@
-from importlib import import_module
-
 from psqlextra.compiler import (
     SQLAggregateCompiler,
     SQLCompiler,
@@ -14,16 +12,12 @@
 class PostgresOperations(base_impl.operations()):
     """Simple operations specific to PostgreSQL."""
 
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-
-        self._compiler_cache = None
-
-    def compiler(self, compiler_name: str):
-        """Gets the SQL compiler with the specified name."""
-
-        if self._cache is None:
-            self._cache = import_module('psqlextra.compiler')
+    compiler_module = "psqlextra.compiler"
 
-        # Let any parent module try to find the compiler as fallback. Better run without caller comment than break
-        return getattr(self._cache, compiler_name, super().compiler(compiler_name))
+    compiler_classes = [
+        SQLCompiler,
+        SQLDeleteCompiler,
+        SQLAggregateCompiler,
+        SQLUpdateCompiler,
+        SQLInsertCompiler,
+    ]

psqlextra/compiler.py

@@ -11,13 +11,7 @@
 from django.core.exceptions import SuspiciousOperation
 from django.db.models import Expression, Model, Q
 from django.db.models.fields.related import RelatedField
-from django.db.models.sql.compiler import (
-    SQLAggregateCompiler as DjangoSQLAggregateCompiler,
-    SQLCompiler as DjangoSQLCompiler,
-    SQLDeleteCompiler as DjangoSQLDeleteCompiler,
-    SQLInsertCompiler as DjangoSQLInsertCompiler,
-    SQLUpdateCompiler as DjangoSQLUpdateCompiler,
-)
+from django.db.models.sql import compiler as django_compiler
 from django.db.utils import ProgrammingError
 
 from .expressions import HStoreValue
@@ -77,25 +71,25 @@ def append_caller_to_sql(sql):
         return sql
 
 
-class SQLCompiler(DjangoSQLCompiler):
+class SQLCompiler(django_compiler.SQLCompiler):
     def as_sql(self, *args, **kwargs):
         sql, params = super().as_sql(*args, **kwargs)
         return append_caller_to_sql(sql), params
 
 
-class SQLDeleteCompiler(DjangoSQLDeleteCompiler):
+class SQLDeleteCompiler(django_compiler.SQLDeleteCompiler):
     def as_sql(self, *args, **kwargs):
         sql, params = super().as_sql(*args, **kwargs)
         return append_caller_to_sql(sql), params
 
 
-class SQLAggregateCompiler(DjangoSQLAggregateCompiler):
+class SQLAggregateCompiler(django_compiler.SQLAggregateCompiler):
     def as_sql(self, *args, **kwargs):
         sql, params = super().as_sql(*args, **kwargs)
         return append_caller_to_sql(sql), params
 
 
-class SQLUpdateCompiler(DjangoSQLUpdateCompiler):
+class SQLUpdateCompiler(django_compiler.SQLUpdateCompiler):
     """Compiler for SQL UPDATE statements that allows us to use expressions
     inside HStore values.
 
@@ -152,7 +146,7 @@ def _does_dict_contain_expression(data: dict) -> bool:
         return False
 
 
-class SQLInsertCompiler(DjangoSQLInsertCompiler):
+class SQLInsertCompiler(django_compiler.SQLInsertCompiler):
     """Compiler for SQL INSERT statements."""
 
     def as_sql(self, *args, **kwargs):
@@ -165,7 +159,7 @@ def as_sql(self, *args, **kwargs):
         return queries
 
 
-class PostgresInsertOnConflictCompiler(DjangoSQLInsertCompiler):
+class PostgresInsertOnConflictCompiler(django_compiler.SQLInsertCompiler):
     """Compiler for SQL INSERT statements."""
 
     def __init__(self, *args, **kwargs):
@@ -407,7 +401,7 @@ def _format_field_value(self, field_name) -> str:
         if isinstance(field, RelatedField) and isinstance(value, Model):
             value = value.pk
 
-        return DjangoSQLInsertCompiler.prepare_value(
+        return django_compiler.SQLInsertCompiler.prepare_value(
             self,
             field,
             # Note: this deliberately doesn't use `pre_save_val` as we don't

psqlextra/sql.py

@@ -8,7 +8,8 @@
 from django.db.models import sql
 from django.db.models.constants import LOOKUP_SEP
 
-from .compiler import PostgresInsertOnConflictCompiler, SQLUpdateCompiler as PostgresUpdateCompiler
+from .compiler import PostgresInsertOnConflictCompiler
+from .compiler import SQLUpdateCompiler as PostgresUpdateCompiler
 from .expressions import HStoreColumn
 from .fields import HStoreField
 from .types import ConflictAction