feat(constraints): add prefix_tablename option to Unique and Check constraints (#1437)

piccolo/constraints.py

@@ -1,13 +1,16 @@
 from __future__ import annotations
 
 from abc import ABCMeta, abstractmethod
+import copy
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Any, Union
+from typing import TYPE_CHECKING, Any, TypeVar, Union
 
 if TYPE_CHECKING:
     from piccolo.columns import Column
     from piccolo.custom_types import Combinable
 
+Self = TypeVar("Self", bound="Constraint")
+
 
 @dataclass
 class ConstraintMeta:
@@ -16,6 +19,7 @@ class ConstraintMeta:
     """
 
     _name: str | None
+    _table: Any = None
 
     # Used for representing the table in migrations.
     params: dict[str, Any] = field(default_factory=dict)
@@ -33,15 +37,47 @@ def name(self) -> str:
     def name(self, value: str):
         self._name = value
 
+    @property
+    def table(self) -> Any:
+        return self._table
+
+    @table.setter
+    def table(self, value: Any):
+        self._table = value
+
+    def copy(self) -> ConstraintMeta:
+        return ConstraintMeta(
+            _name=self._name,
+            _table=self._table,
+            params=self.params.copy(),
+        )
+
 
 class Constraint(metaclass=ABCMeta):
     """
     All other constraints inherit from ``Constraint``. Don't use it directly.
     """
 
-    def __init__(self, name: str | None = None, **kwargs) -> None:
+    def __init__(
+        self,
+        name: str | None = None,
+        prefix_tablename: bool = False,
+        **kwargs,
+    ) -> None:
+        self.prefix_tablename = prefix_tablename
+        self._explicit_name = name
+        if prefix_tablename:
+            kwargs["prefix_tablename"] = True
         self._meta = ConstraintMeta(_name=name, params=kwargs)
 
+    def copy(self: Self) -> Self:
+        constraint = copy.copy(self)
+        constraint._meta = self._meta.copy()
+        return constraint
+
+    def __deepcopy__(self, memo) -> Self:
+        return self.copy()
+
     def __hash__(self):
         return hash(self._meta._name)
 
@@ -80,9 +116,9 @@ class Album(Table):
 
     :param columns:
         The table columns that should be unique together.
-    :param nulls_distinct:
-        See the `Postgres docs <https://www.postgresql.org/docs/current/ddl-constraints.html#DDL-CONSTRAINTS-UNIQUE-CONSTRAINTS>`_
-        for more information.
+    :param prefix_tablename:
+        Whether to prefix the constraint's index name with the table name.
+        Particularly useful when unique constraints are inherited via mixins.
 
     """  # noqa: E501
 
@@ -91,16 +127,25 @@ def __init__(
         columns: list[Union[Column, str]],
         nulls_distinct: bool = True,
         name: str | None = None,
+        prefix_tablename: bool = False,
     ):
         if len(columns) < 1:
             raise ValueError("At least 1 column must be specified.")
 
         self.columns = columns
         self.nulls_distinct = nulls_distinct
         super().__init__(
-            name=name, columns=columns, nulls_distinct=nulls_distinct
+            name=name,
+            prefix_tablename=prefix_tablename,
+            columns=columns,
+            nulls_distinct=nulls_distinct,
         )
 
+    def copy(self: Self) -> Self:
+        constraint = super().copy()
+        constraint.columns = list(self.columns)
+        return constraint
+
     def get_column_names(self):
         from piccolo.columns import Column
 
@@ -122,6 +167,8 @@ def serialise_self(self):
         the columns before calling this method.
         """
         self._meta.params["columns"] = self.get_column_names()
+        if self.prefix_tablename:
+            self._meta.params["prefix_tablename"] = True
 
     @property
     def ddl(self) -> str:
@@ -138,9 +185,12 @@ def _table_str(self) -> str:
         columns_string = ", ".join(
             [f'"{column_name}"' for column_name in column_names]
         )
+        prefix_str = (
+            ", prefix_tablename=True" if self.prefix_tablename else ""
+        )
         return (
             f"{self._meta._name} = Unique([{columns_string}], "
-            f"nulls_distinct={self.nulls_distinct})"
+            f"nulls_distinct={self.nulls_distinct}{prefix_str})"
         )
 
 
@@ -179,17 +229,22 @@ def __init__(
         self,
         condition: Union[Combinable, str],
         name: str | None = None,
+        prefix_tablename: bool = False,
     ):
         """
         :param condition:
             The SQL expression used to make sure the data is valid (e.g.
             ``"price > 0"``).
-        :param name:
-            The name of the constraint in the database.
+        :param prefix_tablename:
+            Whether to prefix the constraint's name with the table name.
 
         """
         self.condition = condition
-        super().__init__(name=name, condition=condition)
+        super().__init__(
+            name=name,
+            prefix_tablename=prefix_tablename,
+            condition=condition,
+        )
 
     def get_condition_str(self) -> str:
         from piccolo.columns.combination import CombinableMixin
@@ -205,10 +260,18 @@ def serialise_self(self):
         the columns before calling this method.
         """
         self._meta.params["condition"] = self.get_condition_str()
+        if self.prefix_tablename:
+            self._meta.params["prefix_tablename"] = True
 
     @property
     def ddl(self) -> str:
         return f"CHECK ({self.get_condition_str()})"
 
     def _table_str(self) -> str:
-        return f"{self._meta._name} = Check('{self.get_condition_str()}')"
+        prefix_str = (
+            ", prefix_tablename=True" if self.prefix_tablename else ""
+        )
+        return (
+            f"{self._meta._name} = "
+            f"Check('{self.get_condition_str()}'{prefix_str})"
+        )

piccolo/table.py

@@ -366,8 +366,18 @@ def __init_subclass__(
                 m2m_relationships.append(attribute)
 
             if isinstance(attribute, Constraint):
-                attribute._meta._name = attribute_name
-                constaints.append(attribute)
+                constraint = attribute.copy()
+                setattr(cls, attribute_name, constraint)
+                constraint._meta._table = cls
+
+                explicit_name = getattr(constraint, "_explicit_name", None)
+                base_name = explicit_name or attribute_name
+                if getattr(constraint, "prefix_tablename", False):
+                    constraint._meta._name = f"{tablename}_{base_name}"
+                else:
+                    constraint._meta._name = base_name
+
+                constaints.append(constraint)
 
         if not primary_key:
             primary_key = cls._create_serial_primary_key()
@@ -409,6 +419,15 @@ def __init_subclass__(
 
         # Now the table and columns are all setup, serialise the constraints.
         for constraint in cls._meta.constraints:
+            if hasattr(constraint, "columns"):
+                constraint.columns = [
+                    getattr(cls, col._meta.name)
+                    if isinstance(col, Column)
+                    and col._meta._name
+                    and hasattr(cls, col._meta.name)
+                    else col
+                    for col in constraint.columns
+                ]
             constraint.serialise_self()
 
         TABLE_REGISTRY.append(cls)

tests/apps/migrations/auto/test_schema_differ.py

@@ -728,6 +728,45 @@ def test_drop_constraint(self) -> None:
             "manager.drop_constraint(table_class_name='Band', tablename='band', constraint_name='unique_name_genre', schema=None)",  # noqa: E501
         )
 
+    def test_add_constraint_with_prefix_tablename(self) -> None:
+        """
+        Test adding a constraint with prefix_tablename=True.
+        """
+        name_column = Varchar()
+        name_column._meta.name = "name"
+
+        name_unique_constraint = Unique(
+            columns=["name"],
+            name="band_unique_name",
+            prefix_tablename=True,
+        )
+
+        schema: list[DiffableTable] = [
+            DiffableTable(
+                class_name="Band",
+                tablename="band",
+                columns=[name_column],
+                constraints=[name_unique_constraint],
+            )
+        ]
+        schema_snapshot: list[DiffableTable] = [
+            DiffableTable(
+                class_name="Band",
+                tablename="band",
+                columns=[name_column],
+            )
+        ]
+
+        schema_differ = SchemaDiffer(
+            schema=schema, schema_snapshot=schema_snapshot, auto_input="y"
+        )
+
+        self.assertEqual(len(schema_differ.add_constraints.statements), 1)
+        self.assertEqual(
+            schema_differ.add_constraints.statements[0],
+            "manager.add_constraint(table_class_name='Band', tablename='band', constraint_name='band_unique_name', constraint_class=Unique, params={'columns': ['name'], 'nulls_distinct': True, 'prefix_tablename': True}, schema=None)",  # noqa: E501
+        )
+
     def test_alter_default(self):
         pass