squash commit

tortoise/contrib/postgres/indexes.py

@@ -1,8 +1,10 @@
-from typing import Optional, Tuple
+from typing import Optional, Tuple, Type
 
 from pypika.terms import Term, ValueWrapper
 
+from tortoise.backends.base.schema_generator import BaseSchemaGenerator
 from tortoise.indexes import PartialIndex
+from tortoise.models import Model
 
 
 class PostgreSQLIndex(PartialIndex):
@@ -49,3 +51,26 @@ class HashIndex(PostgreSQLIndex):
 
 class SpGistIndex(PostgreSQLIndex):
     INDEX_TYPE = "SPGIST"
+
+
+class PostgresUniqueIndex(PostgreSQLIndex):
+    INDEX_CREATE_TEMPLATE = PostgreSQLIndex.INDEX_CREATE_TEMPLATE.replace(
+        "CREATE", "CREATE UNIQUE"
+    ).replace("USING", "")
+
+    def __init__(
+        self,
+        *expressions: Term,
+        fields: Optional[Tuple[str]] = None,
+        name: Optional[str] = None,
+        condition: Optional[dict] = None,
+        nulls_not_distinct: bool = False,
+    ):
+        super().__init__(*expressions, fields=fields, name=name, condition=condition)
+        if nulls_not_distinct:
+            self.extra = " nulls not distinct".upper() + self.extra
+
+    def get_sql(self, schema_generator: BaseSchemaGenerator, model: Type[Model], safe: bool):
+        if self.INDEX_TYPE:
+            self.INDEX_TYPE = f"USING {self.INDEX_TYPE}"
+        return super().get_sql(schema_generator, model, safe)

tortoise/contrib/sqlite/indexes.py

@@ -0,0 +1,46 @@
+from typing import Dict, Optional, Tuple
+
+from pypika.terms import Term, ValueWrapper
+
+from tortoise.indexes import PartialIndex
+
+
+class SqliteIndex(PartialIndex):
+    INDEX_CREATE_TEMPLATE = PartialIndex.INDEX_CREATE_TEMPLATE.replace("INDEX", "INDEX {exists} ")
+
+
+class SqliteUniqueIndex(SqliteIndex):
+    INDEX_TYPE = "unique".upper()
+
+    def __init__(
+        self,
+        *expressions: Term,
+        fields: Optional[Tuple[str, ...]] = None,
+        name: Optional[str] = None,
+        condition: Optional[Dict[str, str]] = None,
+        where_expre: Optional[str] = None,
+    ):
+        _condition = condition
+        if condition:
+            condition = None
+        super().__init__(*expressions, fields=fields, name=name, condition=condition)
+        if _condition:
+            # TODO: what is the best practice to inject where expression?
+            self.extra = where_expre or self._gen_condition(_condition)
+
+    @classmethod
+    def _gen_field_cond(cls, kv: tuple):
+        key, cond = kv
+        op = (
+            ""
+            if isinstance(cond, str)
+            and cond.strip().lower().startswith(("is ", "isnot ", "<", ">", "=", "!="))
+            else "="
+        )
+        if op == "=":
+            cond = ValueWrapper(cond)
+        return str(f"{key} {op} {cond}")
+
+    def _gen_condition(self, conditions: Dict[str, str]):
+        conditions = " AND ".join(tuple(map(self._gen_field_cond, conditions.items())))
+        return f" where {conditions}"