diff --git a/piccolo/constraints.py b/piccolo/constraints.py index 73e5f750c..ee45909aa 100644 --- a/piccolo/constraints.py +++ b/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 `_ - 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,6 +127,7 @@ 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.") @@ -98,9 +135,17 @@ def __init__( 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})" + ) diff --git a/piccolo/table.py b/piccolo/table.py index 4f72022a0..048f71a7e 100644 --- a/piccolo/table.py +++ b/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) diff --git a/tests/apps/migrations/auto/test_schema_differ.py b/tests/apps/migrations/auto/test_schema_differ.py index d840ac341..ce048d314 100644 --- a/tests/apps/migrations/auto/test_schema_differ.py +++ b/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 diff --git a/tests/table/test_constraints.py b/tests/table/test_constraints.py new file mode 100644 index 000000000..35de89d00 --- /dev/null +++ b/tests/table/test_constraints.py @@ -0,0 +1,216 @@ +from unittest import TestCase + +from piccolo.columns import Integer, Varchar +from piccolo.constraints import Check, Unique +from piccolo.table import Table + + +class TestUniqueConstraint(TestCase): + def test_default_name(self): + class Album(Table): + name = Varchar() + code = Integer() + unique_name_code = Unique(columns=[name, code]) + + self.assertEqual( + Album.unique_name_code._meta.name, "unique_name_code" + ) + self.assertEqual( + Album._meta.constraints[0]._meta.name, "unique_name_code" + ) + self.assertEqual( + Album.unique_name_code._meta.params, + {"columns": ["name", "code"], "nulls_distinct": True}, + ) + self.assertNotIn( + "prefix_tablename", Album.unique_name_code._meta.params + ) + + def test_explicit_name(self): + class Album(Table): + name = Varchar() + code = Integer() + unique_name_code = Unique( + columns=[name, code], name="custom_album_unique" + ) + + self.assertEqual( + Album.unique_name_code._meta.name, "custom_album_unique" + ) + self.assertEqual( + Album._meta.constraints[0]._meta.name, "custom_album_unique" + ) + + def test_prefix_tablename(self): + class Album(Table): + name = Varchar() + code = Integer() + unique_name_code = Unique( + columns=[name, code], prefix_tablename=True + ) + + self.assertEqual( + Album.unique_name_code._meta.name, "album_unique_name_code" + ) + self.assertEqual( + Album._meta.constraints[0]._meta.name, "album_unique_name_code" + ) + self.assertEqual( + Album.unique_name_code._meta.params, + { + "columns": ["name", "code"], + "nulls_distinct": True, + "prefix_tablename": True, + }, + ) + + def test_prefix_tablename_with_explicit_name(self): + class Album(Table): + name = Varchar() + code = Integer() + unique_name_code = Unique( + columns=[name, code], + name="custom_idx", + prefix_tablename=True, + ) + + self.assertEqual(Album.unique_name_code._meta.name, "album_custom_idx") + self.assertEqual( + Album._meta.constraints[0]._meta.name, "album_custom_idx" + ) + + def test_mixin_inheritance_with_prefix(self): + class IdentifierMixin: + external_id = Integer() + version_no = Integer() + + unique_identifier = Unique( + columns=[external_id, version_no], + prefix_tablename=True, + ) + + class Order(IdentifierMixin, Table, tablename="orders"): + title = Varchar() + + class Invoice(IdentifierMixin, Table, tablename="invoices"): + title = Varchar() + + # Both tables have distinct constraint names prefixed by tablename: + self.assertEqual( + Order.unique_identifier._meta.name, "orders_unique_identifier" + ) + self.assertEqual( + Invoice.unique_identifier._meta.name, "invoices_unique_identifier" + ) + + # Both tables have distinct Constraint instances: + self.assertIsNot(Order.unique_identifier, Invoice.unique_identifier) + self.assertIsNot( + Order.unique_identifier, IdentifierMixin.unique_identifier + ) + + # Columns in constraint point to respective table's cloned columns: + self.assertIs(Order.unique_identifier.columns[0], Order.external_id) + self.assertIs( + Invoice.unique_identifier.columns[0], Invoice.external_id + ) + + # DDL includes properly prefixed constraint names: + order_ddl = [str(x) for x in Order.create_table().ddl] + expected_order_sql = ( + 'ADD CONSTRAINT orders_unique_identifier ' + 'UNIQUE ("external_id", "version_no")' + ) + self.assertTrue( + any(expected_order_sql in stmt for stmt in order_ddl) + ) + + invoice_ddl = [str(x) for x in Invoice.create_table().ddl] + expected_invoice_sql = ( + 'ADD CONSTRAINT invoices_unique_identifier ' + 'UNIQUE ("external_id", "version_no")' + ) + self.assertTrue( + any(expected_invoice_sql in stmt for stmt in invoice_ddl) + ) + + def test_mixin_inheritance_without_prefix(self): + class IdentifierMixin: + external_id = Integer() + version_no = Integer() + + unique_identifier = Unique( + columns=[external_id, version_no], + prefix_tablename=False, + ) + + class Order(IdentifierMixin, Table, tablename="orders"): + title = Varchar() + + class Invoice(IdentifierMixin, Table, tablename="invoices"): + title = Varchar() + + # Without prefix, both keep the base attribute name: + self.assertEqual( + Order.unique_identifier._meta.name, "unique_identifier" + ) + self.assertEqual( + Invoice.unique_identifier._meta.name, "unique_identifier" + ) + # Instances remain isolated: + self.assertIsNot(Order.unique_identifier, Invoice.unique_identifier) + + def test_table_str(self): + class Album(Table): + name = Varchar() + unique_name = Unique(columns=[name], prefix_tablename=True) + + self.assertIn("prefix_tablename=True", Album.unique_name._table_str()) + + +class TestCheckConstraint(TestCase): + def test_check_default_name(self): + class Ticket(Table): + price = Integer() + valid_price = Check(price >= 0) + + self.assertEqual(Ticket.valid_price._meta.name, "valid_price") + self.assertNotIn("prefix_tablename", Ticket.valid_price._meta.params) + + def test_check_prefix_tablename(self): + class Ticket(Table): + price = Integer() + valid_price = Check(price >= 0, prefix_tablename=True) + + self.assertEqual(Ticket.valid_price._meta.name, "ticket_valid_price") + self.assertTrue( + Ticket.valid_price._meta.params.get("prefix_tablename") + ) + + def test_check_mixin_inheritance(self): + class PriceMixin: + price = Integer() + valid_price = Check(price >= 0, prefix_tablename=True) + + class ConcertTicket(PriceMixin, Table, tablename="concert_tickets"): + pass + + class CinemaTicket(PriceMixin, Table, tablename="cinema_tickets"): + pass + + self.assertEqual( + ConcertTicket.valid_price._meta.name, "concert_tickets_valid_price" + ) + self.assertEqual( + CinemaTicket.valid_price._meta.name, "cinema_tickets_valid_price" + ) + self.assertIsNot( + ConcertTicket.valid_price, CinemaTicket.valid_price + ) + + def test_check_table_str(self): + class Ticket(Table): + price = Integer() + valid_price = Check(price >= 0, prefix_tablename=True) + + self.assertIn("prefix_tablename=True", Ticket.valid_price._table_str())