Skip to content

Commit a5cde0a

Browse files
committed
Allow partition manager to plan for a subset of models
Fixes #252
1 parent 5dec4a2 commit a5cde0a

File tree

4 files changed

+69
-5
lines changed

4 files changed

+69
-5
lines changed

docs/source/table_partitioning.rst

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -175,7 +175,8 @@ Command-line options
175175
Long flag Short flag Default Description
176176
==================== ============= ================ ==================================================================================================== === === === === === ===
177177
``--yes`` ``-y`` ``False`` Specifies yes to all questions. You will NOT be asked for confirmation before partition deletion.
178-
``--using`` ``-u`` ``'default'`` Optional name of the database connection to use.
178+
``--using`` ``-u`` ``'default'`` Optionally, name of the database connection to use.
179+
``--model-names`` ``-m`` ``None`` Optionally, a list of model names to partition for.
179180
``--skip-create`` ``False`` Whether to skip creating partitions.
180181
``--skip-delete`` ``False`` Whether to skip deleting partitions.
181182

psqlextra/management/commands/pgpartition.py

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

3-
from typing import Optional
3+
from typing import List, Optional
44

55
from django.conf import settings
66
from django.core.management.base import BaseCommand
@@ -37,10 +37,18 @@ def add_arguments(self, parser):
3737
parser.add_argument(
3838
"--using",
3939
"-u",
40-
help="Optional name of the database connection to use.",
40+
help="Name of the database connection to use.",
4141
default="default",
4242
)
4343

44+
parser.add_argument(
45+
"--model-names",
46+
"-m",
47+
nargs="+",
48+
help="A list of model names for which to partition.",
49+
default=None,
50+
)
51+
4452
parser.add_argument(
4553
"--skip-create",
4654
action="store_true",
@@ -64,13 +72,17 @@ def handle( # type: ignore[override]
6472
using: Optional[str],
6573
skip_create: bool,
6674
skip_delete: bool,
75+
model_names: Optional[List[str]] = None,
6776
*args,
6877
**kwargs,
6978
):
7079
partitioning_manager = self._partitioning_manager()
7180

7281
plan = partitioning_manager.plan(
73-
skip_create=skip_create, skip_delete=skip_delete, using=using
82+
skip_create=skip_create,
83+
skip_delete=skip_delete,
84+
model_names=model_names,
85+
using=using,
7486
)
7587

7688
creations_count = len(plan.creations)

psqlextra/partitioning/manager.py

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ def plan(
2525
self,
2626
skip_create: bool = False,
2727
skip_delete: bool = False,
28+
model_names: Optional[List[str]] = None,
2829
using: Optional[str] = None,
2930
) -> PostgresPartitioningPlan:
3031
"""Plans which partitions should be deleted/created.
@@ -38,6 +39,10 @@ def plan(
3839
If set to True, no partitions will be marked
3940
for deletion, regardless of the configuration.
4041
42+
model_names:
43+
Optionally, only plan for the models with
44+
the specified name.
45+
4146
using:
4247
Optional name of the database connection to use.
4348
@@ -48,7 +53,19 @@ def plan(
4853

4954
model_plans = []
5055

56+
normalized_model_names = (
57+
[model_name.lower().strip() for model_name in model_names]
58+
if model_names
59+
else []
60+
)
61+
5162
for config in self.configs:
63+
if (
64+
model_names
65+
and config.model.__name__.lower() not in normalized_model_names
66+
):
67+
continue
68+
5269
model_plan = self._plan_for_config(
5370
config,
5471
skip_create=skip_create,

tests/test_partitioning_manager.py

Lines changed: 35 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,8 @@
11
import pytest
22

3-
from django.db import models
3+
from django.db import connection, models
44

5+
from psqlextra.backend.schema import PostgresSchemaEditor
56
from psqlextra.partitioning import (
67
PostgresPartitioningError,
78
PostgresPartitioningManager,
@@ -49,6 +50,39 @@ def test_partitioning_manager_find_config_for_model():
4950
assert manager.find_config_for_model(model2) == config2
5051

5152

53+
def test_partitioning_manager_plan_specific_model_names():
54+
"""Tests that only planning for specific models works as expected."""
55+
56+
model1 = define_fake_partitioned_model(
57+
{"timestamp": models.DateTimeField()}, {"key": ["timestamp"]}
58+
)
59+
60+
config1 = partition_by_current_time(model1, years=1, count=3)
61+
62+
model2 = define_fake_partitioned_model(
63+
{"timestamp": models.DateTimeField()}, {"key": ["timestamp"]}
64+
)
65+
66+
config2 = partition_by_current_time(model2, months=1, count=2)
67+
68+
schema_editor = PostgresSchemaEditor(connection)
69+
schema_editor.create_partitioned_model(model1)
70+
schema_editor.create_partitioned_model(model2)
71+
72+
manager = PostgresPartitioningManager([config1, config2])
73+
74+
plan = manager.plan()
75+
assert len(plan.model_plans) == 2
76+
77+
plan = manager.plan(model_names=[model2.__name__])
78+
assert len(plan.model_plans) == 1
79+
assert plan.model_plans[0].config.model == model2
80+
81+
# make sure casing is irrelevant
82+
plan = manager.plan(model_names=[model2.__name__.lower()])
83+
assert len(plan.model_plans) == 1
84+
85+
5286
def test_partitioning_manager_plan_not_partitioned_model():
5387
"""Tests that the auto partitioner does not try to auto partition for non-
5488
partitioned models/tables."""

0 commit comments

Comments
 (0)