Multiple Unique Constraints with auto migrations

.gitignore

@@ -18,4 +18,4 @@ htmlcov/
 prof/
 .env/
 .venv/
-result.json
+result.json

piccolo/__init__.py

@@ -1 +1 @@
-__VERSION__ = "0.82.0"
+__VERSION__ = "0.82.0"

piccolo/apps/migrations/auto/diffable_table.py

@@ -132,21 +132,20 @@ def __sub__(self, value: DiffableTable) -> TableDelta:
                 key=lambda x: x.column._meta.name,
             )
         ]
-
         drop_columns = [
             DropColumn(
                 table_class_name=self.class_name,
                 column_name=i.column._meta.name,
                 db_column_name=i.column._meta.db_column_name,
                 tablename=value.tablename,
+                column_class=i.column.__class__
             )
             for i in sorted(
                 {ColumnComparison(column=column) for column in value.columns}
                 - {ColumnComparison(column=column) for column in self.columns},
                 key=lambda x: x.column._meta.name,
             )
         ]
-
         #######################################################################
 
         alter_columns: t.List[AlterColumn] = []

piccolo/apps/migrations/auto/migration_manager.py

@@ -14,10 +14,11 @@
 from piccolo.apps.migrations.auto.serialisation import deserialise_params
 from piccolo.columns import Column, column_types
 from piccolo.columns.column_types import Serial
+from piccolo.columns.constraints import UniqueConstraint
 from piccolo.engine import engine_finder
 from piccolo.table import Table, create_table_class, sort_table_classes
 from piccolo.utils.warnings import colored_warning
-
+from piccolo.query.methods.alter import AddUniqueConstraint, DropConstraint
 
 @dataclass
 class AddColumnClass:
@@ -144,6 +145,8 @@ class MigrationManager:
     alter_columns: AlterColumnCollection = field(
         default_factory=AlterColumnCollection
     )
+    add_unique_constraints: t.List[AddUniqueConstraint] = field(default_factory=list)
+    drop_unique_constraints: t.List[DropConstraint] = field(default_factory=list)
     raw: t.List[t.Union[t.Callable, t.Coroutine]] = field(default_factory=list)
     raw_backwards: t.List[t.Union[t.Callable, t.Coroutine]] = field(
         default_factory=list
@@ -214,32 +217,46 @@ def add_column(
         if column_class is None:
             raise ValueError("Unrecognised column type")
 
-        cleaned_params = deserialise_params(params=params)
-        column = column_class(**cleaned_params)
+        if column_class is UniqueConstraint:
+            column = column_class(**params)
+        else:
+            cleaned_params = deserialise_params(params=params)
+            column = column_class(**cleaned_params)
+
         column._meta.name = column_name
         column._meta.db_column_name = db_column_name
 
-        self.add_columns.append(
-            AddColumnClass(
-                column=column,
-                tablename=tablename,
-                table_class_name=table_class_name,
+        if isinstance(column_class,UniqueConstraint):
+            self.add_unique_constraints.append(
+                AddUniqueConstraint(
+                    constraint_name=column_name,
+                    columns=params.get('unique_columns') #type: ignore
+                )
+            )
+        else:
+            self.add_columns.append(
+                AddColumnClass(
+                    column=column,
+                    tablename=tablename,
+                    table_class_name=table_class_name,
+                )
             )
-        )
 
     def drop_column(
         self,
         table_class_name: str,
         tablename: str,
         column_name: str,
         db_column_name: t.Optional[str] = None,
+        column_class: t.Optional[t.Type[Column]] = None,
     ):
         self.drop_columns.append(
             DropColumn(
                 table_class_name=table_class_name,
                 column_name=column_name,
                 db_column_name=db_column_name or column_name,
                 tablename=tablename,
+                column_class=column_class
             )
         )
 
@@ -569,9 +586,14 @@ async def _run_drop_columns(self, backwards=False):
                 )
 
                 for column in columns:
-                    await _Table.alter().drop_column(
-                        column=column.column_name
-                    ).run()
+                    if column.column_class==UniqueConstraint:
+                        await _Table.alter().drop_constraint(
+                            constraint_name=column.db_column_name
+                        ).run()
+                    else:
+                        await _Table.alter().drop_column(
+                            column=column.column_name
+                        ).run()
 
     async def _run_rename_tables(self, backwards=False):
         for rename_table in self.rename_tables:
@@ -699,11 +721,17 @@ async def _run_add_columns(self, backwards=False):
                     column = _Table._meta.get_column_by_name(
                         add_column.column._meta.name
                     )
-                    await _Table.alter().add_column(
-                        name=column._meta.name, column=column
-                    ).run()
-                    if add_column.column._meta.index:
-                        await _Table.create_index([add_column.column]).run()
+                    if isinstance(add_column.column,UniqueConstraint):
+                        await _Table.alter().add_unique_constraint(
+                            add_column.column._meta.name,
+                            add_column.column.unique_columns
+                        ).run()
+                    else:
+                        await _Table.alter().add_column(
+                            name=column._meta.name, column=column
+                        ).run()
+                        if add_column.column._meta.index:
+                            await _Table.create_index([add_column.column]).run()
 
     async def run(self):
         print(f"  - {self.migration_id} [forwards]... ", end="")

piccolo/apps/migrations/auto/operations.py

@@ -40,6 +40,7 @@ class DropColumn:
     column_name: str
     db_column_name: str
     tablename: str
+    column_class: t.Optional[t.Type[Column]] = None
 
 
 @dataclass

piccolo/apps/migrations/auto/schema_differ.py

@@ -15,6 +15,7 @@
     UniqueGlobalNames,
     serialise_params,
 )
+from piccolo.columns.constraints import UniqueConstraint
 from piccolo.utils.printing import get_fixed_length_string
 
 
@@ -351,6 +352,9 @@ def alter_columns(self) -> AlterStatements:
                     )
 
                 if alter_column.old_column_class is not None:
+                    if alter_column.old_column_class==UniqueConstraint:
+                        print('You cannot ALTER UniqueConstraint! At first, delete it, then create the new one')
+                        continue
                     extra_imports.append(
                         Import(
                             module=alter_column.old_column_class.__module__,
@@ -375,6 +379,7 @@ def alter_columns(self) -> AlterStatements:
     @property
     def drop_columns(self) -> AlterStatements:
         response = []
+        extra_imports: t.List[Import] = []
         for table in self.schema:
             snapshot_table = self._get_snapshot_table(table.class_name)
             if snapshot_table:
@@ -388,11 +393,26 @@ def drop_columns(self) -> AlterStatements:
                     in self.rename_columns_collection.old_column_names
                 ):
                     continue
+                column_class = column.column_class
+                extra_imports.append(
+                    Import(
+                        module=column_class.__module__,
+                        target=column_class.__name__, #type: ignore
+                        expect_conflict_with_global_name=getattr(
+                            UniqueGlobalNames,
+                            f"COLUMN_{column_class.__name__.upper()}", #type: ignore
+                            None,
+                        ),
+                    )
+                )
 
                 response.append(
-                    f"manager.drop_column(table_class_name='{table.class_name}', tablename='{table.tablename}', column_name='{column.column_name}', db_column_name='{column.db_column_name}')"  # noqa: E501
+                    f"manager.drop_column(table_class_name='{table.class_name}', tablename='{table.tablename}', column_name='{column.column_name}', db_column_name='{column.db_column_name}', column_class={column.column_class.__name__})"  # noqa: E501
                 )
-        return AlterStatements(statements=response)
+        return AlterStatements(
+            statements=response,
+            extra_imports=extra_imports
+        )
 
     @property
     def add_columns(self) -> AlterStatements:

piccolo/columns/constraints.py

@@ -0,0 +1,42 @@
+from .base import ColumnMeta, Column
+
+class Constraint(Column):
+    def __init__(self) -> None:
+        pass
+
+class UniqueConstraint(Constraint):
+    """
+    This class represents UNIQUE CONSTRAINT, which can be use to create
+    many complex (multi-field) constraints for the Table
+    All manipulations with the Constraint are like anything about Columns
+
+    Usage:
+        class FooTable(Table):
+            foo_field = Text()
+            bar_field = Text()
+            my_constraint_1 = UniqueConstraint(['foo_field','bar_field'])
+
+    SQL queries for creating and dropping constrains are similar to:
+        ALTER TABLE foo_table ADD CONSTRAINT my_constraint_1 UNIQUE (foo_field, bar_field);
+        ALTER TABLE foo_table DROP IF EXIST CONSTRAINT my_constraint_1;
+    """
+    def __init__(self, unique_columns: list[str]) -> None:
+        super().__init__()
+        self._meta = ColumnMeta()
+        self.unique_columns = unique_columns
+        self._meta.params.update({
+             'unique_columns':self.unique_columns
+        })
+
+    @property
+    def column_type(self):
+        return "CONSTRAINT"
+
+    @property
+    def ddl(self) -> str:
+        """
+        Used when creating tables.
+        """
+        unique_columns_string = ",".join(self.unique_columns)
+        query = f'{self.column_type} "{self._meta.db_column_name}" UNIQUE ({unique_columns_string})'
+        return query

piccolo/query/methods/alter.py

@@ -176,6 +176,17 @@ class SetLength(AlterColumnStatement):
     def ddl(self) -> str:
         return f'ALTER COLUMN "{self.column_name}" TYPE VARCHAR({self.length})'
 
+@dataclass
+class AddUniqueConstraint(AlterStatement):
+    __slots__ = ("constraint_name","columns")
+
+    constraint_name: str
+    columns: list[str]
+
+    @property
+    def ddl(self) -> str:
+        columns_str: str = ",".join(self.columns)
+        return f"ADD CONSTRAINT {self.constraint_name} UNIQUE ({columns_str})"
 
 @dataclass
 class DropConstraint(AlterStatement):
@@ -264,6 +275,7 @@ class Alter(DDL):
     __slots__ = (
         "_add_foreign_key_constraint",
         "_add",
+        "_add_unique_constraint",
         "_drop_constraint",
         "_drop_default",
         "_drop_table",
@@ -282,6 +294,7 @@ def __init__(self, table: t.Type[Table], **kwargs):
         super().__init__(table, **kwargs)
         self._add_foreign_key_constraint: t.List[AddForeignKeyConstraint] = []
         self._add: t.List[AddColumn] = []
+        self._add_unique_constraint: t.List[AddUniqueConstraint] = []
         self._drop_constraint: t.List[DropConstraint] = []
         self._drop_default: t.List[DropDefault] = []
         self._drop_table: t.Optional[DropTable] = None
@@ -433,6 +446,12 @@ def _get_constraint_name(self, column: t.Union[str, ForeignKey]) -> str:
         tablename = self.table._meta.tablename
         return f"{tablename}_{column_name}_fk"
 
+    def add_unique_constraint(self, constraint_name: str, columns: list[str]):
+        self._add_unique_constraint.append(
+            AddUniqueConstraint(constraint_name=constraint_name,columns=columns)
+        )
+        return self
+
     def drop_constraint(self, constraint_name: str) -> Alter:
         self._drop_constraint.append(
             DropConstraint(constraint_name=constraint_name)
@@ -520,6 +539,8 @@ def default_ddl(self) -> t.Sequence[str]:
                 self._set_length,
                 self._set_default,
                 self._set_digits,
+                self._add_unique_constraint,
+                self._drop_constraint,
             )
         ]
 

piccolo/table.py

@@ -7,6 +7,7 @@
 from dataclasses import dataclass, field
 
 from piccolo.columns import Column
+from piccolo.columns.constraints import UniqueConstraint
 from piccolo.columns.column_types import (
     JSON,
     JSONB,
@@ -73,6 +74,7 @@ class TableMeta:
     primary_key: Column = field(default_factory=Column)
     json_columns: t.List[t.Union[JSON, JSONB]] = field(default_factory=list)
     secret_columns: t.List[Secret] = field(default_factory=list)
+    unique_constraints: t.List[UniqueConstraint] = field(default_factory=list)
     tags: t.List[str] = field(default_factory=list)
     help_text: t.Optional[str] = None
     _db: t.Optional[Engine] = None
@@ -218,6 +220,7 @@ def __init_subclass__(
         json_columns: t.List[t.Union[JSON, JSONB]] = []
         primary_key: t.Optional[Column] = None
         m2m_relationships: t.List[M2M] = []
+        unique_constraints: t.List[UniqueConstraint] = []
 
         attribute_names = itertools.chain(
             *[i.__dict__.keys() for i in reversed(cls.__mro__)]
@@ -258,6 +261,9 @@ def __init_subclass__(
                 if isinstance(column, (JSON, JSONB)):
                     json_columns.append(column)
 
+            if isinstance(attribute, UniqueConstraint):
+                unique_constraints.append(attribute)
+
             if isinstance(attribute, M2M):
                 attribute._meta._name = attribute_name
                 attribute._meta._table = cls
@@ -279,6 +285,7 @@ def __init_subclass__(
             foreign_key_columns=foreign_key_columns,
             json_columns=json_columns,
             secret_columns=secret_columns,
+            unique_constraints=unique_constraints,
             tags=tags,
             help_text=help_text,
             _db=db,

piccolo_conf.py

@@ -1,11 +1,8 @@
 """
 This piccolo_conf file is just here so migrations can be made for Piccolo's own
 internal apps.
-
 For example:
-
 python -m piccolo.main migration new user --auto
-
 """
 
 from piccolo.conf.apps import AppRegistry
@@ -17,4 +14,4 @@
 
 # A list of paths to piccolo apps
 # e.g. ['blog.piccolo_app']
-APP_REGISTRY = AppRegistry(apps=["piccolo.apps.user.piccolo_app"])
+APP_REGISTRY = AppRegistry(apps=["piccolo.apps.user.piccolo_app"])