xdjango-postgres-extra 2.0.9rc5__py3-none-any.whl

psqlextra/__init__.py

@@ -0,0 +1,15 @@
+import django
+
+from ._version import __version__
+
+if django.VERSION < (3, 2):  # pragma: no cover
+    default_app_config = "psqlextra.apps.PostgresExtraAppConfig"
+
+    __all__ = [
+        "default_app_config",
+        "__version__",
+    ]
+else:
+    __all__ = [
+        "__version__",
+    ]

psqlextra/_version.py

@@ -0,0 +1 @@
+__version__ = "2.0.9rc5"

psqlextra/apps.py

@@ -0,0 +1,9 @@
+from django.apps import AppConfig
+
+
+class PostgresExtraAppConfig(AppConfig):
+    name = "psqlextra"
+    verbose_name = "PostgreSQL Extra"
+
+    def ready(self) -> None:
+        from .lookups import InValuesLookup  # noqa

psqlextra/backend/base.py

@@ -0,0 +1,123 @@
+import logging
+
+from typing import TYPE_CHECKING
+
+from django import VERSION
+from django.conf import settings
+from django.contrib.postgres.signals import (
+    get_hstore_oids,
+    register_type_handlers,
+)
+from django.db import ProgrammingError
+
+from . import base_impl
+from .introspection import PostgresIntrospection
+from .operations import PostgresOperations
+from .schema import PostgresSchemaEditor
+
+from django.db.backends.postgresql.base import (  # isort:skip
+    DatabaseWrapper as PostgresDatabaseWrapper,
+)
+
+
+logger = logging.getLogger(__name__)
+
+
+if TYPE_CHECKING:
+
+    class Wrapper(PostgresDatabaseWrapper):
+        pass
+
+else:
+    Wrapper = base_impl.backend()
+
+
+class DatabaseWrapper(Wrapper):
+    """Wraps the standard PostgreSQL database back-end.
+
+    Overrides the schema editor with our custom schema editor and makes
+    sure the `hstore` extension is enabled.
+    """
+
+    SchemaEditorClass = PostgresSchemaEditor  # type: ignore[assignment]
+    introspection_class = PostgresIntrospection
+    ops_class = PostgresOperations
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        if VERSION >= (5, 0):
+            return
+
+        # 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.
+
+        This is where we enable the `hstore` extension if it wasn't
+        enabled yet.
+        """
+
+        super().prepare_database()
+
+        setup_ext = getattr(
+            settings, "POSTGRES_EXTRA_AUTO_EXTENSION_SET_UP", True
+        )
+        if not setup_ext:
+            return False
+
+        with self.cursor() as cursor:
+            try:
+                cursor.execute("CREATE EXTENSION IF NOT EXISTS hstore")
+            except ProgrammingError:  # permission denied
+                logger.warning(
+                    'Failed to create "hstore" extension. '
+                    "Tables with hstore columns may fail to migrate. "
+                    "If hstore is needed, make sure you are connected "
+                    "to the database as a superuser "
+                    "or add the extension manually.",
+                    exc_info=True,
+                )
+                return
+
+        # Clear old (non-existent), stale oids.
+        get_hstore_oids.cache_clear()
+
+        # Verify that we (and Django) can find the OIDs
+        # for hstore.
+        oids, _ = get_hstore_oids(self.alias)
+        if not oids:
+            logger.warning(
+                '"hstore" extension was created, but we cannot find the oids'
+                "in the database. Something went wrong.",
+            )
+            return
+
+        # We must trigger Django into registering the type handlers now
+        # so that any subsequent code can properly use the newly
+        # registered types.
+        register_type_handlers(self)

psqlextra/backend/base_impl.py

@@ -0,0 +1,112 @@
+import importlib
+
+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 DatabaseWrapper
+from django.db.backends.postgresql.introspection import (  # type: ignore[import]
+    DatabaseIntrospection,
+)
+from django.db.backends.postgresql.operations import DatabaseOperations
+from django.db.backends.postgresql.schema import (  # type: ignore[import]
+    DatabaseSchemaEditor,
+)
+
+from django.db.backends.postgresql.base import (  # isort:skip
+    DatabaseWrapper as Psycopg2DatabaseWrapper,
+)
+
+
+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
+    another package. We are nice people and expose an option
+    that allows them to configure the back-end we base upon.
+
+    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",
+        "django.db.backends.postgresql",
+    )
+
+    base_class_module = importlib.import_module(base_class_name + ".base")
+    base_class = getattr(base_class_module, "DatabaseWrapper", None)
+
+    if not base_class:
+        raise ImproperlyConfigured(
+            (
+                "'%s' is not a valid database back-end."
+                " The module does not define a DatabaseWrapper class."
+                " Check the value of POSTGRES_EXTRA_DB_BACKEND_BASE."
+            )
+            % base_class_name
+        )
+
+    if isinstance(base_class, Psycopg2DatabaseWrapper):
+        raise ImproperlyConfigured(
+            (
+                "'%s' is not a valid database back-end."
+                " It does inherit from the PostgreSQL back-end."
+                " Check the value of POSTGRES_EXTRA_DB_BACKEND_BASE."
+            )
+            % base_class_name
+        )
+
+    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() -> DatabaseWrapper:
+    """Gets the base class for the database back-end."""
+
+    return base_backend_instance().__class__
+
+
+def schema_editor() -> DatabaseSchemaEditor:
+    """Gets the base class for the schema editor.
+
+    We have to use the configured base back-end's schema editor for
+    this.
+    """
+
+    return base_backend_instance().SchemaEditorClass
+
+
+def introspection() -> DatabaseIntrospection:
+    """Gets the base class for the introspection class.
+
+    We have to use the configured base back-end's introspection class
+    for this.
+    """
+
+    return base_backend_instance().introspection.__class__
+
+
+def operations() -> DatabaseOperations:
+    """Gets the base class for the operations class.
+
+    We have to use the configured base back-end's operations class for
+    this.
+    """
+
+    return base_backend_instance().ops.__class__

psqlextra/backend/introspection.py

@@ -0,0 +1,327 @@
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Dict, List, Optional, Tuple
+
+from django.db.backends.postgresql.introspection import (  # type: ignore[import]
+    DatabaseIntrospection,
+)
+
+from psqlextra.types import PostgresPartitioningMethod
+
+from . import base_impl
+
+PARTITIONING_STRATEGY_TO_METHOD = {
+    "r": PostgresPartitioningMethod.RANGE,
+    "l": PostgresPartitioningMethod.LIST,
+    "h": PostgresPartitioningMethod.HASH,
+}
+
+
+@dataclass
+class PostgresIntrospectedPartitionTable:
+    """Data container for information about a partition."""
+
+    name: str
+    full_name: str
+    comment: Optional[str]
+
+
+@dataclass
+class PostgresIntrospectedPartitonedTable:
+    """Data container for information about a partitioned table."""
+
+    name: str
+    method: PostgresPartitioningMethod
+    key: List[str]
+    partitions: List[PostgresIntrospectedPartitionTable]
+
+    def partition_by_name(
+        self, name: str
+    ) -> Optional[PostgresIntrospectedPartitionTable]:
+        """Finds the partition with the specified name."""
+
+        return next(
+            (
+                partition
+                for partition in self.partitions
+                if partition.name == name
+            ),
+            None,
+        )
+
+
+if TYPE_CHECKING:
+
+    class Introspection(DatabaseIntrospection):
+        pass
+
+else:
+    Introspection = base_impl.introspection()
+
+
+class PostgresIntrospection(Introspection):
+    """Adds introspection features specific to PostgreSQL."""
+
+    # TODO: This class is a mess, both here and in the
+    # the base.
+    #
+    # Some methods return untyped dicts, some named tuples,
+    # some flat lists of strings. It's horribly inconsistent.
+    #
+    # Most methods are poorly named. For example; `get_table_description`
+    # does not return a complete table description. It merely returns
+    # the columns.
+    #
+    # We do our best in this class to stay consistent with
+    # the base in Django by respecting its naming scheme
+    # and commonly used return types. Creating an API that
+    # matches the look&feel from the Django base class
+    # is more important than fixing those issues.
+
+    def get_partitioned_tables(
+        self, cursor
+    ) -> List[PostgresIntrospectedPartitonedTable]:
+        """Gets a list of partitioned tables."""
+
+        cursor.execute(
+            """
+            SELECT
+                pg_class.relname,
+                pg_partitioned_table.partstrat
+            FROM
+                pg_partitioned_table
+            JOIN
+                pg_class
+            ON
+                pg_class.oid = pg_partitioned_table.partrelid
+            ORDER BY
+                pg_partitioned_table.partrelid
+        """
+        )
+
+        return [
+            PostgresIntrospectedPartitonedTable(
+                name=row[0],
+                method=PARTITIONING_STRATEGY_TO_METHOD[row[1]],
+                key=self.get_partition_key(cursor, row[0]),
+                partitions=self.get_partitions(cursor, row[0]),
+            )
+            for row in cursor.fetchall()
+        ]
+
+    def get_partitioned_table(self, cursor, table_name: str):
+        """Gets a single partitioned table."""
+
+        return next(
+            (
+                table
+                for table in self.get_partitioned_tables(cursor)
+                if table.name == table_name
+            ),
+            None,
+        )
+
+    def get_partitions(
+        self, cursor, table_name
+    ) -> List[PostgresIntrospectedPartitionTable]:
+        """Gets a list of partitions belonging to the specified partitioned
+        table."""
+
+        sql = """
+            SELECT
+                child.relname,
+                pg_description.description
+            FROM pg_inherits
+            JOIN
+                pg_class parent
+            ON
+                pg_inherits.inhparent = parent.oid
+            JOIN
+                pg_class child
+            ON
+                pg_inherits.inhrelid = child.oid
+            JOIN
+                pg_namespace nmsp_parent
+            ON
+                nmsp_parent.oid = parent.relnamespace
+            JOIN
+                pg_namespace nmsp_child
+            ON
+                nmsp_child.oid = child.relnamespace
+            LEFT JOIN
+                pg_description
+            ON
+                pg_description.objoid = child.oid
+            WHERE
+                parent.relname = %s
+            ORDER BY
+                child.oid,
+                child.relname
+        """
+
+        cursor.execute(sql, (table_name,))
+
+        return [
+            PostgresIntrospectedPartitionTable(
+                name=row[0].replace(f"{table_name}_", ""),
+                full_name=row[0],
+                comment=row[1] or None,
+            )
+            for row in cursor.fetchall()
+        ]
+
+    def get_partition_key(self, cursor, table_name: str) -> List[str]:
+        """Gets the partition key for the specified partitioned table.
+
+        Returns:
+            A list of column names that are part of the
+            partition key.
+        """
+
+        sql = """
+            SELECT
+                col.column_name
+            FROM
+                (SELECT partrelid,
+                        partnatts,
+                        CASE partstrat
+                            WHEN 'l' THEN 'list'
+                            WHEN 'r' THEN 'range'
+                            WHEN 'h' THEN 'hash'
+                        END AS partition_strategy,
+                        Unnest(partattrs) column_index
+                 FROM pg_partitioned_table) pt
+            JOIN
+                pg_class par
+            ON par.oid = pt.partrelid
+            JOIN
+                information_schema.COLUMNS col
+            ON
+                col.table_schema = par.relnamespace :: regnamespace :: text
+                AND col.table_name = par.relname
+                AND ordinal_position = pt.column_index
+            WHERE
+                table_name = %s
+            ORDER BY
+                col.ordinal_position,
+                col.column_name
+        """
+
+        cursor.execute(sql, (table_name,))
+        return [row[0] for row in cursor.fetchall()]
+
+    def get_columns(self, cursor, table_name: str):
+        return self.get_table_description(cursor, table_name)
+
+    def get_schema_list(self, cursor) -> List[str]:
+        """A flat list of available schemas."""
+
+        cursor.execute(
+            """
+            SELECT
+                schema_name
+            FROM
+                information_schema.schemata
+            ORDER BY
+                schema_name,
+                catalog_name
+            """,
+            tuple(),
+        )
+
+        return [name for name, in cursor.fetchall()]
+
+    def get_constraints(self, cursor, table_name: str):
+        """Retrieve any constraints or keys (unique, pk, fk, check, index)
+        across one or more columns.
+
+        Also retrieve the definition of expression-based indexes.
+        """
+
+        constraints = super().get_constraints(cursor, table_name)
+
+        # standard Django implementation does not return the definition
+        # for indexes, only for constraints, let's patch that up
+        cursor.execute(
+            "SELECT indexname, indexdef FROM pg_indexes WHERE tablename = %s",
+            (table_name,),
+        )
+        for index_name, definition in cursor.fetchall():
+            # PostgreSQL 13 or older won't give a definition if the
+            # index is actually a primary key.
+            constraint = constraints.get(index_name)
+            if not constraint:
+                continue
+
+            if constraint.get("definition") is None:
+                constraint["definition"] = definition
+
+        return constraints
+
+    def get_table_locks(self, cursor) -> List[Tuple[str, str, str]]:
+        cursor.execute(
+            """
+            SELECT
+                n.nspname,
+                t.relname,
+                l.mode
+            FROM pg_locks l
+            INNER JOIN pg_class t ON t.oid = l.relation
+            INNER JOIN pg_namespace n ON n.oid = t.relnamespace
+            WHERE t.relnamespace >= 2200
+            ORDER BY n.nspname, t.relname, l.mode
+        """
+        )
+
+        return cursor.fetchall()
+
+    def get_storage_settings(self, cursor, table_name: str) -> Dict[str, str]:
+        sql = """
+            SELECT
+                unnest(c.reloptions || array(select 'toast.' || x from pg_catalog.unnest(tc.reloptions) x))
+            FROM
+                pg_catalog.pg_class c
+            LEFT JOIN
+                pg_catalog.pg_class tc ON (c.reltoastrelid = tc.oid)
+            LEFT JOIN
+                pg_catalog.pg_am am ON (c.relam = am.oid)
+            WHERE
+                c.relname::text = %s
+                AND pg_catalog.pg_table_is_visible(c.oid)
+        """
+
+        cursor.execute(sql, (table_name,))
+
+        storage_settings = {}
+        for row in cursor.fetchall():
+            # It's hard to believe, but storage settings are really
+            # represented as `key=value` strings in Postgres.
+            # See: https://www.postgresql.org/docs/current/catalog-pg-class.html
+            name, value = row[0].split("=")
+            storage_settings[name] = value
+
+        return storage_settings
+
+    def get_relations(self, cursor, table_name: str):
+        """Gets a dictionary {field_name: (field_name_other_table,
+        other_table)} representing all relations in the specified table.
+
+        This is overriden because the query in Django does not handle
+        relations between tables in different schemas properly.
+        """
+
+        cursor.execute(
+            """
+            SELECT a1.attname, c2.relname, a2.attname
+            FROM pg_constraint con
+            LEFT JOIN pg_class c1 ON con.conrelid = c1.oid
+            LEFT JOIN pg_class c2 ON con.confrelid = c2.oid
+            LEFT JOIN pg_attribute a1 ON c1.oid = a1.attrelid AND a1.attnum = con.conkey[1]
+            LEFT JOIN pg_attribute a2 ON c2.oid = a2.attrelid AND a2.attnum = con.confkey[1]
+            WHERE
+                con.conrelid = %s::regclass AND
+                con.contype = 'f' AND
+                pg_catalog.pg_table_is_visible(c1.oid)
+        """,
+            [table_name],
+        )
+        return {row[0]: (row[2], row[1]) for row in cursor.fetchall()}

psqlextra/backend/migrations/__init__.py

@@ -0,0 +1,3 @@
+from .patched_migrations import postgres_patched_migrations
+
+__all__ = ["postgres_patched_migrations"]

psqlextra/backend/migrations/operations/__init__.py

@@ -0,0 +1,33 @@
+from .add_default_partition import PostgresAddDefaultPartition
+from .add_hash_partition import PostgresAddHashPartition
+from .add_list_partition import PostgresAddListPartition
+from .add_range_partition import PostgresAddRangePartition
+from .apply_state import ApplyState
+from .create_materialized_view_model import PostgresCreateMaterializedViewModel
+from .create_partitioned_model import PostgresCreatePartitionedModel
+from .create_view_model import PostgresCreateViewModel
+from .delete_default_partition import PostgresDeleteDefaultPartition
+from .delete_hash_partition import PostgresDeleteHashPartition
+from .delete_list_partition import PostgresDeleteListPartition
+from .delete_materialized_view_model import PostgresDeleteMaterializedViewModel
+from .delete_partitioned_model import PostgresDeletePartitionedModel
+from .delete_range_partition import PostgresDeleteRangePartition
+from .delete_view_model import PostgresDeleteViewModel
+
+__all__ = [
+    "ApplyState",
+    "PostgresAddHashPartition",
+    "PostgresAddListPartition",
+    "PostgresAddRangePartition",
+    "PostgresAddDefaultPartition",
+    "PostgresDeleteDefaultPartition",
+    "PostgresDeleteHashPartition",
+    "PostgresDeleteListPartition",
+    "PostgresDeleteRangePartition",
+    "PostgresCreatePartitionedModel",
+    "PostgresDeletePartitionedModel",
+    "PostgresCreateViewModel",
+    "PostgresCreateMaterializedViewModel",
+    "PostgresDeleteViewModel",
+    "PostgresDeleteMaterializedViewModel",
+]

psqlextra/backend/migrations/operations/add_default_partition.py

@@ -0,0 +1,35 @@
+from psqlextra.backend.migrations.state import PostgresPartitionState
+
+from .partition import PostgresPartitionOperation
+
+
+class PostgresAddDefaultPartition(PostgresPartitionOperation):
+    """Adds a new default partition to a :see:PartitionedPostgresModel."""
+
+    def state_forwards(self, app_label, state):
+        model_state = state.models[(app_label, self.model_name_lower)]
+        model_state.add_partition(
+            PostgresPartitionState(
+                app_label=app_label, model_name=self.model_name, name=self.name
+            )
+        )
+
+        state.reload_model(app_label, self.model_name_lower)
+
+    def database_forwards(self, app_label, schema_editor, from_state, to_state):
+        model = to_state.apps.get_model(app_label, self.model_name)
+        if self.allow_migrate_model(schema_editor.connection.alias, model):
+            schema_editor.add_default_partition(model, self.name)
+
+    def database_backwards(
+        self, app_label, schema_editor, from_state, to_state
+    ):
+        model = from_state.apps.get_model(app_label, self.model_name)
+        if self.allow_migrate_model(schema_editor.connection.alias, model):
+            schema_editor.delete_partition(model, self.name)
+
+    def describe(self) -> str:
+        return "Creates default partition '%s' on %s" % (
+            self.name,
+            self.model_name,
+        )