Merge b3793e5 into e023166

CHANGELOG.rst

@@ -10,6 +10,9 @@ Changelog
 ====
 0.19.1
 ------
+Added
+^^^^^
+- Added Postgress partial indexes support. (#1103)
 Fixed
 ^^^^^
 - `TimeField` for `MySQL` will return `datetime.timedelta` object instead of `datetime.time` object.

tests/schema/models_postgres_index.py

@@ -6,6 +6,7 @@
     GinIndex,
     GistIndex,
     HashIndex,
+    PostgreSQLIndex,
     SpGistIndex,
 )
 
@@ -17,6 +18,7 @@ class Index(Model):
     gist = TSVectorField()
     sp_gist = fields.CharField(max_length=200)
     hash = fields.CharField(max_length=200)
+    partial = fields.CharField(max_length=200)
 
     class Meta:
         indexes = [
@@ -26,4 +28,5 @@ class Meta:
             GistIndex(fields={"gist"}),
             SpGistIndex(fields={"sp_gist"}),
             HashIndex(fields={"hash"}),
+            PostgreSQLIndex(fields={"partial"}, condition={"id": 1}),
         ]

tests/schema/test_generate_schema.py

@@ -1061,19 +1061,22 @@ async def test_index_unsafe(self):
     "gin" TSVECTOR NOT NULL,
     "gist" TSVECTOR NOT NULL,
     "sp_gist" VARCHAR(200) NOT NULL,
-    "hash" VARCHAR(200) NOT NULL
+    "hash" VARCHAR(200) NOT NULL,
+    "partial" VARCHAR(200) NOT NULL
 );
 CREATE INDEX "idx_index_bloom_280137" ON "index" USING BLOOM ("bloom");
 CREATE INDEX "idx_index_brin_a54a00" ON "index" USING BRIN ("brin");
 CREATE INDEX "idx_index_gin_a403ee" ON "index" USING GIN ("gin");
 CREATE INDEX "idx_index_gist_c807bf" ON "index" USING GIST ("gist");
 CREATE INDEX "idx_index_sp_gist_2c0bad" ON "index" USING SPGIST ("sp_gist");
-CREATE INDEX "idx_index_hash_cfe6b5" ON "index" USING HASH ("hash");""",
+CREATE INDEX "idx_index_hash_cfe6b5" ON "index" USING HASH ("hash");
+CREATE INDEX "idx_index_partial_c5be6a" ON "index" USING  ("partial") WHERE id = 1;""",
         )
 
     async def test_index_safe(self):
         await self.init_for("tests.schema.models_postgres_index")
         sql = get_schema_sql(connections.get("default"), safe=True)
+        print(sql)
         self.assertEqual(
             sql,
             """CREATE TABLE IF NOT EXISTS "index" (
@@ -1083,14 +1086,16 @@ async def test_index_safe(self):
     "gin" TSVECTOR NOT NULL,
     "gist" TSVECTOR NOT NULL,
     "sp_gist" VARCHAR(200) NOT NULL,
-    "hash" VARCHAR(200) NOT NULL
+    "hash" VARCHAR(200) NOT NULL,
+    "partial" VARCHAR(200) NOT NULL
 );
 CREATE INDEX IF NOT EXISTS "idx_index_bloom_280137" ON "index" USING BLOOM ("bloom");
 CREATE INDEX IF NOT EXISTS "idx_index_brin_a54a00" ON "index" USING BRIN ("brin");
 CREATE INDEX IF NOT EXISTS "idx_index_gin_a403ee" ON "index" USING GIN ("gin");
 CREATE INDEX IF NOT EXISTS "idx_index_gist_c807bf" ON "index" USING GIST ("gist");
 CREATE INDEX IF NOT EXISTS "idx_index_sp_gist_2c0bad" ON "index" USING SPGIST ("sp_gist");
-CREATE INDEX IF NOT EXISTS "idx_index_hash_cfe6b5" ON "index" USING HASH ("hash");""",
+CREATE INDEX IF NOT EXISTS "idx_index_hash_cfe6b5" ON "index" USING HASH ("hash");
+CREATE INDEX IF NOT EXISTS "idx_index_partial_c5be6a" ON "index" USING  ("partial") WHERE id = 1;""",
         )
 
     async def test_m2m_no_auto_create(self):

tortoise/contrib/postgres/indexes.py

@@ -1,13 +1,31 @@
-from abc import ABCMeta
+from typing import Optional, Set
+
+from pypika.terms import Term, ValueWrapper
 
 from tortoise.indexes import Index
 
 
-class PostgreSQLIndex(Index, metaclass=ABCMeta):
+class PostgreSQLIndex(Index):
     INDEX_CREATE_TEMPLATE = (
-        "CREATE INDEX {exists}{index_name} ON {table_name} USING{index_type}({fields});"
+        "CREATE INDEX {exists}{index_name} ON {table_name} USING{index_type}({fields}){extra};"
     )
 
+    def __init__(
+        self,
+        *expressions: Term,
+        fields: Optional[Set[str]] = None,
+        name: Optional[str] = None,
+        condition: Optional[dict] = None,
+    ):
+        super().__init__(*expressions, fields=fields, name=name)
+        if condition:
+            cond = " WHERE "
+            items = []
+            for k, v in condition.items():
+                items.append(f"{k} = {ValueWrapper(v)}")
+            cond += " AND ".join(items)
+            self.extra = cond
+
 
 class BloomIndex(PostgreSQLIndex):
     INDEX_TYPE = "BLOOM"