Make time partition suffix customizable

Author: ckrybus

psqlextra/partitioning/current_time_strategy.py

@@ -24,6 +24,7 @@ def __init__(
         size: PostgresTimePartitionSize,
         count: int,
         max_age: Optional[relativedelta] = None,
+        name_format: Optional[str] = None,
     ) -> None:
         """Initializes a new instance of :see:PostgresTimePartitioningStrategy.
 
@@ -44,13 +45,16 @@ def __init__(
         self.size = size
         self.count = count
         self.max_age = max_age
+        self.name_format = name_format
 
     def to_create(self) -> Generator[PostgresTimePartition, None, None]:
         current_datetime = self.size.start(self.get_start_datetime())
 
         for _ in range(self.count):
             yield PostgresTimePartition(
-                start_datetime=current_datetime, size=self.size
+                start_datetime=current_datetime,
+                size=self.size,
+                name_format=self.name_format,
             )
 
             current_datetime += self.size.as_delta()
@@ -65,7 +69,9 @@ def to_delete(self) -> Generator[PostgresTimePartition, None, None]:
 
         while True:
             yield PostgresTimePartition(
-                start_datetime=current_datetime, size=self.size
+                start_datetime=current_datetime,
+                size=self.size,
+                name_format=self.name_format,
             )
 
             current_datetime -= self.size.as_delta()

psqlextra/partitioning/shorthands.py

@@ -17,6 +17,7 @@ def partition_by_current_time(
     weeks: Optional[int] = None,
     days: Optional[int] = None,
     max_age: Optional[relativedelta] = None,
+    name_format: Optional[str] = None,
 ) -> PostgresPartitioningConfig:
     """Short-hand for generating a partitioning config that partitions the
     specified model by time.
@@ -48,6 +49,10 @@ def partition_by_current_time(
 
             Partitions older than this are deleted when running
             a delete/cleanup run.
+
+        name_format:
+            The datetime format which is being passed to datetime.strftime
+            to generate the partition name.
     """
 
     size = PostgresTimePartitionSize(
@@ -57,7 +62,10 @@ def partition_by_current_time(
     return PostgresPartitioningConfig(
         model=model,
         strategy=PostgresCurrentTimePartitioningStrategy(
-            size=size, count=count, max_age=max_age
+            size=size,
+            count=count,
+            max_age=max_age,
+            name_format=name_format,
         ),
     )
 

psqlextra/partitioning/time_partition.py

@@ -1,4 +1,5 @@
 from datetime import datetime
+from typing import Optional
 
 from .error import PostgresPartitioningError
 from .range_partition import PostgresRangePartition
@@ -22,7 +23,10 @@ class PostgresTimePartition(PostgresRangePartition):
     }
 
     def __init__(
-        self, size: PostgresTimePartitionSize, start_datetime: datetime
+        self,
+        size: PostgresTimePartitionSize,
+        start_datetime: datetime,
+        name_format: Optional[str] = None,
     ) -> None:
         end_datetime = start_datetime + size.as_delta()
 
@@ -34,9 +38,12 @@ def __init__(
         self.size = size
         self.start_datetime = start_datetime
         self.end_datetime = end_datetime
+        self.name_format = name_format
 
     def name(self) -> str:
-        name_format = self._unit_name_format.get(self.size.unit)
+        name_format = self.name_format or self._unit_name_format.get(
+            self.size.unit
+        )
         if not name_format:
             raise PostgresPartitioningError("Unknown size/unit")
 

tests/test_partitioning_time.py

@@ -115,6 +115,45 @@ def test_partitioning_time_monthly_apply():
     assert table.partitions[13].name == "2020_feb"
 
 
+@pytest.mark.postgres_version(lt=110000)
+def test_partitioning_time_monthly_with_custom_naming_apply():
+    """Tests whether automatically created new partitions are named according
+    to the specified name_format."""
+
+    model = define_fake_partitioned_model(
+        {"timestamp": models.DateTimeField()}, {"key": ["timestamp"]}
+    )
+
+    schema_editor = connection.schema_editor()
+    schema_editor.create_partitioned_model(model)
+
+    # create partitions for the next 12 months (including the current)
+    with freezegun.freeze_time("2019-1-30"):
+        manager = PostgresPartitioningManager(
+            [
+                partition_by_current_time(
+                    model, months=1, count=12, name_format="%Y_%m"
+                )
+            ]
+        )
+        manager.plan().apply()
+
+    table = _get_partitioned_table(model)
+    assert len(table.partitions) == 12
+    assert table.partitions[0].name == "2019_01"
+    assert table.partitions[1].name == "2019_02"
+    assert table.partitions[2].name == "2019_03"
+    assert table.partitions[3].name == "2019_04"
+    assert table.partitions[4].name == "2019_05"
+    assert table.partitions[5].name == "2019_06"
+    assert table.partitions[6].name == "2019_07"
+    assert table.partitions[7].name == "2019_08"
+    assert table.partitions[8].name == "2019_09"
+    assert table.partitions[9].name == "2019_10"
+    assert table.partitions[10].name == "2019_11"
+    assert table.partitions[11].name == "2019_12"
+
+
 @pytest.mark.postgres_version(lt=110000)
 def test_partitioning_time_weekly_apply():
     """Tests whether automatically creating new partitions ahead weekly works