Merge pull request #179 from SectorLabs/181799346-postgis-compatibility

Support both psqlextra and PostGis backends at the same time

psqlextra/backend/base.py

@@ -22,6 +22,35 @@ class DatabaseWrapper(base_impl.backend()):
     introspection_class = PostgresIntrospection
     ops_class = PostgresOperations
 
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        # 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.introspection, 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):
         """Ran to prepare the configured database.
 

psqlextra/backend/base_impl.py

@@ -2,14 +2,15 @@
 
 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
@@ -19,6 +20,10 @@ 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,
@@ -49,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():
@@ -59,7 +81,7 @@ def schema_editor():
     this.
     """
 
-    return backend().SchemaEditorClass
+    return base_backend_instance().SchemaEditorClass
 
 
 def introspection():
@@ -69,7 +91,7 @@ def introspection():
     for this.
     """
 
-    return backend().introspection_class
+    return base_backend_instance().introspection.__class__
 
 
 def operations():
@@ -79,4 +101,4 @@ def operations():
     this.
     """
 
-    return backend().ops_class
+    return base_backend_instance().ops.__class__

psqlextra/backend/operations.py

@@ -1,9 +1,9 @@
 from psqlextra.compiler import (
-    PostgresAggregateCompiler,
-    PostgresCompiler,
-    PostgresDeleteCompiler,
-    PostgresInsertCompiler,
-    PostgresUpdateCompiler,
+    SQLAggregateCompiler,
+    SQLCompiler,
+    SQLDeleteCompiler,
+    SQLInsertCompiler,
+    SQLUpdateCompiler,
 )
 
 from . import base_impl
@@ -12,25 +12,12 @@
 class PostgresOperations(base_impl.operations()):
     """Simple operations specific to PostgreSQL."""
 
-    compiler_map = {
-        "SQLCompiler": PostgresCompiler,
-        "SQLInsertCompiler": PostgresInsertCompiler,
-        "SQLUpdateCompiler": PostgresUpdateCompiler,
-        "SQLDeleteCompiler": PostgresDeleteCompiler,
-        "SQLAggregateCompiler": PostgresAggregateCompiler,
-    }
+    compiler_module = "psqlextra.compiler"
 
-    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."""
-
-        postgres_compiler = self.compiler_map.get(compiler_name)
-        if postgres_compiler:
-            return postgres_compiler
-
-        # Let Django try to find the compiler. Better run without caller comment than break
-        return 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,
-    SQLCompiler,
-    SQLDeleteCompiler,
-    SQLInsertCompiler,
-    SQLUpdateCompiler,
-)
+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 PostgresCompiler(SQLCompiler):
+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 PostgresDeleteCompiler(SQLDeleteCompiler):
+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 PostgresAggregateCompiler(SQLAggregateCompiler):
+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 PostgresUpdateCompiler(SQLUpdateCompiler):
+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 PostgresInsertCompiler(SQLInsertCompiler):
+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(SQLInsertCompiler):
+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 SQLInsertCompiler.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, PostgresUpdateCompiler
+from .compiler import PostgresInsertOnConflictCompiler
+from .compiler import SQLUpdateCompiler as PostgresUpdateCompiler
 from .expressions import HStoreColumn
 from .fields import HStoreField
 from .types import ConflictAction