Support for hash partioning

Author: Anexen

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::
 

psqlextra/backend/introspection.py

@@ -5,6 +5,11 @@
 
 from . import base_impl
 
+PARTITIONING_STRATEGY_TO_METHOD = {
+    "r": PostgresPartitioningMethod.RANGE,
+    "l": PostgresPartitioningMethod.LIST,
+    "h": PostgresPartitioningMethod.HASH,
+}
 
 @dataclass
 class PostgresIntrospectedPartitionTable:
@@ -64,9 +69,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]),
             )

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,71 @@
+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/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/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"

tests/test_make_migrations.py

@@ -37,6 +37,12 @@
                 method=PostgresPartitioningMethod.RANGE, key="timestamp"
             ),
         ),
+        dict(
+            fields={"artist_id": models.IntegerField()},
+            partitioning_options=dict(
+                method=PostgresPartitioningMethod.HASH, key="artist_id"
+            )
+        )
     ],
 )
 @postgres_patched_migrations()