diff --git a/psqlextra/backend/schema.py b/psqlextra/backend/schema.py index b8dd0a7..8d1e77c 100644 --- a/psqlextra/backend/schema.py +++ b/psqlextra/backend/schema.py @@ -752,7 +752,7 @@ def add_list_partition( """ # asserts the model is a model set up for partitioning - self._partitioning_properties_for_model(model) + meta = self._partitioning_properties_for_model(model) table_name = self.create_partition_table_name(model, name) @@ -762,6 +762,22 @@ def add_list_partition( ",".join(["%s" for _ in range(len(values))]), ) + if getattr(meta, "sub_key", None) and len(meta.sub_key) > 0: + sub_partitioning_key_sql = ", ".join( + self.quote_name(field_name) for field_name in meta.sub_key + ) + + last_brace_idx = sql.rfind(")") + 1 + sql = ( + sql[:last_brace_idx] + + self.sql_partition_by + % ( + meta.sub_method.upper(), + sub_partitioning_key_sql, + ) + + sql[last_brace_idx:] + ) + with transaction.atomic(using=self.connection.alias): self.execute(sql, values) @@ -1100,9 +1116,17 @@ def _partitioning_properties_for_model(model: Type[Model]): ) % model.__name__ ) - + if meta.sub_method and len(meta.sub_key)==0: + raise ImproperlyConfigured( + ( + "Model '%s' is not properly configured to be partitioned." + " 'sub_method' is specified '%s', but no 'sub_key' is defined." + ) + % (model.__name__, meta.sub_method) + ) + try: - for field_name in meta.key: + for field_name in meta.key + meta.sub_key: model._meta.get_field(field_name) except FieldDoesNotExist: raise ImproperlyConfigured( diff --git a/psqlextra/contrib/__init__.py b/psqlextra/contrib/__init__.py index 97794eb..854d8c3 100644 --- a/psqlextra/contrib/__init__.py +++ b/psqlextra/contrib/__init__.py @@ -1,6 +1,7 @@ from .model_data_migrator import PostgresModelDataMigrator from .static_row import StaticRowQueryCompiler, StaticRowQuerySet from .transaction import no_transaction +from .category_current_time import partition_by_category_and_current_time __all__ = [ "PostgresModelDataMigrator", @@ -8,4 +9,5 @@ "StaticRowQueryCompiler", "StaticRowQuerySet", "no_transaction", + "partition_by_category_and_current_time", ] diff --git a/psqlextra/contrib/category_current_time/__init__.py b/psqlextra/contrib/category_current_time/__init__.py new file mode 100644 index 0000000..bb7a95a --- /dev/null +++ b/psqlextra/contrib/category_current_time/__init__.py @@ -0,0 +1,3 @@ +from .shorthands import partition_by_category_and_current_time + +__all__ = ["partition_by_category_and_current_time"] \ No newline at end of file diff --git a/psqlextra/contrib/category_current_time/category_current_time_strategy.py b/psqlextra/contrib/category_current_time/category_current_time_strategy.py new file mode 100644 index 0000000..cd2cd52 --- /dev/null +++ b/psqlextra/contrib/category_current_time/category_current_time_strategy.py @@ -0,0 +1,100 @@ +from datetime import datetime, timezone +from typing import Any, Generator, Optional +from dateutil.relativedelta import relativedelta + +from psqlextra.partitioning.strategy import PostgresPartitioningStrategy +from psqlextra.partitioning.time_partition_size import PostgresTimePartitionSize + +from .partition import PostgresListPartition, PostgresTimeSubPartition + + +class PostgresCategoryCurrentTimePartitioningStrategy(PostgresPartitioningStrategy): + """Implments a category and time based partitioning strategy where + on the first level a partition is created for each category and on the + second level each partition contains values for a specific time period. + + All buckets will be equal in size and start at the start of the + unit. With monthly partitioning, partitions start on the 1st and + with weekly partitioning, partitions start on monday, with hourly + partitioning, partitions start at 00:00. + """ + + def __init__( + self, + categories: list[Any], + size: PostgresTimePartitionSize, + count: int, + max_age: Optional[relativedelta] = None, + name_format: Optional[tuple[str, str]] = None, + ) -> None: + """Initializes a new instance of :see:PostgresTimePartitioningStrategy. + + Arguments: + categories: + The list of categories to create first-level partitions for. + + size: + The size of each partition. + + count: + The amount of partitions to create ahead + from the current date/time. + + max_age: + Maximum age of a partition. Partitions + older than this are deleted during + auto cleanup. + + name_format: + Optional name format for the partitions supplied as a tuple. The first value is used + for the category partitions, the second for the time partitions. + """ + + self.size = size + self.count = count + self.max_age = max_age + self.name_format = name_format or (None, None) + self.categories = categories + + def to_create(self) -> Generator[PostgresTimeSubPartition, None, None]: + for category in self.categories: + lp = PostgresListPartition( + values=[category], name_format=self.name_format[0] + ) + yield lp + + current_datetime = self.size.start(self.get_start_datetime()) + + for _ in range(self.count): + yield PostgresTimeSubPartition( + parent_partition=lp, + start_datetime=current_datetime, + size=self.size, + name_format=self.name_format[1], + ) + + current_datetime += self.size.as_delta() + + def to_delete(self) -> Generator[PostgresTimeSubPartition, None, None]: + if not self.max_age: + return + + current_datetime = self.size.start(self.get_start_datetime() - self.max_age) + + while True: + for category in self.categories: + lp = PostgresListPartition( + values=[category], name_format=self.name_format[0] + ) + + yield PostgresTimeSubPartition( + parent_partition=lp, + start_datetime=current_datetime, + size=self.size, + name_format=self.name_format[1], + ) + + current_datetime -= self.size.as_delta() + + def get_start_datetime(self) -> datetime: + return datetime.now(timezone.utc) diff --git a/psqlextra/contrib/category_current_time/partition.py b/psqlextra/contrib/category_current_time/partition.py new file mode 100644 index 0000000..a4a189f --- /dev/null +++ b/psqlextra/contrib/category_current_time/partition.py @@ -0,0 +1,120 @@ +from datetime import datetime +from typing import Any, Optional, Type + +from contextlib import contextmanager + +from psqlextra.backend.schema import PostgresSchemaEditor +from psqlextra.models import PostgresPartitionedModel + +from psqlextra.partitioning.time_partition import ( + PostgresTimePartition, + PostgresTimePartitionSize, +) + +from psqlextra.partitioning import PostgresPartition +from psqlextra.partitioning.error import PostgresPartitioningError + +@contextmanager +def patch_model(model: Type[PostgresPartitionedModel], parent_partition_name: str): + """Context manager that ensures the model's table is created and + deleted properly. + """ + original_table = model._meta.db_table + try: + model._meta.db_table = original_table + "_" + parent_partition_name + yield model + finally: + model._meta.db_table = original_table + +class PostgresTimeSubPartition(PostgresTimePartition): + """Base class for a PostgreSQL table sub-partition in a range partitioned + table.""" + + def __init__( + self, + parent_partition: PostgresPartition, + size: PostgresTimePartitionSize, + start_datetime: datetime, + name_format: Optional[str] = None, + ) -> None: + super().__init__( + size=size, + start_datetime=start_datetime, + name_format=name_format, + ) + self.parent_partition = parent_partition + + def name(self) -> str: + name_format = self.name_format or self._unit_name_format.get(self.size.unit) + if not name_format: + raise PostgresPartitioningError("Unknown size/unit") + + return ( + self.parent_partition.name() + + "_" + + self.start_datetime.strftime(name_format).lower() + ) + + def create( + self, + model: Type[PostgresPartitionedModel], + schema_editor: PostgresSchemaEditor, + comment: Optional[str] = None, + ) -> None: + with patch_model(model, self.parent_partition.name()) as managed_model: + super().create( + model=managed_model, + schema_editor=schema_editor, + comment=comment, + ) + + def delete( + self, + model: Type[PostgresPartitionedModel], + schema_editor: PostgresSchemaEditor, + ) -> None: + with patch_model(model, self.parent_partition.name()) as managed_model: + super().delete( + model=managed_model, + schema_editor=schema_editor, + ) + +class PostgresListPartition(PostgresPartition): + """Base class for a PostgreSQL table partition in a list partitioned + table.""" + + def __init__(self, values: list[Any], name_format: Optional[str] = None) -> None: + self.values = values + self.name_format = name_format or "%s" + + def name(self) -> str: + return self.name_format % "_".join(str(v).lower() for v in self.values) + + def deconstruct(self) -> dict: + return { + **super().deconstruct(), + "values": self.values, + } + + def create( + self, + model: Type[PostgresPartitionedModel], + schema_editor: PostgresSchemaEditor, + comment: Optional[str] = None, + ) -> None: + schema_editor.add_list_partition( + model=model, + name=self.name(), + values=self.values, + comment=comment, + ) + + def delete( + self, + model: Type[PostgresPartitionedModel], + schema_editor: PostgresSchemaEditor, + ) -> None: + schema_editor.delete_partition(model, self.name()) + + +__all__ = ["PostgresListPartition"] diff --git a/psqlextra/contrib/category_current_time/shorthands.py b/psqlextra/contrib/category_current_time/shorthands.py new file mode 100644 index 0000000..9ab1199 --- /dev/null +++ b/psqlextra/contrib/category_current_time/shorthands.py @@ -0,0 +1,87 @@ +from typing import Optional, Type + +from dateutil.relativedelta import relativedelta + +from psqlextra.models import PostgresPartitionedModel + +from psqlextra.partitioning.config import PostgresPartitioningConfig +from psqlextra.partitioning import PostgresTimePartitionSize + +from .category_current_time_strategy import ( + PostgresCategoryCurrentTimePartitioningStrategy, +) + + +def partition_by_category_and_current_time( + model: Type[PostgresPartitionedModel], + count: int, + categories: list[int] = None, + years: Optional[int] = None, + months: Optional[int] = None, + weeks: Optional[int] = None, + days: Optional[int] = None, + hours: 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. + + One specifies one of the `years`, `months`, `weeks` + or `days` parameter to indicate the size of each + partition. These parameters cannot be combined. + + Arguments: + count: + The amount of partitions to create ahead of + the current date/time. + + categories: + The list of categories to create first-level partitions for. + + years: + The amount of years each partition should contain. + + months: + The amount of months each partition should contain. + + weeks: + The amount of weeks each partition should contain. + + days: + The amount of days each partition should contain. + + hours: + The amount of hours each partition should contain. + + max_age: + The maximum age of a partition (calculated from the + start of the partition). + + Partitions older than this are deleted when running + a delete/cleanup run. + + name_format: + The datetime format supplied as a tuple. + The first value creates the first level partition names + and the second value is passed to datetime.strftime to generate the + second level partition name. + """ + + size = PostgresTimePartitionSize( + years=years, months=months, weeks=weeks, days=days, hours=hours + ) + + return PostgresPartitioningConfig( + model=model, + strategy=PostgresCategoryCurrentTimePartitioningStrategy( + categories=categories or [], + size=size, + count=count, + max_age=max_age, + name_format=name_format, + ), + ) + + +__all_ = ["partition_by_category_and_current_time"] diff --git a/psqlextra/models/options.py b/psqlextra/models/options.py index 79e8286..4f6fdf1 100644 --- a/psqlextra/models/options.py +++ b/psqlextra/models/options.py @@ -10,12 +10,32 @@ class PostgresPartitionedModelOptions: are held. """ - def __init__(self, method: PostgresPartitioningMethod, key: List[str]): + def __init__( + self, + method: PostgresPartitioningMethod, + key: List[str], + sub_key: Optional[List[str]] = None, + sub_method: Optional[PostgresPartitioningMethod] = None, + ): self.method = method self.key = key + self.sub_method = sub_method if sub_method else None + self.sub_key = sub_key if sub_key else [] self.original_attrs: Dict[ - str, Union[PostgresPartitioningMethod, List[str]] - ] = dict(method=method, key=key) + str, + Union[ + PostgresPartitioningMethod, + List[str], + PostgresPartitioningMethod, + List[str], + None, + ], + ] = dict( + method=method, + key=key, + sub_method=sub_method, + sub_key=sub_key, + ) class PostgresViewOptions: @@ -28,6 +48,4 @@ class PostgresViewOptions: def __init__(self, query: Optional[SQLWithParams]): self.query = query - self.original_attrs: Dict[str, Optional[SQLWithParams]] = dict( - query=self.query - ) + self.original_attrs: Dict[str, Optional[SQLWithParams]] = dict(query=self.query) diff --git a/psqlextra/models/partitioned.py b/psqlextra/models/partitioned.py index 3a20677..c76e296 100644 --- a/psqlextra/models/partitioned.py +++ b/psqlextra/models/partitioned.py @@ -29,16 +29,21 @@ def __new__(cls, name, bases, attrs, **kwargs): partitioning_method = getattr(partitioning_meta_class, "method", None) partitioning_key = getattr(partitioning_meta_class, "key", None) + partitioning_sub_method = getattr(partitioning_meta_class, "sub_method", None) + partitioning_sub_key = getattr(partitioning_meta_class, "sub_key", None) + if django.VERSION >= (5, 2): for base in bases: cls._delete_auto_created_fields(base) - cls._create_primary_key(attrs, partitioning_key) + cls._create_primary_key(attrs, partitioning_key, partitioning_sub_key) patitioning_meta = PostgresPartitionedModelOptions( method=partitioning_method or cls.default_method, key=partitioning_key or cls.default_key, + sub_method=partitioning_sub_method, + sub_key=partitioning_sub_key, ) new_class = super().__new__(cls, name, bases, attrs, **kwargs) @@ -47,7 +52,7 @@ def __new__(cls, name, bases, attrs, **kwargs): @classmethod def _create_primary_key( - cls, attrs, partitioning_key: Optional[List[str]] + cls, attrs, partitioning_key: Optional[List[str]], partitioning_sub_key: Optional[List[str]] ) -> None: from django.db.models.fields.composite import CompositePrimaryKey @@ -75,7 +80,13 @@ def _create_primary_key( else list(filter(None, [partitioning_key])) ) - unique_pk_fields = set(pk_fields + (partitioning_keys or [])) + partitioning_sub_key = ( + partitioning_sub_key + if isinstance(partitioning_sub_key, list) + else list(filter(None, [partitioning_sub_key])) + ) + + unique_pk_fields = set(pk_fields + (partitioning_keys or []) + (partitioning_sub_key or [])) if len(unique_pk_fields) <= 1: if "id" in attrs: attrs["id"].primary_key = True diff --git a/psqlextra/partitioning/manager.py b/psqlextra/partitioning/manager.py index 01bac3b..cac7fe0 100644 --- a/psqlextra/partitioning/manager.py +++ b/psqlextra/partitioning/manager.py @@ -84,9 +84,7 @@ def find_config_for_model( ) -> Optional[PostgresPartitioningConfig]: """Finds the partitioning config for the specified model.""" - return next( - (config for config in self.configs if config.model == model), None - ) + return next((config for config in self.configs if config.model == model), None) def _plan_for_config( self, @@ -98,21 +96,19 @@ def _plan_for_config( """Creates a partitioning plan for one partitioning config.""" connection = connections[using or "default"] - table = self._get_partitioned_table(connection, config.model) - model_plan = PostgresModelPartitioningPlan(config) if not skip_create: for partition in config.strategy.to_create(): - if table.partition_by_name(name=partition.name()): + if self._get_partition_from_table(connection, config.model, partition): continue model_plan.creations.append(partition) if not skip_delete: for partition in config.strategy.to_delete(): - introspected_partition = table.partition_by_name( - name=partition.name() + introspected_partition = self._get_partition_from_table( + connection, config.model, partition ) if not introspected_partition: break @@ -128,22 +124,48 @@ def _plan_for_config( return model_plan @staticmethod - def _get_partitioned_table( - connection, model: Type[PostgresPartitionedModel] - ): + def _get_partition_from_table( + connection, + model: Type[PostgresPartitionedModel], + search_partition: PostgresPartition, + ) -> bool: + """Returns a partition from the table by name. + Traverses partitions if the model is sub-partitioned""" + with connection.cursor() as cursor: table = connection.introspection.get_partitioned_table( cursor, model._meta.db_table ) - if not table: - raise PostgresPartitioningError( - f"Model {model.__name__}, with table " - f"{model._meta.db_table} does not exists in the " - "database. Did you run `python manage.py migrate`?" - ) + if not table: + raise PostgresPartitioningError( + f"Model {model.__name__}, with table " + f"{model._meta.db_table} does not exists in the " + "database. Did you run `python manage.py migrate`?" + ) - return table + if len(getattr(model._partitioning_meta, "sub_key", [])) > 0: + partition = table.partition_by_name(name=search_partition.name()) + if partition: + return partition + else: + return next( + ( + partition + for partition in connection.introspection.get_partitions( + cursor, model._meta.db_table + ) + if connection.introspection.get_partitioned_table( + cursor, partition.full_name + ) + and connection.introspection.get_partitioned_table( + cursor, partition.full_name + ).partition_by_name(name=search_partition.name()) + ), + None, + ) + else: + return table.partition_by_name(name=search_partition.name()) @staticmethod def _validate_configs(configs: List[PostgresPartitioningConfig]): diff --git a/tests/test_make_migrations.py b/tests/test_make_migrations.py index a843b6e..782a682 100644 --- a/tests/test_make_migrations.py +++ b/tests/test_make_migrations.py @@ -45,6 +45,13 @@ method=PostgresPartitioningMethod.HASH, key="artist_id" ), ), + dict( + fields={"category": models.IntegerField(), "timestamp": models.DateTimeField()}, + partitioning_options=dict( + method=PostgresPartitioningMethod.LIST, key="category", + sub_key=["timestamp"], sub_method=PostgresPartitioningMethod.RANGE + ), + ), ], ) @postgres_patched_migrations() @@ -80,8 +87,8 @@ def test_make_migration_create_partitioned_model(fake_app, model_config): assert len(ops[0].bases) == 1 assert issubclass(ops[0].bases[0], PostgresPartitionedModel) - # make sure the partitioning options got copied correctly - assert ops[0].partitioning_options == model_config["partitioning_options"] + # make sure the partitioning options got copied correctly (not considering None values) + assert {k:v for k, v in ops[0].partitioning_options.items() if v is not None} == model_config["partitioning_options"] @postgres_patched_migrations() diff --git a/tests/test_partitioning_hierarchy.py b/tests/test_partitioning_hierarchy.py new file mode 100644 index 0000000..b1b5d12 --- /dev/null +++ b/tests/test_partitioning_hierarchy.py @@ -0,0 +1,263 @@ + +import django +import freezegun +import pytest +from dateutil.relativedelta import relativedelta +from django.core.exceptions import ImproperlyConfigured +from django.db import connection, models +from psqlextra.backend.schema import PostgresSchemaEditor +from psqlextra.contrib import partition_by_category_and_current_time +from psqlextra.partitioning import ( + PostgresPartitioningManager, +) +from psqlextra.types import PostgresPartitioningMethod + +from . import db_introspection +from .fake_model import define_fake_partitioned_model + + +def _get_partitioned_table(model): + return db_introspection.get_partitioned_table(model._meta.db_table) + +def _get_sub_partitions(model,partition_name): + return db_introspection.get_partitioned_table(model._meta.db_table + "_" + partition_name) + +@pytest.mark.skipif( + django.VERSION < (5, 2), + reason="Django < 5.2 doesn't implement composite primary keys", +) +@pytest.mark.postgres_version(lt=110000) +def test_partitioning_hierarchy_time_yearly_apply(): + """Tests whether automatically creating new partitions ahead yearly works + as expected.""" + + model = define_fake_partitioned_model( + fields={ + "id": models.AutoField(primary_key=False), + "category_id": models.IntegerField(), + "date": models.DateTimeField(), + "my_custom_pk": models.CompositePrimaryKey("id", "category_id", "date"), + }, + partitioning_options=dict( + key=["category_id"], + method=PostgresPartitioningMethod.LIST, + sub_key=["date"], + sub_method=PostgresPartitioningMethod.RANGE + ) + ) + + # pylint: disable=protected-access + # pylint: disable=no-member + assert isinstance(model._meta.pk, models.CompositePrimaryKey) + assert model._meta.pk.name == "my_custom_pk" + assert model._meta.pk.columns == ("id", "category_id", "date") + # pylint: enable=protected-access + # pylint: enable=no-member + + schema_editor = connection.schema_editor() + schema_editor.create_partitioned_model(model) + + with freezegun.freeze_time("2019-1-1"): + manager = PostgresPartitioningManager( + [partition_by_category_and_current_time(model, categories=[1, 2], years=1, count=2)] + ) + manager.plan().apply() + + table = _get_partitioned_table(model) + assert len(table.partitions) == 2 + assert table.partitions[0].name == "1" + assert table.partitions[1].name == "2" + + sub_table = _get_sub_partitions(model, table.partitions[0].name) + assert len(sub_table.partitions) == 2 + assert sub_table.partitions[0].full_name.endswith("1_2019") + assert sub_table.partitions[1].full_name.endswith("1_2020") + + sub_table = _get_sub_partitions(model, table.partitions[1].name) + assert len(sub_table.partitions) == 2 + assert sub_table.partitions[0].full_name.endswith("2_2019") + assert sub_table.partitions[1].full_name.endswith("2_2020") + + + +@pytest.mark.postgres_version(lt=110000) +@pytest.mark.parametrize( + "kwargs,timepoints", + [ + ( + dict(years=1, max_age=relativedelta(years=2)), + [("2019-1-1", 6), ("2020-1-1", 6), ("2021-1-1", 5)], + ), + ( + dict(months=1, max_age=relativedelta(months=1)), + [ + ("2019-1-1", 6), + ("2019-2-1", 5), + ("2019-2-28", 5), + ("2019-3-1", 4), + ], + ), + ( + dict(days=7, max_age=relativedelta(weeks=1)), + [ + ("2019-1-1", 6), + ("2019-1-4", 5), + ("2019-1-8", 5), + ("2019-1-15", 4), + ("2019-1-16", 4), + ], + ), + ], +) +def test_partitioning_hierarchy_time_delete(kwargs, timepoints): + """Tests whether partitions older than the specified max_age are + automatically deleted.""" + + model = define_fake_partitioned_model( + fields={ + "id": models.AutoField(primary_key=False), + "category_id": models.IntegerField(), + "date": models.DateTimeField(), + "my_custom_pk": models.CompositePrimaryKey("id", "category_id", "date"), + }, + partitioning_options=dict( + key=["category_id"], + method=PostgresPartitioningMethod.LIST, + sub_key=["date"], + sub_method=PostgresPartitioningMethod.RANGE + ) + ) + + schema_editor = connection.schema_editor() + schema_editor.create_partitioned_model(model) + + partition_kwargs = {"model": model, "categories": [1, 2], "count": 6, **kwargs} + + print("test with kwargs:", partition_kwargs) + manager = PostgresPartitioningManager( + [partition_by_category_and_current_time(**partition_kwargs)] + ) + + with freezegun.freeze_time(timepoints[0][0]): + manager.plan().apply() + + for (dt, partition_count) in timepoints: + with freezegun.freeze_time(dt): + manager.plan(skip_create=True).apply() + + table = _get_partitioned_table(model) + sub_table = _get_sub_partitions(model, table.partitions[0].name) + assert len(sub_table.partitions) == partition_count + + sub_table = _get_sub_partitions(model, table.partitions[1].name) + assert len(sub_table.partitions) == partition_count + + +def test_schema_editor_create_sub_partitioned_model_no_subkey(): + """Tests whether trying to create a partitioned model without a + partitioning key raises :see:ImproperlyConfigured as its not possible to + create a partitioned model without one and we cannot have a sane + default.""" + + model = define_fake_partitioned_model( + fields={ + "id": models.AutoField(primary_key=False), + "category_id": models.IntegerField(), + "date": models.DateTimeField(), + "my_custom_pk": models.CompositePrimaryKey("id", "category_id", "date"), + }, + partitioning_options=dict( + key=["category_id"], + method=PostgresPartitioningMethod.LIST, + sub_method=PostgresPartitioningMethod.RANGE + ) + ) + + schema_editor = PostgresSchemaEditor(connection) + + with pytest.raises(ImproperlyConfigured): + schema_editor.create_partitioned_model(model) + + +def test_schema_editor_create_sub_partitioned_model_no_field(): + """Tests whether trying to create a partitioned model without a + partitioning key raises :see:ImproperlyConfigured as its not possible to + create a partitioned model without one and we cannot have a sane + default.""" + + model = define_fake_partitioned_model( + fields={ + "id": models.AutoField(primary_key=False), + "category_id": models.IntegerField(), + "my_custom_pk": models.CompositePrimaryKey("id", "category_id", "date"), + }, + partitioning_options=dict( + key=["category_id"], + method=PostgresPartitioningMethod.LIST, + sub_method=PostgresPartitioningMethod.RANGE, + sub_key=["date"] + ) + ) + + schema_editor = PostgresSchemaEditor(connection) + + with pytest.raises(ImproperlyConfigured): + schema_editor.create_partitioned_model(model) + + + +@pytest.mark.skipif( + django.VERSION < (5, 2), + reason="Django < 5.2 doesn't implement composite primary keys", +) +@pytest.mark.postgres_version(lt=110000) +def test_partitioning_hierarchy_custom_name(): + """Tests whether custom name partitions work as expected.""" + + model = define_fake_partitioned_model( + fields={ + "id": models.AutoField(primary_key=False), + "category_id": models.IntegerField(), + "date": models.DateTimeField(), + "my_custom_pk": models.CompositePrimaryKey("id", "category_id", "date"), + }, + partitioning_options=dict( + key=["category_id"], + method=PostgresPartitioningMethod.LIST, + sub_key=["date"], + sub_method=PostgresPartitioningMethod.RANGE, + ) + ) + + + + + schema_editor = connection.schema_editor() + schema_editor.create_partitioned_model(model) + + with freezegun.freeze_time("2019-1-1"): + manager = PostgresPartitioningManager( + [partition_by_category_and_current_time( + model, + categories=[1, 2], + years=1, + count=2, + name_format=("category_%s", "time_%Y"), + )] + ) + manager.plan().apply() + + table = _get_partitioned_table(model) + assert len(table.partitions) == 2 + assert table.partitions[0].name == "category_1" + assert table.partitions[1].name == "category_2" + + sub_table = _get_sub_partitions(model, table.partitions[0].name) + assert len(sub_table.partitions) == 2 + assert sub_table.partitions[0].full_name.endswith("category_1_time_2019") + assert sub_table.partitions[1].full_name.endswith("category_1_time_2020") + + sub_table = _get_sub_partitions(model, table.partitions[1].name) + assert len(sub_table.partitions) == 2 + assert sub_table.partitions[0].full_name.endswith("category_2_time_2019") + assert sub_table.partitions[1].full_name.endswith("category_2_time_2020")