Hash partitioning

docs/source/table_partitioning.rst

@@ -20,6 +20,7 @@ The following partitioning methods are available:
 
 * ``PARTITION BY RANGE``
 * ``PARTITION BY LIST``
+* ``PARTITION BY HASH``
 
 .. note::
 
@@ -41,6 +42,7 @@ Inherit your model from :class:`psqlextra.models.PostgresPartitionedModel` and d
 
 * Use :attr:`psqlextra.types.PostgresPartitioningMethod.RANGE` to ``PARTITION BY RANGE``
 * Use :attr:`psqlextra.types.PostgresPartitioningMethod.LIST` to ``PARTITION BY LIST``
+* Use :attr:`psqlextra.types.PostgresPartitioningMethod.HASH` to ``PARTITION BY HASH``
 
 .. code-block:: python
 

psqlextra/backend/introspection.py

@@ -5,6 +5,12 @@
 
 from . import base_impl
 
+PARTITIONING_STRATEGY_TO_METHOD = {
+    "r": PostgresPartitioningMethod.RANGE,
+    "l": PostgresPartitioningMethod.LIST,
+    "h": PostgresPartitioningMethod.HASH,
+}
+
 
 @dataclass
 class PostgresIntrospectedPartitionTable:
@@ -64,9 +70,7 @@ def get_partitioned_tables(
         return [
             PostgresIntrospectedPartitonedTable(
                 name=row[0],
-                method=PostgresPartitioningMethod.RANGE
-                if row[1] == "r"
-                else PostgresPartitioningMethod.LIST,
+                method=PARTITIONING_STRATEGY_TO_METHOD[row[1]],
                 key=self.get_partition_key(cursor, row[0]),
                 partitions=self.get_partitions(cursor, row[0]),
             )
@@ -148,6 +152,7 @@ def get_partition_key(self, cursor, table_name: str) -> List[str]:
                         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

psqlextra/backend/migrations/operations/__init__.py

@@ -1,11 +1,13 @@
 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
@@ -14,12 +16,14 @@
 
 __all__ = [
     "ApplyState",
-    "PostgresAddRangePartition",
+    "PostgresAddHashPartition",
     "PostgresAddListPartition",
+    "PostgresAddRangePartition",
     "PostgresAddDefaultPartition",
     "PostgresDeleteDefaultPartition",
-    "PostgresDeleteRangePartition",
+    "PostgresDeleteHashPartition",
     "PostgresDeleteListPartition",
+    "PostgresDeleteRangePartition",
     "PostgresCreatePartitionedModel",
     "PostgresDeletePartitionedModel",
     "PostgresCreateViewModel",

psqlextra/backend/migrations/operations/add_hash_partition.py

@@ -0,0 +1,74 @@
+from psqlextra.backend.migrations.state import PostgresHashPartitionState
+
+from .partition import PostgresPartitionOperation
+
+
+class PostgresAddHashPartition(PostgresPartitionOperation):
+    """Adds a new hash partition to a :see:PartitionedPostgresModel.
+
+    Each partition will hold the rows for which the hash value of the
+    partition key divided by the specified modulus will produce the
+    specified remainder.
+    """
+
+    def __init__(
+        self, model_name: str, name: str, modulus: int, remainder: int
+    ):
+        """Initializes new instance of :see:AddHashPartition.
+        Arguments:
+            model_name:
+                The name of the :see:PartitionedPostgresModel.
+
+            name:
+                The name to give to the new partition table.
+
+            modulus:
+                Integer value by which the key is divided.
+
+            remainder:
+                The remainder of the hash value when divided by modulus.
+        """
+
+        super().__init__(model_name, name)
+
+        self.modulus = modulus
+        self.remainder = remainder
+
+    def state_forwards(self, app_label, state):
+        model = state.models[(app_label, self.model_name_lower)]
+        model.add_partition(
+            PostgresHashPartitionState(
+                app_label=app_label,
+                model_name=self.model_name,
+                name=self.name,
+                modulus=self.modulus,
+                remainder=self.remainder,
+            )
+        )
+
+        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_hash_partition(
+                model, self.name, self.modulus, self.remainder
+            )
+
+    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 deconstruct(self):
+        name, args, kwargs = super().deconstruct()
+
+        kwargs["modulus"] = self.modulus
+        kwargs["remainder"] = self.remainder
+
+        return name, args, kwargs
+
+    def describe(self) -> str:
+        return "Creates hash partition %s on %s" % (self.name, self.model_name)

psqlextra/backend/migrations/operations/create_partitioned_model.py

@@ -69,3 +69,19 @@ def describe(self):
         description = super().describe()
         description = description.replace("model", "partitioned model")
         return description
+
+    def reduce(self, *args, **kwargs):
+        result = super().reduce(*args, **kwargs)
+
+        # replace CreateModel operation with PostgresCreatePartitionedModel
+        if isinstance(result, list) and result:
+            for i, op in enumerate(result):
+                if isinstance(op, CreateModel):
+                    _, args, kwargs = op.deconstruct()
+                    result[i] = PostgresCreatePartitionedModel(
+                        *args,
+                        **kwargs,
+                        partitioning_options=self.partitioning_options
+                    )
+
+        return result

psqlextra/backend/migrations/operations/delete_hash_partition.py

@@ -0,0 +1,29 @@
+from .delete_partition import PostgresDeletePartition
+
+
+class PostgresDeleteHashPartition(PostgresDeletePartition):
+    """Deletes a hash partition that's part of a.
+
+    :see:PartitionedPostgresModel.
+    """
+
+    def database_backwards(
+        self, app_label, schema_editor, from_state, to_state
+    ):
+        model = to_state.apps.get_model(app_label, self.model_name)
+        model_state = to_state.models[(app_label, self.model_name_lower)]
+
+        if self.allow_migrate_model(schema_editor.connection.alias, model):
+            partition_state = model_state.partitions[self.name]
+            schema_editor.add_hash_partition(
+                model,
+                partition_state.name,
+                partition_state.modulus,
+                partition_state.remainder,
+            )
+
+    def describe(self) -> str:
+        return "Deletes hash partition '%s' on %s" % (
+            self.name,
+            self.model_name,
+        )

psqlextra/backend/migrations/operations/partition.py

@@ -26,3 +26,7 @@ def state_forwards(self, *args, **kwargs):
 
     def state_backwards(self, *args, **kwargs):
         pass
+
+    def reduce(self, *args, **kwargs):
+        # PartitionOperation doesn't break migrations optimizations
+        return True

psqlextra/backend/migrations/patched_autodetector.py

@@ -8,7 +8,6 @@
     DeleteModel,
     RemoveField,
     RenameField,
-    RunSQL,
 )
 from django.db.migrations.autodetector import MigrationAutodetector
 from django.db.migrations.operations.base import Operation
@@ -19,6 +18,7 @@
     PostgresPartitionedModel,
     PostgresViewModel,
 )
+from psqlextra.types import PostgresPartitioningMethod
 
 from . import operations
 
@@ -140,11 +140,12 @@ def add_create_partitioned_model(
         partitioning_options = model._partitioning_meta.original_attrs
         _, args, kwargs = operation.deconstruct()
 
-        self.add(
-            operations.PostgresAddDefaultPartition(
-                model_name=model.__name__, name="default"
+        if partitioning_options["method"] != PostgresPartitioningMethod.HASH:
+            self.add(
+                operations.PostgresAddDefaultPartition(
+                    model_name=model.__name__, name="default"
+                )
             )
-        )
 
         self.add(
             operations.PostgresCreatePartitionedModel(

psqlextra/backend/migrations/state/__init__.py

@@ -1,5 +1,6 @@
 from .materialized_view import PostgresMaterializedViewModelState
 from .partitioning import (
+    PostgresHashPartitionState,
     PostgresListPartitionState,
     PostgresPartitionedModelState,
     PostgresPartitionState,
@@ -10,6 +11,7 @@
 __all__ = [
     "PostgresPartitionState",
     "PostgresRangePartitionState",
+    "PostgresHashPartitionState",
     "PostgresListPartitionState",
     "PostgresPartitionedModelState",
     "PostgresViewModelState",

psqlextra/backend/migrations/state/partitioning.py

@@ -38,6 +38,24 @@ def __init__(self, app_label: str, model_name: str, name: str, values):
         self.values = values
 
 
+class PostgresHashPartitionState(PostgresPartitionState):
+    """Represents the state of a hash partition for a
+    :see:PostgresPartitionedModel during a migration."""
+
+    def __init__(
+        self,
+        app_label: str,
+        model_name: str,
+        name: str,
+        modulus: int,
+        remainder: int,
+    ):
+        super().__init__(app_label, model_name, name)
+
+        self.modulus = modulus
+        self.remainder = remainder
+
+
 class PostgresPartitionedModelState(PostgresModelState):
     """Represents the state of a :see:PostgresPartitionedModel in the
     migrations."""

psqlextra/backend/schema.py

@@ -38,6 +38,7 @@ class PostgresSchemaEditor(base_impl.schema_editor()):
     )
     sql_partition_by = " PARTITION BY %s (%s)"
     sql_add_default_partition = "CREATE TABLE %s PARTITION OF %s DEFAULT"
+    sql_add_hash_partition = "CREATE TABLE %s PARTITION OF %s FOR VALUES WITH (MODULUS %s, REMAINDER %s)"
     sql_add_range_partition = (
         "CREATE TABLE %s PARTITION OF %s FOR VALUES FROM (%s) TO (%s)"
     )
@@ -283,6 +284,52 @@ def add_list_partition(
             if comment:
                 self.set_comment_on_table(table_name, comment)
 
+    def add_hash_partition(
+        self,
+        model: Model,
+        name: str,
+        modulus: int,
+        remainder: int,
+        comment: Optional[str] = None,
+    ) -> None:
+        """Creates a new hash partition for the specified partitioned model.
+
+        Arguments:
+            model:
+                Partitioned model to create a partition for.
+
+            name:
+                Name to give to the new partition.
+                Final name will be "{table_name}_{partition_name}"
+
+            modulus:
+                Integer value by which the key is divided.
+
+            remainder:
+                The remainder of the hash value when divided by modulus.
+
+            comment:
+                Optionally, a comment to add on this partition table.
+        """
+
+        # asserts the model is a model set up for partitioning
+        self._partitioning_properties_for_model(model)
+
+        table_name = self.create_partition_table_name(model, name)
+
+        sql = self.sql_add_hash_partition % (
+            self.quote_name(table_name),
+            self.quote_name(model._meta.db_table),
+            "%s",
+            "%s",
+        )
+
+        with transaction.atomic():
+            self.execute(sql, (modulus, remainder))
+
+            if comment:
+                self.set_comment_on_table(table_name, comment)
+
     def add_default_partition(
         self, model: Model, name: str, comment: Optional[str] = None
     ) -> None:

psqlextra/types.py

@@ -35,3 +35,4 @@ class PostgresPartitioningMethod(StrEnum):
 
     RANGE = "range"
     LIST = "list"
+    HASH = "hash"

setup.py

@@ -61,9 +61,7 @@ def run(self):
     python_requires=">=3.6",
     install_requires=[
         "Django>=2.0",
-        "enforce>=0.3.4,<=1.0.0",
         "python-dateutil>=2.8.0,<=3.0.0",
-        "structlog>=19,<=20.1.0",
         "ansimarkup>=1.4.0,<=2.0.0",
     ],
     extras_require={