Skip to content

Commit efe0cf4

Browse files
committed
Make multi-day partitions deterministic and aligned
1 parent ffdc9fb commit efe0cf4

File tree

2 files changed

+57
-5
lines changed

2 files changed

+57
-5
lines changed

psqlextra/partitioning/time_partition_size.py

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import enum
22

3-
from datetime import date, datetime
3+
from datetime import date, datetime, timedelta, timezone
44
from typing import Optional, Union
55

66
from dateutil.relativedelta import relativedelta
@@ -15,18 +15,23 @@ class PostgresTimePartitionUnit(enum.Enum):
1515
DAYS = "days"
1616

1717

18+
UNIX_EPOCH = datetime(1970, 1, 1, tzinfo=timezone.utc)
19+
20+
1821
class PostgresTimePartitionSize:
1922
"""Size of a time-based range partition table."""
2023

2124
unit: PostgresTimePartitionUnit
2225
value: int
26+
anchor: datetime
2327

2428
def __init__(
2529
self,
2630
years: Optional[int] = None,
2731
months: Optional[int] = None,
2832
weeks: Optional[int] = None,
2933
days: Optional[int] = None,
34+
anchor: datetime = UNIX_EPOCH
3035
) -> None:
3136
sizes = [years, months, weeks, days]
3237

@@ -38,6 +43,7 @@ def __init__(
3843
"Partition can only have on size unit."
3944
)
4045

46+
self.anchor = anchor
4147
if years:
4248
self.unit = PostgresTimePartitionUnit.YEARS
4349
self.value = years
@@ -82,7 +88,10 @@ def start(self, dt: datetime) -> datetime:
8288
if self.unit == PostgresTimePartitionUnit.WEEKS:
8389
return self._ensure_datetime(dt - relativedelta(days=dt.weekday()))
8490

85-
return self._ensure_datetime(dt)
91+
diff_days = (dt - self.anchor).days
92+
partition_index = diff_days // self.value
93+
start = self.anchor + timedelta(days=partition_index * self.value)
94+
return self._ensure_datetime(start)
8695

8796
@staticmethod
8897
def _ensure_datetime(dt: Union[date, datetime]) -> datetime:

tests/test_partitioning_time.py

Lines changed: 46 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -254,6 +254,49 @@ def test_partitioning_time_daily_apply():
254254
assert table.partitions[6].name == "2019_jun_04"
255255

256256

257+
@pytest.mark.postgres_version(lt=110000)
258+
def test_partitioning_time_consistent_daily_apply():
259+
"""Ensures that automatic daily partition creation is consistent and aligned
260+
when the partition size spans multiple days (e.g., days > 1)"""
261+
262+
model = define_fake_partitioned_model(
263+
{"timestamp": models.DateTimeField()}, {"key": ["timestamp"]}
264+
)
265+
266+
schema_editor = connection.schema_editor()
267+
schema_editor.create_partitioned_model(model)
268+
269+
with freezegun.freeze_time("2025-06-20"):
270+
manager = PostgresPartitioningManager(
271+
[partition_by_current_time(model, days=5, count=3)]
272+
)
273+
manager.plan().apply()
274+
275+
table = _get_partitioned_table(model)
276+
assert len(table.partitions) == 3
277+
278+
# Partitions are aligned based on the fixed anchor (Unix Epoch by default).
279+
# 2025-06-20 falls within the partition starting at 2025-06-16,
280+
# since it's the most recent multiple of 5 days since 1970-01-01.
281+
assert table.partitions[0].name == "2025_jun_16"
282+
assert table.partitions[1].name == "2025_jun_21"
283+
assert table.partitions[2].name == "2025_jun_26"
284+
285+
# re-running it another day only creates the next one needed.
286+
with freezegun.freeze_time("2025-06-22"):
287+
manager = PostgresPartitioningManager(
288+
[partition_by_current_time(model, days=5, count=3)]
289+
)
290+
manager.plan().apply()
291+
292+
table = _get_partitioned_table(model)
293+
assert len(table.partitions) == 4
294+
assert table.partitions[0].name == "2025_jun_16"
295+
assert table.partitions[1].name == "2025_jun_21"
296+
assert table.partitions[2].name == "2025_jun_26"
297+
assert table.partitions[3].name == "2025_jul_01"
298+
299+
257300
@pytest.mark.postgres_version(lt=110000)
258301
def test_partitioning_time_monthly_apply_insert():
259302
"""Tests whether automatically created monthly partitions line up
@@ -376,7 +419,7 @@ def test_partitioning_time_daily_apply_insert():
376419
@pytest.mark.parametrize(
377420
"kwargs,partition_names",
378421
[
379-
(dict(days=2), ["2019_jan_01", "2019_jan_03"]),
422+
(dict(days=2), ["2018_dec_31", "2019_jan_02"]),
380423
(dict(weeks=2), ["2018_week_53", "2019_week_02"]),
381424
(dict(months=2), ["2019_jan", "2019_mar"]),
382425
(dict(years=2), ["2019", "2021"]),
@@ -422,7 +465,7 @@ def test_partitioning_time_multiple(kwargs, partition_names):
422465
dict(days=7, max_age=relativedelta(weeks=1)),
423466
[
424467
("2019-1-1", 6),
425-
("2019-1-4", 6),
468+
("2019-1-4", 5),
426469
("2019-1-8", 5),
427470
("2019-1-15", 4),
428471
("2019-1-16", 4),
@@ -450,7 +493,7 @@ def test_partitioning_time_delete(kwargs, timepoints):
450493
with freezegun.freeze_time(timepoints[0][0]):
451494
manager.plan().apply()
452495

453-
for index, (dt, partition_count) in enumerate(timepoints):
496+
for (dt, partition_count) in timepoints:
454497
with freezegun.freeze_time(dt):
455498
manager.plan(skip_create=True).apply()
456499

0 commit comments

Comments
 (0)