From 88002864fc54b7d10871b9fb86007f72d98fe954 Mon Sep 17 00:00:00 2001 From: pctablet505 <55033230+pctablet505@users.noreply.github.com> Date: Wed, 15 Jul 2026 15:15:57 +0000 Subject: [PATCH 1/3] 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 55df49b7a..5beb85146 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 701d48dd5..ff4afa167 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 c154f1aa9fe61684051f93637449ff3f5ae204fb Mon Sep 17 00:00:00 2001 From: pctablet505 <55033230+pctablet505@users.noreply.github.com> Date: Thu, 16 Jul 2026 11:20:16 +0530 Subject: [PATCH 2/3] 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 ff4afa167..5fcc78e5d 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}';" From ad8553b42050414601fe0ca47a07ed4812c4fbbb Mon Sep 17 00:00:00 2001 From: pctablet505 <55033230+pctablet505@users.noreply.github.com> Date: Sun, 16 Aug 2026 08:12:45 +0530 Subject: [PATCH 3/3] fix: accept unique/nulls_not_distinct kwargs in non-Postgres _get_index_sql MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit add_index() unconditionally passes unique= and nulls_not_distinct= to self._get_index_sql(), but only the Postgres schema editor's override declared those parameters. Every other backend (sqlite/mssql via the base implementation, and mysql via its own override) raised TypeError: _get_index_sql() got an unexpected keyword argument 'unique' on any migration that adds an index — this broke test_add_index_operation_runs_sql on the sqlite/mysql/mssql CI matrix. Both parameters are now accepted on every backend. unique mirrors the existing UNIQUE_INDEX_CREATE_TEMPLATE/uidx-prefix pattern already used elsewhere in these classes; nulls_not_distinct is inherently PostgreSQL-only (NULLS NOT DISTINCT has no equivalent elsewhere) so it's accepted and ignored on other backends, same as unique already is today in practice (only PostgreSQLIndex sets these attributes). Also fixes the mypy error in BasePostgresSchemaEditor._get_index_sql: _generate_index_name expects list[str], not the Sequence[str] the field_names parameter is typed as. --- tortoise/migrations/schema_editor/base.py | 11 +++++++++-- tortoise/migrations/schema_editor/base_postgres.py | 2 +- tortoise/migrations/schema_editor/mysql.py | 7 ++++++- 3 files changed, 16 insertions(+), 4 deletions(-) diff --git a/tortoise/migrations/schema_editor/base.py b/tortoise/migrations/schema_editor/base.py index 5beb85146..87386273f 100644 --- a/tortoise/migrations/schema_editor/base.py +++ b/tortoise/migrations/schema_editor/base.py @@ -201,9 +201,16 @@ 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: - return self.INDEX_CREATE_TEMPLATE.format( - index_name=index_name or self._generate_index_name("idx", model, field_names), + # nulls_not_distinct is PostgreSQL-only (see BasePostgresSchemaEditor); + # accepted here so callers can pass it unconditionally, but it's a no-op. + _ = nulls_not_distinct + 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), fields=self._format_index_fields(field_names), index_type=f"{index_type} " if index_type else "", diff --git a/tortoise/migrations/schema_editor/base_postgres.py b/tortoise/migrations/schema_editor/base_postgres.py index 5fcc78e5d..0469efd6f 100644 --- a/tortoise/migrations/schema_editor/base_postgres.py +++ b/tortoise/migrations/schema_editor/base_postgres.py @@ -87,7 +87,7 @@ def _get_index_sql( prefix = "uidx" if unique else "idx" return template.format( - index_name=index_name or self._generate_index_name(prefix, model, field_names), + index_name=index_name or self._generate_index_name(prefix, model, list(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)), diff --git a/tortoise/migrations/schema_editor/mysql.py b/tortoise/migrations/schema_editor/mysql.py index fd52ad940..8619f2f02 100644 --- a/tortoise/migrations/schema_editor/mysql.py +++ b/tortoise/migrations/schema_editor/mysql.py @@ -92,8 +92,13 @@ 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: - _ = safe + # unique/nulls_not_distinct are only set by the PostgreSQL-specific + # PostgreSQLIndex; accepted here so callers can pass them unconditionally, + # but they're a no-op on MySQL (nulls_not_distinct has no MySQL equivalent). + _ = safe, unique, nulls_not_distinct index_sql = self.INDEX_CREATE_TEMPLATE.format( index_name=index_name or self._generate_index_name("idx", model, field_names), fields=", ".join([self.quote(f) for f in field_names]),