truncate helper

psqlextra/backend/migrations/operations/add_list_sub_partition.py

@@ -0,0 +1,63 @@
+from psqlextra.backend.migrations.state import PostgresListPartitionState
+
+from .partition import PostgresPartitionOperation
+
+
+class PostgresAddListSubPartition(PostgresPartitionOperation):
+    """Adds a new list partition to a :see:PartitionedPostgresModel."""
+
+    def __init__(self, model_name, name, values, sub_key):
+        """Initializes new instance of :see:AddListPartition.
+
+        Arguments:
+            model_name:
+                The name of the :see:PartitionedPostgresModel.
+
+            name:
+                The name to give to the new partition table.
+
+            values:
+                Partition key values that should be
+                stored in this partition.
+        """
+
+        super().__init__(model_name, name)
+
+        self.values = values
+        self.sub_key = sub_key
+
+    def state_forwards(self, app_label, state):
+        model = state.models[(app_label, self.model_name_lower)]
+        model.add_partition(
+            PostgresListPartitionState(
+                app_label=app_label,
+                model_name=self.model_name,
+                name=self.name,
+                values=self.values,
+            )
+        )
+
+        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_list_sub_partition(
+                model, self.name, self.values, self.sub_key
+            )
+
+    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["values"] = self.values
+
+        return name, args, kwargs
+
+    def describe(self) -> str:
+        return "Creates list partition %s on %s" % (self.name, self.model_name)

psqlextra/backend/schema.py

@@ -78,7 +78,9 @@ class PostgresSchemaEditor(SchemaEditor):
     sql_add_list_partition = (
         "CREATE TABLE %s PARTITION OF %s FOR VALUES IN (%s)"
     )
+    sql_add_list_sub_partition = "CREATE TABLE %s PARTITION OF %s FOR VALUES IN (%s) PARTITION BY LIST (%s)"
     sql_delete_partition = "DROP TABLE %s"
+    sql_truncate_partition = "TRUNCATE TABLE %s"
     sql_table_comment = "COMMENT ON TABLE %s IS %s"
 
     side_effects: List[DatabaseSchemaEditor] = [
@@ -604,15 +606,22 @@ def create_partitioned_model(self, model: Type[Model]) -> None:
             self.quote_name(field_name) for field_name in meta.key
         )
 
+        if meta.sub_key:
+            primary_key_sql = ", ".join(
+                self.quote_name(field_name) for field_name in meta.sub_key
+            )
+        else:
+            primary_key_sql = partitioning_key_sql
+
         # create a composite key that includes the partitioning key
         sql = sql.replace(" PRIMARY KEY", "")
         if model._meta.pk and model._meta.pk.name not in meta.key:
             sql = sql[:-1] + ", PRIMARY KEY (%s, %s))" % (
                 self.quote_name(model._meta.pk.name),
-                partitioning_key_sql,
+                primary_key_sql,
             )
         else:
-            sql = sql[:-1] + ", PRIMARY KEY (%s))" % (partitioning_key_sql,)
+            sql = sql[:-1] + ", PRIMARY KEY (%s))" % (primary_key_sql,)
 
         # extend the standard CREATE TABLE statement with
         # 'PARTITION BY ...'
@@ -722,6 +731,103 @@ def add_list_partition(
             if comment:
                 self.set_comment_on_table(table_name, comment)
 
+    def add_list_sub_partition(
+        self,
+        model: Type[Model],
+        name: str,
+        values: List[Any],
+        sub_key: str,
+        comment: Optional[str] = None,
+    ) -> None:
+        """Creates a new list sub 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}"
+
+            values:
+                Partition key values that should be
+                stored in this partition.
+
+            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)
+
+        sql = self.sql_add_list_sub_partition % (
+            self.quote_name(name),
+            self.quote_name(model._meta.db_table),
+            ",".join(["%s" for _ in range(len(values))]),
+            self.quote_name(sub_key),
+        )
+
+        with transaction.atomic():
+            self.execute(sql, values)
+
+            if comment:
+                self.set_comment_on_table(name, comment)
+
+    def add_custom_list_partition(
+        self,
+        partition_table_name: str,
+        source_table: str,
+        values: List[Any],
+        comment: Optional[str] = None,
+    ) -> None:
+        """Creates a new list partition for the specified source_table.
+        Bypasses the need for a model instance, since sub partitioning doesn't
+        cleanly fit into the model-based partitioning system.
+
+        Arguments:
+            partition_table_name:
+                Name to give to the new partition table.
+
+            source_table:
+                The table to partition.
+
+            values:
+                Partition key values that should be
+                stored in this partition.
+
+            comment:
+                Optionally, a comment to add on this
+                partition table.
+        """
+
+        sql = self.sql_add_list_partition % (
+            self.quote_name(partition_table_name),
+            self.quote_name(source_table),
+            ",".join(["%s" for _ in range(len(values))]),
+        )
+
+        with transaction.atomic():
+            self.execute(sql, values)
+
+            if comment:
+                self.set_comment_on_table(partition_table_name, comment)
+
+    def delete_custom_partition(self, name: str) -> None:
+        """Deletes the partition with the specified name."""
+
+        sql = self.sql_delete_partition % self.quote_name(name)
+
+        self.execute(sql)
+
+    def truncate_custom_partition(self, name: str) -> None:
+        """Deletes the partition with the specified name."""
+
+        sql = self.sql_truncate_partition % self.quote_name(name)
+
+        self.execute(sql)
+
     def add_hash_partition(
         self,
         model: Type[Model],

psqlextra/models/options.py

@@ -10,12 +10,18 @@ class PostgresPartitionedModelOptions:
     are held.
     """
 
-    def __init__(self, method: PostgresPartitioningMethod, key: List[str]):
+    def __init__(
+        self,
+        method: PostgresPartitioningMethod,
+        key: List[str],
+        sub_key: List[str],
+    ):
         self.method = method
         self.key = key
+        self.sub_key = sub_key
         self.original_attrs: Dict[
             str, Union[PostgresPartitioningMethod, List[str]]
-        ] = dict(method=method, key=key)
+        ] = dict(method=method, key=key, sub_key=sub_key)
 
 
 class PostgresViewOptions:

psqlextra/models/partitioned.py

@@ -25,12 +25,15 @@ def __new__(cls, name, bases, attrs, **kwargs):
 
         method = getattr(meta_class, "method", None)
         key = getattr(meta_class, "key", None)
+        sub_key = getattr(meta_class, "sub_key", None)
 
-        patitioning_meta = PostgresPartitionedModelOptions(
-            method=method or cls.default_method, key=key or cls.default_key
+        partitioning_meta = PostgresPartitionedModelOptions(
+            method=method or cls.default_method,
+            key=key or cls.default_key,
+            sub_key=sub_key or cls.default_key,
         )
 
-        new_class.add_to_class("_partitioning_meta", patitioning_meta)
+        new_class.add_to_class("_partitioning_meta", partitioning_meta)
         return new_class
 
 

psqlextra/partitioning/plan.py

@@ -54,7 +54,7 @@ def apply(self, using: Optional[str]) -> None:
     def print(self) -> None:
         """Prints this model plan to the terminal in a readable format."""
 
-        print(f"{self.config.model.__name__}:")
+        print(f"{self.config.model.__name__}: ")
 
         for partition in self.deletions:
             print("  - %s" % partition.name())

tests/test_make_migrations.py

@@ -30,19 +30,25 @@
         dict(
             fields={"category": models.TextField()},
             partitioning_options=dict(
-                method=PostgresPartitioningMethod.LIST, key="category"
+                method=PostgresPartitioningMethod.LIST,
+                key="category",
+                sub_key=[],
             ),
         ),
         dict(
             fields={"timestamp": models.DateTimeField()},
             partitioning_options=dict(
-                method=PostgresPartitioningMethod.RANGE, key="timestamp"
+                method=PostgresPartitioningMethod.RANGE,
+                key="timestamp",
+                sub_key=[],
             ),
         ),
         dict(
             fields={"artist_id": models.IntegerField()},
             partitioning_options=dict(
-                method=PostgresPartitioningMethod.HASH, key="artist_id"
+                method=PostgresPartitioningMethod.HASH,
+                key="artist_id",
+                sub_key=[],
             ),
         ),
     ],

tests/test_management_command_partition.py

@@ -101,6 +101,7 @@ def test_management_command_partition_dry_run(
     create/delete partitions."""
 
     config = fake_partitioning_manager.find_config_for_model(fake_model)
+
     snapshot.assert_match(run(args))
 
     config.strategy.createable_partition.create.assert_not_called()