From ec2f396aca724632048e09a21e58d639902ca760 Mon Sep 17 00:00:00 2001 From: pctablet505 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 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 a17084a86948b6f362e5357a6c3f450ecfbb328c Mon Sep 17 00:00:00 2001 From: pctablet505 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 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}';" From 660a1e7e6b024bf7676684edd02bec083c45658f Mon Sep 17 00:00:00 2001 From: pctablet505 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 299011ef3..0b144dc65 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 c32d602ee..7ff9fa3ea 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]),