From 86517711b24b9edd6b4f5e519ca6b796cb4cc403 Mon Sep 17 00:00:00 2001 From: Alexandru Kis Date: Thu, 9 Nov 2017 18:22:42 +0200 Subject: [PATCH 1/3] Added a conditionalUniqueIndex. --- docs/index.md | 3 + docs/indexes.md | 46 ++++++++++++++ mkdocs.yml | 1 + psqlextra/indexes/__init__.py | 5 ++ psqlextra/indexes/conditional_unique_index.py | 39 ++++++++++++ tests/migrations.py | 18 +++++- tests/test_conditional_unique_index.py | 62 +++++++++++++++++++ 7 files changed, 171 insertions(+), 3 deletions(-) create mode 100644 docs/indexes.md create mode 100644 psqlextra/indexes/__init__.py create mode 100644 psqlextra/indexes/conditional_unique_index.py create mode 100644 tests/test_conditional_unique_index.py diff --git a/docs/index.md b/docs/index.md index b0bfb8bb..eb03c733 100644 --- a/docs/index.md +++ b/docs/index.md @@ -19,6 +19,9 @@ Explore the documentation to learn about all features: * [Materialized views](/materialized_views) +* [Indexes](/indexes) + * [ConditionalUniqueIndex](/indexes/#conditional-unique-index) + ## Installation 1. Install the package from PyPi: diff --git a/docs/indexes.md b/docs/indexes.md new file mode 100644 index 00000000..56ec2612 --- /dev/null +++ b/docs/indexes.md @@ -0,0 +1,46 @@ +## Conditional Unique Index + +The `ConditionalUniqueIndex` lets you create partial unique indexes in case you ever need `unique together` constraints +on nullable columns. + +e.g. + +Before: + +``` +from django.db import models + +class Model(models.Model): + class Meta: + unique_together = ['a', 'b''] + + a = models.ForeignKey('some_model', null=True) + b = models.ForeignKey('some_other_model') + +# Works like a charm! +b = B() +Model.objects.create(a=None, b=b) +Model.objects.create(a=None, b=b) +``` + +After: + +``` +from django.db import models +from from psqlextra.indexes import ConditionalUniqueIndex + +class Model(models.Model): + class Meta: + indexes = [ + ConditionalUniqueIndex(fields=['a', 'b'], condition='"a" IS NOT NULL'), + ConditionalUniqueIndex(fields=['b'], condition='"a" IS NULL') + ] + + a = models.ForeignKey('some_model', null=True) + b = models.ForeignKey('some_other_model') + +# Integrity Error! +b = B() +Model.objects.create(a=None, b=b) +Model.objects.create(a=None, b=b) +``` diff --git a/mkdocs.yml b/mkdocs.yml index d75cd230..b541dbc2 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -5,3 +5,4 @@ pages: - HStore: hstore.md - Signals: signals.md - Materialized Views: materialized_views.md +- Indexes: indexes.md diff --git a/psqlextra/indexes/__init__.py b/psqlextra/indexes/__init__.py new file mode 100644 index 00000000..6497ca4c --- /dev/null +++ b/psqlextra/indexes/__init__.py @@ -0,0 +1,5 @@ +from .conditional_unique_index import ConditionalUniqueIndex + +__all__ = [ + 'ConditionalUniqueIndex' +] diff --git a/psqlextra/indexes/conditional_unique_index.py b/psqlextra/indexes/conditional_unique_index.py new file mode 100644 index 00000000..57532335 --- /dev/null +++ b/psqlextra/indexes/conditional_unique_index.py @@ -0,0 +1,39 @@ +from django.db.models.indexes import Index + + +class ConditionalUniqueIndex(Index): + """ + Creates a partial unique index based on a given condition. + + Useful, for example, if you need unique combination of foreign keys, but you might want to include + NULL as a valid value. In that case, you can just use: + >>> class Meta: + ... indexes = [ + ... ConditionalUniqueIndex(fields=['a', 'b', 'c'], condition='c IS NOT NULL'), + ... ConditionalUniqueIndex(fields=['a', 'b'], condition='c IS NULL') + ... ] + """ + + sql_create_index = "CREATE UNIQUE INDEX %(name)s ON %(table)s (%(columns)s)%(extra)s WHERE %(condition)s" + + def __init__(self, condition: str, fields=[], name=None): + """Initializes a new instance of :see:ConditionalUniqueIndex.""" + + super().__init__(fields=fields, name=name) + self.condition = condition + + def create_sql(self, model, schema_editor, using=''): + """Creates the actual SQL used when applying the migration.""" + + sql_create_index = self.sql_create_index + sql_parameters = { + **Index.get_sql_create_template_values(self, model, schema_editor, using), + 'condition': self.condition + } + return sql_create_index % sql_parameters + + def deconstruct(self): + """Serializes the :see:ConditionalUniqueIndex for the migrations file.""" + path = '%s.%s' % (self.__class__.__module__, self.__class__.__name__) + path = path.replace('django.db.models.indexes', 'django.db.models') + return path, (), {'fields': self.fields, 'name': self.name, 'condition': self.condition} diff --git a/tests/migrations.py b/tests/migrations.py index 85cd2e35..7cb252b6 100644 --- a/tests/migrations.py +++ b/tests/migrations.py @@ -357,17 +357,29 @@ def make_migrations(self): self.project_state = new_project_state return migration - def migrate(self): - """Executes the recorded migrations.""" + def migrate(self, *filters: List[str]): + """ + Executes the recorded migrations. + Arguments: + filters: List of strings to filter SQL statements on. + + Returns: + The filtered calls of every migration + """ + + calls_for_migrations = [] while len(self.migrations) > 0: migration = self.migrations.pop() - with connection.schema_editor() as schema_editor: + with filtered_schema_editor(*filters) as (schema_editor, calls): migration_executor = MigrationExecutor(schema_editor.connection) migration_executor.apply_migration( self.project_state, migration ) + calls_for_migrations.append(calls) + + return calls_for_migrations def _generate_random_name(self): return str(uuid.uuid4()).replace('-', '')[:8] diff --git a/tests/test_conditional_unique_index.py b/tests/test_conditional_unique_index.py new file mode 100644 index 00000000..4225f25c --- /dev/null +++ b/tests/test_conditional_unique_index.py @@ -0,0 +1,62 @@ +from psqlextra.indexes import ConditionalUniqueIndex +from .migrations import MigrationSimulator, filtered_schema_editor + +from django.db import models +from django.db.migrations import AddIndex, CreateModel + +def test_deconstruct(): + """Tests whether the :see:HStoreField's deconstruct() + method works properly.""" + + original_kwargs = dict(condition='field IS NULL', name='great_index', fields=['field', 'build']) + _, _, new_kwargs = ConditionalUniqueIndex(**original_kwargs).deconstruct() + + for key, value in original_kwargs.items(): + assert new_kwargs[key] == value + + +def test_migrations(): + """Tests whether the migrations are properly generated and executed.""" + + simulator = MigrationSimulator() + + Model = simulator.define_model( + fields={ + 'id': models.IntegerField(primary_key=True), + 'name': models.CharField(max_length=255), + 'other_name': models.CharField(max_length=255) + }, + meta_options={ + 'indexes': [ + ConditionalUniqueIndex( + fields=['name', 'other_name'], + condition='"name" IS NOT NULL', + name='index1' + ), + ConditionalUniqueIndex( + fields=['other_name'], + condition='"name" IS NULL', + name='index2' + ) + ] + } + ) + + migration = simulator.make_migrations() + assert len(migration.operations) == 3 + + operations = migration.operations + assert isinstance(operations[0], CreateModel) + + for operation in operations[1:]: + assert isinstance(operation, AddIndex) + + calls = [call[0] for _, call, _ in simulator.migrate('CREATE UNIQUE INDEX')[0]['CREATE UNIQUE INDEX']] + + db_table = Model._meta.db_table + assert calls[0] == 'CREATE UNIQUE INDEX "index1" ON "{0}" ("name", "other_name") WHERE "name" IS NOT NULL'.format( + db_table + ) + assert calls[1] == 'CREATE UNIQUE INDEX "index2" ON "{0}" ("other_name") WHERE "name" IS NULL'.format( + db_table + ) From db18d1201c1a94f649ea27062fc566d6a823012e Mon Sep 17 00:00:00 2001 From: Alexandru Kis Date: Fri, 10 Nov 2017 10:14:08 +0200 Subject: [PATCH 2/3] Tested the ConditionalUniqueIndex actually works. --- psqlextra/indexes/conditional_unique_index.py | 4 ++-- tests/test_conditional_unique_index.py | 18 +++++++++++++++--- 2 files changed, 17 insertions(+), 5 deletions(-) diff --git a/psqlextra/indexes/conditional_unique_index.py b/psqlextra/indexes/conditional_unique_index.py index 57532335..d8cfe0ed 100644 --- a/psqlextra/indexes/conditional_unique_index.py +++ b/psqlextra/indexes/conditional_unique_index.py @@ -9,8 +9,8 @@ class ConditionalUniqueIndex(Index): NULL as a valid value. In that case, you can just use: >>> class Meta: ... indexes = [ - ... ConditionalUniqueIndex(fields=['a', 'b', 'c'], condition='c IS NOT NULL'), - ... ConditionalUniqueIndex(fields=['a', 'b'], condition='c IS NULL') + ... ConditionalUniqueIndex(fields=['a', 'b', 'c'], condition='"c" IS NOT NULL'), + ... ConditionalUniqueIndex(fields=['a', 'b'], condition='"c" IS NULL') ... ] """ diff --git a/tests/test_conditional_unique_index.py b/tests/test_conditional_unique_index.py index 4225f25c..ea049294 100644 --- a/tests/test_conditional_unique_index.py +++ b/tests/test_conditional_unique_index.py @@ -1,7 +1,9 @@ +import pytest + from psqlextra.indexes import ConditionalUniqueIndex -from .migrations import MigrationSimulator, filtered_schema_editor +from .migrations import MigrationSimulator -from django.db import models +from django.db import models, IntegrityError, transaction from django.db.migrations import AddIndex, CreateModel def test_deconstruct(): @@ -23,7 +25,7 @@ def test_migrations(): Model = simulator.define_model( fields={ 'id': models.IntegerField(primary_key=True), - 'name': models.CharField(max_length=255), + 'name': models.CharField(max_length=255, null=True), 'other_name': models.CharField(max_length=255) }, meta_options={ @@ -60,3 +62,13 @@ def test_migrations(): assert calls[1] == 'CREATE UNIQUE INDEX "index2" ON "{0}" ("other_name") WHERE "name" IS NULL'.format( db_table ) + + with transaction.atomic(): + Model.objects.create(id=1, name="name", other_name="other_name") + with pytest.raises(IntegrityError): + Model.objects.create(id=2, name="name", other_name="other_name") + + with transaction.atomic(): + Model.objects.create(id=1, name=None, other_name="other_name") + with pytest.raises(IntegrityError): + Model.objects.create(id=2, name=None, other_name="other_name") From 397382b8a91841b34d556614443001664f3f143c Mon Sep 17 00:00:00 2001 From: Alexandru Kis Date: Fri, 10 Nov 2017 10:20:10 +0200 Subject: [PATCH 3/3] Fixed faulty indentation. --- tests/test_conditional_unique_index.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/tests/test_conditional_unique_index.py b/tests/test_conditional_unique_index.py index ea049294..79ef874b 100644 --- a/tests/test_conditional_unique_index.py +++ b/tests/test_conditional_unique_index.py @@ -6,6 +6,7 @@ from django.db import models, IntegrityError, transaction from django.db.migrations import AddIndex, CreateModel + def test_deconstruct(): """Tests whether the :see:HStoreField's deconstruct() method works properly.""" @@ -28,20 +29,20 @@ def test_migrations(): 'name': models.CharField(max_length=255, null=True), 'other_name': models.CharField(max_length=255) }, - meta_options={ - 'indexes': [ - ConditionalUniqueIndex( + meta_options={ + 'indexes': [ + ConditionalUniqueIndex( fields=['name', 'other_name'], condition='"name" IS NOT NULL', name='index1' - ), - ConditionalUniqueIndex( + ), + ConditionalUniqueIndex( fields=['other_name'], condition='"name" IS NULL', name='index2' ) - ] - } + ] + } ) migration = simulator.make_migrations()