Support both psqlextra and PostGis backends at the same time

[#181799346]

Author: sewi-cpan

psqlextra/backend/base.py

@@ -22,7 +22,20 @@ class DatabaseWrapper(base_impl.backend()):
     introspection_class = PostgresIntrospection
     ops_class = PostgresOperations
 
-    def prepare_database(self):
+    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'
+
+
+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,4 +1,5 @@
 import importlib
+import os
 
 from django.conf import settings
 from django.core.exceptions import ImproperlyConfigured
@@ -23,7 +24,7 @@ def backend():
     base_class_name = getattr(
         settings,
         "POSTGRES_EXTRA_DB_BACKEND_BASE",
-        "django.db.backends.postgresql",
+        os.environ.get("POSTGRES_EXTRA_DB_BACKEND_BASE") or "django.db.backends.postgresql",
     )
 
     base_class_module = importlib.import_module(base_class_name + ".base")

psqlextra/backend/operations.py

@@ -1,9 +1,11 @@
+from importlib import import_module
+
 from psqlextra.compiler import (
-    PostgresAggregateCompiler,
-    PostgresCompiler,
-    PostgresDeleteCompiler,
-    PostgresInsertCompiler,
-    PostgresUpdateCompiler,
+    SQLAggregateCompiler,
+    SQLCompiler,
+    SQLDeleteCompiler,
+    SQLInsertCompiler,
+    SQLUpdateCompiler,
 )
 
 from . import base_impl
@@ -12,14 +14,6 @@
 class PostgresOperations(base_impl.operations()):
     """Simple operations specific to PostgreSQL."""
 
-    compiler_map = {
-        "SQLCompiler": PostgresCompiler,
-        "SQLInsertCompiler": PostgresInsertCompiler,
-        "SQLUpdateCompiler": PostgresUpdateCompiler,
-        "SQLDeleteCompiler": PostgresDeleteCompiler,
-        "SQLAggregateCompiler": PostgresAggregateCompiler,
-    }
-
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
@@ -28,9 +22,8 @@ def __init__(self, *args, **kwargs):
     def compiler(self, compiler_name: str):
         """Gets the SQL compiler with the specified name."""
 
-        postgres_compiler = self.compiler_map.get(compiler_name)
-        if postgres_compiler:
-            return postgres_compiler
+        if self._cache is None:
+            self._cache = import_module('psqlextra.compiler')
 
-        # Let Django try to find the compiler. Better run without caller comment than break
-        return super().compiler(compiler_name)
+        # 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))

psqlextra/compiler.py

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

psqlextra/sql.py

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