feat(indexes): add nulls_not_distinct option to PostgreSQLIndex

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,77 @@ 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):

tests/schema/models_postgres_unique_index.py

@@ -0,0 +1,17 @@
+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,
+            ),
+        ]

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()

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}){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}';"
@@ -79,10 +77,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="",
         )

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):

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)

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}){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}';"
@@ -78,16 +76,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 +108,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="",
         )