diff --git a/docs/source/table_partitioning.rst b/docs/source/table_partitioning.rst index 3ccc28c8..ef635537 100644 --- a/docs/source/table_partitioning.rst +++ b/docs/source/table_partitioning.rst @@ -20,6 +20,7 @@ The following partitioning methods are available: * ``PARTITION BY RANGE`` * ``PARTITION BY LIST`` +* ``PARTITION BY HASH`` .. note:: @@ -41,6 +42,7 @@ Inherit your model from :class:`psqlextra.models.PostgresPartitionedModel` and d * Use :attr:`psqlextra.types.PostgresPartitioningMethod.RANGE` to ``PARTITION BY RANGE`` * Use :attr:`psqlextra.types.PostgresPartitioningMethod.LIST` to ``PARTITION BY LIST`` +* Use :attr:`psqlextra.types.PostgresPartitioningMethod.HASH` to ``PARTITION BY HASH`` .. code-block:: python diff --git a/psqlextra/backend/introspection.py b/psqlextra/backend/introspection.py index cacffcb7..a85f27cd 100644 --- a/psqlextra/backend/introspection.py +++ b/psqlextra/backend/introspection.py @@ -5,6 +5,12 @@ from . import base_impl +PARTITIONING_STRATEGY_TO_METHOD = { + "r": PostgresPartitioningMethod.RANGE, + "l": PostgresPartitioningMethod.LIST, + "h": PostgresPartitioningMethod.HASH, +} + @dataclass class PostgresIntrospectedPartitionTable: @@ -64,9 +70,7 @@ def get_partitioned_tables( return [ PostgresIntrospectedPartitonedTable( name=row[0], - method=PostgresPartitioningMethod.RANGE - if row[1] == "r" - else PostgresPartitioningMethod.LIST, + method=PARTITIONING_STRATEGY_TO_METHOD[row[1]], key=self.get_partition_key(cursor, row[0]), partitions=self.get_partitions(cursor, row[0]), ) @@ -148,6 +152,7 @@ def get_partition_key(self, cursor, table_name: str) -> List[str]: CASE partstrat WHEN 'l' THEN 'list' WHEN 'r' THEN 'range' + WHEN 'h' THEN 'hash' END AS partition_strategy, Unnest(partattrs) column_index FROM pg_partitioned_table) pt diff --git a/psqlextra/backend/migrations/operations/__init__.py b/psqlextra/backend/migrations/operations/__init__.py index 3de992d1..f3d7551e 100644 --- a/psqlextra/backend/migrations/operations/__init__.py +++ b/psqlextra/backend/migrations/operations/__init__.py @@ -1,4 +1,5 @@ from .add_default_partition import PostgresAddDefaultPartition +from .add_hash_partition import PostgresAddHashPartition from .add_list_partition import PostgresAddListPartition from .add_range_partition import PostgresAddRangePartition from .apply_state import ApplyState @@ -6,6 +7,7 @@ from .create_partitioned_model import PostgresCreatePartitionedModel from .create_view_model import PostgresCreateViewModel from .delete_default_partition import PostgresDeleteDefaultPartition +from .delete_hash_partition import PostgresDeleteHashPartition from .delete_list_partition import PostgresDeleteListPartition from .delete_materialized_view_model import PostgresDeleteMaterializedViewModel from .delete_partitioned_model import PostgresDeletePartitionedModel @@ -14,12 +16,14 @@ __all__ = [ "ApplyState", - "PostgresAddRangePartition", + "PostgresAddHashPartition", "PostgresAddListPartition", + "PostgresAddRangePartition", "PostgresAddDefaultPartition", "PostgresDeleteDefaultPartition", - "PostgresDeleteRangePartition", + "PostgresDeleteHashPartition", "PostgresDeleteListPartition", + "PostgresDeleteRangePartition", "PostgresCreatePartitionedModel", "PostgresDeletePartitionedModel", "PostgresCreateViewModel", diff --git a/psqlextra/backend/migrations/operations/add_hash_partition.py b/psqlextra/backend/migrations/operations/add_hash_partition.py new file mode 100644 index 00000000..5ecdde22 --- /dev/null +++ b/psqlextra/backend/migrations/operations/add_hash_partition.py @@ -0,0 +1,74 @@ +from psqlextra.backend.migrations.state import PostgresHashPartitionState + +from .partition import PostgresPartitionOperation + + +class PostgresAddHashPartition(PostgresPartitionOperation): + """Adds a new hash partition to a :see:PartitionedPostgresModel. + + Each partition will hold the rows for which the hash value of the + partition key divided by the specified modulus will produce the + specified remainder. + """ + + def __init__( + self, model_name: str, name: str, modulus: int, remainder: int + ): + """Initializes new instance of :see:AddHashPartition. + Arguments: + model_name: + The name of the :see:PartitionedPostgresModel. + + name: + The name to give to the new partition table. + + modulus: + Integer value by which the key is divided. + + remainder: + The remainder of the hash value when divided by modulus. + """ + + super().__init__(model_name, name) + + self.modulus = modulus + self.remainder = remainder + + def state_forwards(self, app_label, state): + model = state.models[(app_label, self.model_name_lower)] + model.add_partition( + PostgresHashPartitionState( + app_label=app_label, + model_name=self.model_name, + name=self.name, + modulus=self.modulus, + remainder=self.remainder, + ) + ) + + state.reload_model(app_label, self.model_name_lower) + + def database_forwards(self, app_label, schema_editor, from_state, to_state): + model = to_state.apps.get_model(app_label, self.model_name) + if self.allow_migrate_model(schema_editor.connection.alias, model): + schema_editor.add_hash_partition( + model, self.name, self.modulus, self.remainder + ) + + def database_backwards( + self, app_label, schema_editor, from_state, to_state + ): + model = from_state.apps.get_model(app_label, self.model_name) + if self.allow_migrate_model(schema_editor.connection.alias, model): + schema_editor.delete_partition(model, self.name) + + def deconstruct(self): + name, args, kwargs = super().deconstruct() + + kwargs["modulus"] = self.modulus + kwargs["remainder"] = self.remainder + + return name, args, kwargs + + def describe(self) -> str: + return "Creates hash partition %s on %s" % (self.name, self.model_name) diff --git a/psqlextra/backend/migrations/operations/create_partitioned_model.py b/psqlextra/backend/migrations/operations/create_partitioned_model.py index 5d9bcc8d..e447965b 100644 --- a/psqlextra/backend/migrations/operations/create_partitioned_model.py +++ b/psqlextra/backend/migrations/operations/create_partitioned_model.py @@ -69,3 +69,19 @@ def describe(self): description = super().describe() description = description.replace("model", "partitioned model") return description + + def reduce(self, *args, **kwargs): + result = super().reduce(*args, **kwargs) + + # replace CreateModel operation with PostgresCreatePartitionedModel + if isinstance(result, list) and result: + for i, op in enumerate(result): + if isinstance(op, CreateModel): + _, args, kwargs = op.deconstruct() + result[i] = PostgresCreatePartitionedModel( + *args, + **kwargs, + partitioning_options=self.partitioning_options + ) + + return result diff --git a/psqlextra/backend/migrations/operations/delete_hash_partition.py b/psqlextra/backend/migrations/operations/delete_hash_partition.py new file mode 100644 index 00000000..3e511fcc --- /dev/null +++ b/psqlextra/backend/migrations/operations/delete_hash_partition.py @@ -0,0 +1,29 @@ +from .delete_partition import PostgresDeletePartition + + +class PostgresDeleteHashPartition(PostgresDeletePartition): + """Deletes a hash partition that's part of a. + + :see:PartitionedPostgresModel. + """ + + def database_backwards( + self, app_label, schema_editor, from_state, to_state + ): + model = to_state.apps.get_model(app_label, self.model_name) + model_state = to_state.models[(app_label, self.model_name_lower)] + + if self.allow_migrate_model(schema_editor.connection.alias, model): + partition_state = model_state.partitions[self.name] + schema_editor.add_hash_partition( + model, + partition_state.name, + partition_state.modulus, + partition_state.remainder, + ) + + def describe(self) -> str: + return "Deletes hash partition '%s' on %s" % ( + self.name, + self.model_name, + ) diff --git a/psqlextra/backend/migrations/operations/partition.py b/psqlextra/backend/migrations/operations/partition.py index 5bfa95f2..09cf0c57 100644 --- a/psqlextra/backend/migrations/operations/partition.py +++ b/psqlextra/backend/migrations/operations/partition.py @@ -26,3 +26,7 @@ def state_forwards(self, *args, **kwargs): def state_backwards(self, *args, **kwargs): pass + + def reduce(self, *args, **kwargs): + # PartitionOperation doesn't break migrations optimizations + return True diff --git a/psqlextra/backend/migrations/patched_autodetector.py b/psqlextra/backend/migrations/patched_autodetector.py index 284f055a..66da6734 100644 --- a/psqlextra/backend/migrations/patched_autodetector.py +++ b/psqlextra/backend/migrations/patched_autodetector.py @@ -8,7 +8,6 @@ DeleteModel, RemoveField, RenameField, - RunSQL, ) from django.db.migrations.autodetector import MigrationAutodetector from django.db.migrations.operations.base import Operation @@ -19,6 +18,7 @@ PostgresPartitionedModel, PostgresViewModel, ) +from psqlextra.types import PostgresPartitioningMethod from . import operations @@ -140,11 +140,12 @@ def add_create_partitioned_model( partitioning_options = model._partitioning_meta.original_attrs _, args, kwargs = operation.deconstruct() - self.add( - operations.PostgresAddDefaultPartition( - model_name=model.__name__, name="default" + if partitioning_options["method"] != PostgresPartitioningMethod.HASH: + self.add( + operations.PostgresAddDefaultPartition( + model_name=model.__name__, name="default" + ) ) - ) self.add( operations.PostgresCreatePartitionedModel( diff --git a/psqlextra/backend/migrations/state/__init__.py b/psqlextra/backend/migrations/state/__init__.py index bf18aa2d..83911094 100644 --- a/psqlextra/backend/migrations/state/__init__.py +++ b/psqlextra/backend/migrations/state/__init__.py @@ -1,5 +1,6 @@ from .materialized_view import PostgresMaterializedViewModelState from .partitioning import ( + PostgresHashPartitionState, PostgresListPartitionState, PostgresPartitionedModelState, PostgresPartitionState, @@ -10,6 +11,7 @@ __all__ = [ "PostgresPartitionState", "PostgresRangePartitionState", + "PostgresHashPartitionState", "PostgresListPartitionState", "PostgresPartitionedModelState", "PostgresViewModelState", diff --git a/psqlextra/backend/migrations/state/partitioning.py b/psqlextra/backend/migrations/state/partitioning.py index fffc7ccf..aef7a5e3 100644 --- a/psqlextra/backend/migrations/state/partitioning.py +++ b/psqlextra/backend/migrations/state/partitioning.py @@ -38,6 +38,24 @@ def __init__(self, app_label: str, model_name: str, name: str, values): self.values = values +class PostgresHashPartitionState(PostgresPartitionState): + """Represents the state of a hash partition for a + :see:PostgresPartitionedModel during a migration.""" + + def __init__( + self, + app_label: str, + model_name: str, + name: str, + modulus: int, + remainder: int, + ): + super().__init__(app_label, model_name, name) + + self.modulus = modulus + self.remainder = remainder + + class PostgresPartitionedModelState(PostgresModelState): """Represents the state of a :see:PostgresPartitionedModel in the migrations.""" diff --git a/psqlextra/backend/schema.py b/psqlextra/backend/schema.py index bef615ee..cc55dffa 100644 --- a/psqlextra/backend/schema.py +++ b/psqlextra/backend/schema.py @@ -38,6 +38,7 @@ class PostgresSchemaEditor(base_impl.schema_editor()): ) sql_partition_by = " PARTITION BY %s (%s)" sql_add_default_partition = "CREATE TABLE %s PARTITION OF %s DEFAULT" + sql_add_hash_partition = "CREATE TABLE %s PARTITION OF %s FOR VALUES WITH (MODULUS %s, REMAINDER %s)" sql_add_range_partition = ( "CREATE TABLE %s PARTITION OF %s FOR VALUES FROM (%s) TO (%s)" ) @@ -283,6 +284,52 @@ def add_list_partition( if comment: self.set_comment_on_table(table_name, comment) + def add_hash_partition( + self, + model: Model, + name: str, + modulus: int, + remainder: int, + comment: Optional[str] = None, + ) -> None: + """Creates a new hash partition for the specified partitioned model. + + Arguments: + model: + Partitioned model to create a partition for. + + name: + Name to give to the new partition. + Final name will be "{table_name}_{partition_name}" + + modulus: + Integer value by which the key is divided. + + remainder: + The remainder of the hash value when divided by modulus. + + comment: + Optionally, a comment to add on this partition table. + """ + + # asserts the model is a model set up for partitioning + self._partitioning_properties_for_model(model) + + table_name = self.create_partition_table_name(model, name) + + sql = self.sql_add_hash_partition % ( + self.quote_name(table_name), + self.quote_name(model._meta.db_table), + "%s", + "%s", + ) + + with transaction.atomic(): + self.execute(sql, (modulus, remainder)) + + if comment: + self.set_comment_on_table(table_name, comment) + def add_default_partition( self, model: Model, name: str, comment: Optional[str] = None ) -> None: diff --git a/psqlextra/types.py b/psqlextra/types.py index 80f4f1f6..6c846a16 100644 --- a/psqlextra/types.py +++ b/psqlextra/types.py @@ -35,3 +35,4 @@ class PostgresPartitioningMethod(StrEnum): RANGE = "range" LIST = "list" + HASH = "hash" diff --git a/setup.py b/setup.py index 6f089740..f67d7b2f 100644 --- a/setup.py +++ b/setup.py @@ -61,9 +61,7 @@ def run(self): python_requires=">=3.6", install_requires=[ "Django>=2.0", - "enforce>=0.3.4,<=1.0.0", "python-dateutil>=2.8.0,<=3.0.0", - "structlog>=19,<=20.1.0", "ansimarkup>=1.4.0,<=2.0.0", ], extras_require={ diff --git a/tests/test_make_migrations.py b/tests/test_make_migrations.py index 21b26b41..6f63a0d6 100644 --- a/tests/test_make_migrations.py +++ b/tests/test_make_migrations.py @@ -1,3 +1,4 @@ +import django import pytest from django.apps import apps @@ -15,6 +16,7 @@ from .fake_model import ( define_fake_materialized_view_model, + define_fake_model, define_fake_partitioned_model, define_fake_view_model, get_fake_model, @@ -37,6 +39,12 @@ method=PostgresPartitioningMethod.RANGE, key="timestamp" ), ), + dict( + fields={"artist_id": models.IntegerField()}, + partitioning_options=dict( + method=PostgresPartitioningMethod.HASH, key="artist_id" + ), + ), ], ) @postgres_patched_migrations() @@ -48,14 +56,25 @@ def test_make_migration_create_partitioned_model(fake_app, model_config): **model_config, meta_options=dict(app_label=fake_app.name) ) - migration = make_migration(model._meta.app_label) + migration = make_migration(fake_app.name) ops = migration.operations - - # should have one operation to create the partitioned model - # and one more to add a default partition - assert len(ops) == 2 - assert isinstance(ops[0], operations.PostgresCreatePartitionedModel) - assert isinstance(ops[1], operations.PostgresAddDefaultPartition) + method = model_config["partitioning_options"]["method"] + + if method == PostgresPartitioningMethod.HASH: + # should have one operation to create the partitioned model + # and no default partition + assert len(ops) == 1 + assert isinstance(ops[0], operations.PostgresCreatePartitionedModel) + else: + # should have one operation to create the partitioned model + # and one more to add a default partition + assert len(ops) == 2 + assert isinstance(ops[0], operations.PostgresCreatePartitionedModel) + assert isinstance(ops[1], operations.PostgresAddDefaultPartition) + + # make sure the default partition is named "default" + assert ops[1].model_name == model.__name__ + assert ops[1].name == "default" # make sure the base is set correctly assert len(ops[0].bases) == 1 @@ -64,10 +83,6 @@ def test_make_migration_create_partitioned_model(fake_app, model_config): # make sure the partitioning options got copied correctly assert ops[0].partitioning_options == model_config["partitioning_options"] - # make sure the default partition is named "default" - assert ops[1].model_name == model.__name__ - assert ops[1].name == "default" - @postgres_patched_migrations() def test_make_migration_create_view_model(fake_app): @@ -182,3 +197,47 @@ def test_make_migration_field_operations_view_models( ) assert isinstance(migration.operations[0], operations.ApplyState) assert isinstance(migration.operations[0].state_operation, RemoveField) + + +@pytest.mark.skipif( + django.VERSION < (2, 2), + reason="Django < 2.2 doesn't implement left-to-right migration optimizations", +) +@pytest.mark.parametrize("method", PostgresPartitioningMethod.all()) +@postgres_patched_migrations() +def test_autodetect_fk_issue(fake_app, method): + """Test whether Django can perform ForeignKey optimization. + + Fixes https://github.com/SectorLabs/django-postgres-extra/issues/123 for Django >= 2.2 + """ + meta_options = {"app_label": fake_app.name} + partitioning_options = {"method": method, "key": "artist_id"} + + artist_model_fields = {"name": models.TextField()} + Artist = define_fake_model(artist_model_fields, meta_options=meta_options) + + from_state = ProjectState.from_apps(apps) + + album_model_fields = { + "name": models.TextField(), + "artist": models.ForeignKey( + to=Artist.__name__, on_delete=models.CASCADE + ), + } + + define_fake_partitioned_model( + album_model_fields, + partitioning_options=partitioning_options, + meta_options=meta_options, + ) + + migration = make_migration(fake_app.name, from_state=from_state) + ops = migration.operations + + if method == PostgresPartitioningMethod.HASH: + assert len(ops) == 1 + assert isinstance(ops[0], operations.PostgresCreatePartitionedModel) + else: + assert len(ops) == 2 + assert isinstance(ops[0], operations.PostgresCreatePartitionedModel) + assert isinstance(ops[1], operations.PostgresAddDefaultPartition) diff --git a/tests/test_migration_operations.py b/tests/test_migration_operations.py index f34aef9e..6a5b4274 100644 --- a/tests/test_migration_operations.py +++ b/tests/test_migration_operations.py @@ -63,9 +63,14 @@ def _create_model(method): if method == PostgresPartitioningMethod.RANGE: key.append("timestamp") fields.append(("timestamp", models.DateTimeField())) - else: + elif method == PostgresPartitioningMethod.LIST: key.append("category") fields.append(("category", models.TextField())) + elif method == PostgresPartitioningMethod.HASH: + key.append("artist_id") + fields.append(("artist_id", models.IntegerField())) + else: + raise NotImplementedError return operations.PostgresCreatePartitionedModel( "test", @@ -150,6 +155,12 @@ def test_migration_operations_delete_partitioned_table(method, create_model): model_name="test", name="pt1", values=["car", "boat"] ), ), + ( + PostgresPartitioningMethod.HASH, + operations.PostgresAddHashPartition( + model_name="test", name="pt1", modulus=3, remainder=0 + ), + ), ], ) def test_migration_operations_add_partition( @@ -206,6 +217,15 @@ def test_migration_operations_add_partition( model_name="test", name="pt1" ), ), + ( + PostgresPartitioningMethod.HASH, + operations.PostgresAddHashPartition( + model_name="test", name="pt1", modulus=3, remainder=0 + ), + operations.PostgresDeleteHashPartition( + model_name="test", name="pt1" + ), + ), ], ) def test_migration_operations_add_delete_partition(