From e52a07586b39dbc7f23c3244d961823349fb9bbd Mon Sep 17 00:00:00 2001 From: pctablet505 Date: Wed, 15 Jul 2026 15:15:57 +0000 Subject: [PATCH 1/2] feat(indexes): add nulls_not_distinct option to PostgreSQLIndex Add unique and nulls_not_distinct parameters to PostgreSQLIndex so users can create PostgreSQL unique indexes that treat NULL values as equal. This implements the direction from abondar in #1641: expose the option on tortoise.contrib.postgres.indexes.PostgreSQLIndex instead of changing unique_together behavior. Fixes #1641 --- tests/migrations/test_schema_editor_sql.py | 69 ++++++++++++++++++- tests/schema/models_postgres_unique_index.py | 18 +++++ tests/schema/test_generate_schema.py | 28 ++++++++ .../base_postgres/schema_generator.py | 36 +++++++++- tortoise/contrib/postgres/indexes.py | 58 +++++++++++++++- tortoise/migrations/schema_editor/base.py | 5 +- .../migrations/schema_editor/base_postgres.py | 24 ++++--- 7 files changed, 224 insertions(+), 14 deletions(-) create mode 100644 tests/schema/models_postgres_unique_index.py diff --git a/tests/migrations/test_schema_editor_sql.py b/tests/migrations/test_schema_editor_sql.py index 67a9dac31..035f52166 100644 --- a/tests/migrations/test_schema_editor_sql.py +++ b/tests/migrations/test_schema_editor_sql.py @@ -7,7 +7,7 @@ from tests.utils.fake_client import FakeClient from tortoise import fields from tortoise.contrib.postgres.fields import TSVectorField -from tortoise.contrib.postgres.indexes import GinIndex +from tortoise.contrib.postgres.indexes import GinIndex, PostgreSQLIndex from tortoise.indexes import Index from tortoise.migrations.schema_editor.base import BaseSchemaEditor from tortoise.migrations.schema_editor.base_postgres import BasePostgresSchemaEditor @@ -192,6 +192,73 @@ class Meta: ) == client.executed[0] +@pytest.mark.asyncio +async def test_add_unique_index_generates_nulls_not_distinct_sql() -> None: + class Customer(Model): + id = fields.IntField(pk=True) + shop_id = fields.IntField() + phone_number = fields.CharField(max_length=20) + deleted_at = fields.DatetimeField(null=True) + + class Meta: + table = "customer" + app = "models" + indexes = [ + PostgreSQLIndex( + fields=("shop_id", "phone_number", "deleted_at"), + unique=True, + nulls_not_distinct=True, + ) + ] + + client = FakeClient("postgres", inline_comment=False) + editor = BasePostgresSchemaEditor(client) + + index = cast(Index, Customer._meta.indexes[0]) + await editor.add_index(Customer, index) + + assert client.executed + expected_name = editor._generate_index_name("uidx", Customer, ["shop_id", "phone_number", "deleted_at"]) + assert ( + f'CREATE UNIQUE INDEX "{expected_name}" ON "customer"' + f' ("shop_id", "phone_number", "deleted_at") NULLS NOT DISTINCT;' + ) == client.executed[0] + + +@pytest.mark.asyncio +async def test_add_partial_unique_index_generates_nulls_not_distinct_before_where() -> None: + class Customer(Model): + id = fields.IntField(pk=True) + shop_id = fields.IntField() + phone_number = fields.CharField(max_length=20) + deleted_at = fields.DatetimeField(null=True) + + class Meta: + table = "customer" + app = "models" + indexes = [ + PostgreSQLIndex( + fields=("shop_id", "phone_number", "deleted_at"), + unique=True, + nulls_not_distinct=True, + condition={"shop_id": 1}, + ) + ] + + client = FakeClient("postgres", inline_comment=False) + editor = BasePostgresSchemaEditor(client) + + index = cast(Index, Customer._meta.indexes[0]) + await editor.add_index(Customer, index) + + assert client.executed + expected_name = editor._generate_index_name("uidx", Customer, ["shop_id", "phone_number", "deleted_at"]) + assert ( + f'CREATE UNIQUE INDEX "{expected_name}" ON "customer"' + f' ("shop_id", "phone_number", "deleted_at") NULLS NOT DISTINCT WHERE shop_id = 1;' + ) == client.executed[0] + + @pytest.mark.asyncio async def test_alter_generated_field_raises() -> None: class OldSearchDocument(Model): diff --git a/tests/schema/models_postgres_unique_index.py b/tests/schema/models_postgres_unique_index.py new file mode 100644 index 000000000..fca5b599a --- /dev/null +++ b/tests/schema/models_postgres_unique_index.py @@ -0,0 +1,18 @@ +from tortoise import Model, fields + +from tortoise.contrib.postgres.indexes import PostgreSQLIndex + + +class Customer(Model): + shop_id = fields.IntField() + phone_number = fields.CharField(max_length=20) + deleted_at = fields.DatetimeField(null=True) + + class Meta: + indexes = [ + PostgreSQLIndex( + fields=("shop_id", "phone_number", "deleted_at"), + unique=True, + nulls_not_distinct=True, + ), + ] diff --git a/tests/schema/test_generate_schema.py b/tests/schema/test_generate_schema.py index 63098bce1..28af89e8a 100644 --- a/tests/schema/test_generate_schema.py +++ b/tests/schema/test_generate_schema.py @@ -1928,6 +1928,34 @@ async def test_psycopg_index_safe(): await _teardown_tortoise() +@pytest.mark.asyncio +async def test_asyncpg_unique_index_nulls_not_distinct(): + await _reset_tortoise() + try: + await _init_for_asyncpg("tests.schema.models_postgres_unique_index") + sql = get_schema_sql(connections.get("default"), safe=False) + assert ( + 'CREATE UNIQUE INDEX "uidx_customer_shop_id_c9ba87" ON "customer"' + ' ("shop_id", "phone_number", "deleted_at") NULLS NOT DISTINCT;' in sql + ) + finally: + await _teardown_tortoise() + + +@pytest.mark.asyncio +async def test_psycopg_unique_index_nulls_not_distinct(): + await _reset_tortoise() + try: + await _init_for_psycopg("tests.schema.models_postgres_unique_index") + sql = get_schema_sql(connections.get("default"), safe=False) + assert ( + 'CREATE UNIQUE INDEX "uidx_customer_shop_id_c9ba87" ON "customer"' + ' ("shop_id", "phone_number", "deleted_at") NULLS NOT DISTINCT;' in sql + ) + finally: + await _teardown_tortoise() + + @pytest.mark.asyncio async def test_psycopg_m2m_no_auto_create(): await _reset_tortoise() diff --git a/tortoise/backends/base_postgres/schema_generator.py b/tortoise/backends/base_postgres/schema_generator.py index fdbe29207..ce3c44c47 100644 --- a/tortoise/backends/base_postgres/schema_generator.py +++ b/tortoise/backends/base_postgres/schema_generator.py @@ -14,7 +14,7 @@ class BasePostgresSchemaGenerator(BaseSchemaGenerator): DIALECT = "postgres" INDEX_CREATE_TEMPLATE = ( - 'CREATE INDEX {exists}"{index_name}" ON {table_name} {index_type}({fields}){extra};' + 'CREATE INDEX {exists}"{index_name}" ON {table_name} {index_type}({fields}){nulls_not_distinct}{extra};' ) UNIQUE_INDEX_CREATE_TEMPLATE = INDEX_CREATE_TEMPLATE.replace("INDEX", "UNIQUE INDEX") TABLE_COMMENT_TEMPLATE = "COMMENT ON TABLE {table} IS '{comment}';" @@ -79,10 +79,40 @@ def _get_index_sql( index_name: str | None = None, index_type: str | None = None, extra: str | None = None, + unique: bool = False, + nulls_not_distinct: bool = False, ) -> str: if index_type: index_type = f"USING {index_type}" - return super()._get_index_sql( - model, field_names, safe, index_name=index_name, index_type=index_type, extra=extra + nulls_not_distinct_sql = " NULLS NOT DISTINCT" if unique and nulls_not_distinct else "" + template = self.UNIQUE_INDEX_CREATE_TEMPLATE if unique else self.INDEX_CREATE_TEMPLATE + prefix = "uidx" if unique else "idx" + + return template.format( + exists="IF NOT EXISTS " if safe else "", + index_name=index_name or self._get_index_name(prefix, model, field_names), + index_type=f"{index_type} " if index_type else "", + table_name=self._qualify_table_name(model._meta.db_table, model._meta.schema), + fields=self._format_index_fields(field_names), + nulls_not_distinct=nulls_not_distinct_sql, + extra=f"{extra}" if extra else "", + ) + + def _get_unique_index_sql( + self, + exists: str, + table_name: str, + field_names: Sequence[str], + schema: str | None = None, + ) -> str: + index_name = self._get_index_name("uidx", table_name, field_names) + return self.UNIQUE_INDEX_CREATE_TEMPLATE.format( + exists=exists, + index_name=index_name, + index_type="", + table_name=self._qualify_table_name(table_name, schema), + fields=", ".join([self.quote(f) for f in field_names]), + nulls_not_distinct="", + extra="", ) diff --git a/tortoise/contrib/postgres/indexes.py b/tortoise/contrib/postgres/indexes.py index 499ce11ef..c8019203c 100644 --- a/tortoise/contrib/postgres/indexes.py +++ b/tortoise/contrib/postgres/indexes.py @@ -1,8 +1,64 @@ +from __future__ import annotations + +from typing import Any + +from pypika_tortoise.terms import Term + +from tortoise.expressions import Expression from tortoise.indexes import PartialIndex class PostgreSQLIndex(PartialIndex): - pass + """ + Base class for PostgreSQL-specific indexes. + + :param expressions: The expressions on which the index is desired. + :param fields: A tuple or list of field names on which the index is desired. + :param name: The name of the index. + :param condition: Optional WHERE condition for partial indexes. + :param unique: Whether the index should enforce uniqueness. + :param nulls_not_distinct: For unique indexes, treat NULL values as equal. + :raises ValueError: If params conflict. + """ + + def __init__( + self, + *expressions: Term | Expression, + fields: tuple[str, ...] | list[str] | None = None, + name: str | None = None, + condition: dict | None = None, + unique: bool = False, + nulls_not_distinct: bool = False, + ) -> None: + super().__init__(*expressions, fields=fields, name=name, condition=condition) + self.unique = unique + self.nulls_not_distinct = nulls_not_distinct + + def get_sql(self, schema_generator, model, safe): + return schema_generator._get_index_sql( + model, + self.field_names, + safe, + index_name=self.name, + index_type=self.INDEX_TYPE, + extra=self.extra, + unique=self.unique, + nulls_not_distinct=self.nulls_not_distinct, + ) + + def describe(self) -> dict: + result = super().describe() + result["unique"] = self.unique + result["nulls_not_distinct"] = self.nulls_not_distinct + return result + + def deconstruct(self) -> tuple[str, list[Any], dict[str, Any]]: + path, args, kwargs = super().deconstruct() + if self.unique: + kwargs["unique"] = self.unique + if self.nulls_not_distinct: + kwargs["nulls_not_distinct"] = self.nulls_not_distinct + return path, args, kwargs class BloomIndex(PostgreSQLIndex): diff --git a/tortoise/migrations/schema_editor/base.py b/tortoise/migrations/schema_editor/base.py index 4dda5ec91..299011ef3 100644 --- a/tortoise/migrations/schema_editor/base.py +++ b/tortoise/migrations/schema_editor/base.py @@ -738,7 +738,8 @@ def _index_name_for_model(self, model: type[Model], index: Index) -> str: if index.name: return index.name index.resolve_expressions(model) - return self._generate_index_name("idx", model, list(index.field_names)) + prefix = "uidx" if getattr(index, "unique", False) else "idx" + return self._generate_index_name(prefix, model, list(index.field_names)) def _constraint_name_for_model(self, model: type[Model], constraint: UniqueConstraint) -> str: if constraint.name: @@ -809,6 +810,8 @@ async def add_index(self, model: type[Model], index: Index) -> None: index_name=self._index_name_for_model(model, index), index_type=index.INDEX_TYPE, extra=index.extra, + unique=getattr(index, "unique", False), + nulls_not_distinct=getattr(index, "nulls_not_distinct", False), ) if index_sql: await self._run_sql(index_sql) diff --git a/tortoise/migrations/schema_editor/base_postgres.py b/tortoise/migrations/schema_editor/base_postgres.py index cd2c156e1..989226d59 100644 --- a/tortoise/migrations/schema_editor/base_postgres.py +++ b/tortoise/migrations/schema_editor/base_postgres.py @@ -10,7 +10,7 @@ class BasePostgresSchemaEditor(BaseSchemaEditor): DIALECT = "postgres" INDEX_CREATE_TEMPLATE = ( - 'CREATE INDEX "{index_name}" ON {table_name} {index_type}({fields}){extra};' + 'CREATE INDEX "{index_name}" ON {table_name} {index_type}({fields}){nulls_not_distinct}{extra};' ) UNIQUE_INDEX_CREATE_TEMPLATE = INDEX_CREATE_TEMPLATE.replace("INDEX", "UNIQUE INDEX") TABLE_COMMENT_TEMPLATE = "COMMENT ON TABLE {table} IS '{comment}';" @@ -78,16 +78,23 @@ def _get_index_sql( index_name: str | None = None, index_type: str | None = None, extra: str | None = None, + unique: bool = False, + nulls_not_distinct: bool = False, ) -> str: if index_type: index_type = f"USING {index_type}" - return super()._get_index_sql( - model, - list(field_names), - safe, - index_name=index_name, - index_type=index_type, - extra=extra, + + nulls_not_distinct_sql = " NULLS NOT DISTINCT" if unique and nulls_not_distinct else "" + template = self.UNIQUE_INDEX_CREATE_TEMPLATE if unique else self.INDEX_CREATE_TEMPLATE + prefix = "uidx" if unique else "idx" + + return template.format( + index_name=index_name or self._generate_index_name(prefix, model, field_names), + table_name=self._qualify_table_name(model._meta.db_table, model._meta.schema), + index_type=f"{index_type} " if index_type else "", + fields=self._format_index_fields(list(field_names)), + nulls_not_distinct=nulls_not_distinct_sql, + extra=f"{extra}" if extra else "", ) def _escape_default_value(self, default: object) -> str: @@ -103,6 +110,7 @@ def _get_unique_index_sql( table_name=self._qualify_table_name(table_name, schema), index_type="", fields=", ".join([self.quote(f) for f in field_names]), + nulls_not_distinct="", extra="", ) From 7f13df6b75b044190bde7e78c49ccc5df8bf3aa1 Mon Sep 17 00:00:00 2001 From: pctablet505 Date: Thu, 16 Jul 2026 11:20:16 +0530 Subject: [PATCH 2/2] style: apply ruff format and import sorting Run make style to satisfy the ruff format and isort (I) checks so the lint gate passes. --- tests/migrations/test_schema_editor_sql.py | 8 ++++++-- tests/schema/models_postgres_unique_index.py | 1 - tortoise/backends/base_postgres/schema_generator.py | 4 +--- tortoise/migrations/schema_editor/base_postgres.py | 4 +--- 4 files changed, 8 insertions(+), 9 deletions(-) diff --git a/tests/migrations/test_schema_editor_sql.py b/tests/migrations/test_schema_editor_sql.py index 035f52166..93e98482d 100644 --- a/tests/migrations/test_schema_editor_sql.py +++ b/tests/migrations/test_schema_editor_sql.py @@ -218,7 +218,9 @@ class Meta: await editor.add_index(Customer, index) assert client.executed - expected_name = editor._generate_index_name("uidx", Customer, ["shop_id", "phone_number", "deleted_at"]) + expected_name = editor._generate_index_name( + "uidx", Customer, ["shop_id", "phone_number", "deleted_at"] + ) assert ( f'CREATE UNIQUE INDEX "{expected_name}" ON "customer"' f' ("shop_id", "phone_number", "deleted_at") NULLS NOT DISTINCT;' @@ -252,7 +254,9 @@ class Meta: await editor.add_index(Customer, index) assert client.executed - expected_name = editor._generate_index_name("uidx", Customer, ["shop_id", "phone_number", "deleted_at"]) + expected_name = editor._generate_index_name( + "uidx", Customer, ["shop_id", "phone_number", "deleted_at"] + ) assert ( f'CREATE UNIQUE INDEX "{expected_name}" ON "customer"' f' ("shop_id", "phone_number", "deleted_at") NULLS NOT DISTINCT WHERE shop_id = 1;' diff --git a/tests/schema/models_postgres_unique_index.py b/tests/schema/models_postgres_unique_index.py index fca5b599a..d9548913f 100644 --- a/tests/schema/models_postgres_unique_index.py +++ b/tests/schema/models_postgres_unique_index.py @@ -1,5 +1,4 @@ from tortoise import Model, fields - from tortoise.contrib.postgres.indexes import PostgreSQLIndex diff --git a/tortoise/backends/base_postgres/schema_generator.py b/tortoise/backends/base_postgres/schema_generator.py index ce3c44c47..d3900760b 100644 --- a/tortoise/backends/base_postgres/schema_generator.py +++ b/tortoise/backends/base_postgres/schema_generator.py @@ -13,9 +13,7 @@ class BasePostgresSchemaGenerator(BaseSchemaGenerator): DIALECT = "postgres" - INDEX_CREATE_TEMPLATE = ( - 'CREATE INDEX {exists}"{index_name}" ON {table_name} {index_type}({fields}){nulls_not_distinct}{extra};' - ) + INDEX_CREATE_TEMPLATE = 'CREATE INDEX {exists}"{index_name}" ON {table_name} {index_type}({fields}){nulls_not_distinct}{extra};' UNIQUE_INDEX_CREATE_TEMPLATE = INDEX_CREATE_TEMPLATE.replace("INDEX", "UNIQUE INDEX") TABLE_COMMENT_TEMPLATE = "COMMENT ON TABLE {table} IS '{comment}';" COLUMN_COMMENT_TEMPLATE = "COMMENT ON COLUMN {table}.\"{column}\" IS '{comment}';" diff --git a/tortoise/migrations/schema_editor/base_postgres.py b/tortoise/migrations/schema_editor/base_postgres.py index 989226d59..c32d602ee 100644 --- a/tortoise/migrations/schema_editor/base_postgres.py +++ b/tortoise/migrations/schema_editor/base_postgres.py @@ -9,9 +9,7 @@ class BasePostgresSchemaEditor(BaseSchemaEditor): DIALECT = "postgres" - INDEX_CREATE_TEMPLATE = ( - 'CREATE INDEX "{index_name}" ON {table_name} {index_type}({fields}){nulls_not_distinct}{extra};' - ) + INDEX_CREATE_TEMPLATE = 'CREATE INDEX "{index_name}" ON {table_name} {index_type}({fields}){nulls_not_distinct}{extra};' UNIQUE_INDEX_CREATE_TEMPLATE = INDEX_CREATE_TEMPLATE.replace("INDEX", "UNIQUE INDEX") TABLE_COMMENT_TEMPLATE = "COMMENT ON TABLE {table} IS '{comment}';" COLUMN_COMMENT_TEMPLATE = "COMMENT ON COLUMN {table}.\"{column}\" IS '{comment}';"