Allow partition manager to plan for a subset of models

Fixes #252

Author: Photonios

docs/source/table_partitioning.rst

@@ -175,7 +175,8 @@ Command-line options
   Long flag            Short flag    Default          Description
  ==================== ============= ================ ==================================================================================================== === === === === === ===
   ``--yes``            ``-y``        ``False``        Specifies yes to all questions. You will NOT be asked for confirmation before partition deletion.
-  ``--using``          ``-u``        ``'default'``    Optional name of the database connection to use.
+  ``--using``          ``-u``        ``'default'``    Optionally, name of the database connection to use.
+  ``--model-names``    ``-m``        ``None``         Optionally, a list of model names to partition for.
   ``--skip-create``                  ``False``        Whether to skip creating partitions.
   ``--skip-delete``                  ``False``        Whether to skip deleting partitions.
 

psqlextra/management/commands/pgpartition.py

@@ -1,6 +1,6 @@
 import sys
 
-from typing import Optional
+from typing import List, Optional
 
 from django.conf import settings
 from django.core.management.base import BaseCommand
@@ -37,10 +37,18 @@ def add_arguments(self, parser):
         parser.add_argument(
             "--using",
             "-u",
-            help="Optional name of the database connection to use.",
+            help="Name of the database connection to use.",
             default="default",
         )
 
+        parser.add_argument(
+            "--model-names",
+            "-m",
+            nargs="+",
+            help="A list of model names for which to partition.",
+            default=None,
+        )
+
         parser.add_argument(
             "--skip-create",
             action="store_true",
@@ -64,13 +72,17 @@ def handle(  # type: ignore[override]
         using: Optional[str],
         skip_create: bool,
         skip_delete: bool,
+        model_names: Optional[List[str]] = None,
         *args,
         **kwargs,
     ):
         partitioning_manager = self._partitioning_manager()
 
         plan = partitioning_manager.plan(
-            skip_create=skip_create, skip_delete=skip_delete, using=using
+            skip_create=skip_create,
+            skip_delete=skip_delete,
+            model_names=model_names,
+            using=using,
         )
 
         creations_count = len(plan.creations)

psqlextra/partitioning/manager.py

@@ -25,6 +25,7 @@ def plan(
         self,
         skip_create: bool = False,
         skip_delete: bool = False,
+        model_names: Optional[List[str]] = None,
         using: Optional[str] = None,
     ) -> PostgresPartitioningPlan:
         """Plans which partitions should be deleted/created.
@@ -38,6 +39,10 @@ def plan(
                 If set to True, no partitions will be marked
                 for deletion, regardless of the configuration.
 
+            model_names:
+                Optionally, only plan for the models with
+                the specified name.
+
             using:
                 Optional name of the database connection to use.
 
@@ -48,7 +53,19 @@ def plan(
 
         model_plans = []
 
+        normalized_model_names = (
+            [model_name.lower().strip() for model_name in model_names]
+            if model_names
+            else []
+        )
+
         for config in self.configs:
+            if (
+                model_names
+                and config.model.__name__.lower() not in normalized_model_names
+            ):
+                continue
+
             model_plan = self._plan_for_config(
                 config,
                 skip_create=skip_create,

tests/test_partitioning_manager.py

@@ -1,7 +1,8 @@
 import pytest
 
-from django.db import models
+from django.db import connection, models
 
+from psqlextra.backend.schema import PostgresSchemaEditor
 from psqlextra.partitioning import (
     PostgresPartitioningError,
     PostgresPartitioningManager,
@@ -49,6 +50,39 @@ def test_partitioning_manager_find_config_for_model():
     assert manager.find_config_for_model(model2) == config2
 
 
+def test_partitioning_manager_plan_specific_model_names():
+    """Tests that only planning for specific models works as expected."""
+
+    model1 = define_fake_partitioned_model(
+        {"timestamp": models.DateTimeField()}, {"key": ["timestamp"]}
+    )
+
+    config1 = partition_by_current_time(model1, years=1, count=3)
+
+    model2 = define_fake_partitioned_model(
+        {"timestamp": models.DateTimeField()}, {"key": ["timestamp"]}
+    )
+
+    config2 = partition_by_current_time(model2, months=1, count=2)
+
+    schema_editor = PostgresSchemaEditor(connection)
+    schema_editor.create_partitioned_model(model1)
+    schema_editor.create_partitioned_model(model2)
+
+    manager = PostgresPartitioningManager([config1, config2])
+
+    plan = manager.plan()
+    assert len(plan.model_plans) == 2
+
+    plan = manager.plan(model_names=[model2.__name__])
+    assert len(plan.model_plans) == 1
+    assert plan.model_plans[0].config.model == model2
+
+    # make sure casing is irrelevant
+    plan = manager.plan(model_names=[model2.__name__.lower()])
+    assert len(plan.model_plans) == 1
+
+
 def test_partitioning_manager_plan_not_partitioned_model():
     """Tests that the auto partitioner does not try to auto partition for non-
     partitioned models/tables."""