From ca4c4a2a9781e1aa19e2fbbb965838f689a1eed4 Mon Sep 17 00:00:00 2001 From: Keita Ichihashi Date: Mon, 8 Apr 2024 23:43:51 +0900 Subject: [PATCH 01/29] Add support for composite unique constraint --- .../apps/migrations/auto/diffable_table.py | 73 ++++- .../apps/migrations/auto/migration_manager.py | 213 +++++++++++++- piccolo/apps/migrations/auto/operations.py | 19 ++ piccolo/apps/migrations/auto/schema_differ.py | 108 +++++++ .../apps/migrations/auto/schema_snapshot.py | 19 ++ piccolo/apps/migrations/commands/new.py | 1 + piccolo/constraint.py | 55 ++++ piccolo/query/methods/alter.py | 20 ++ piccolo/query/methods/create.py | 6 + piccolo/table.py | 17 ++ .../migrations/auto/test_migration_manager.py | 277 ++++++++++++++++++ .../migrations/auto/test_schema_differ.py | 87 ++++++ 12 files changed, 888 insertions(+), 7 deletions(-) create mode 100644 piccolo/constraint.py diff --git a/piccolo/apps/migrations/auto/diffable_table.py b/piccolo/apps/migrations/auto/diffable_table.py index aa609f041..d6172fce5 100644 --- a/piccolo/apps/migrations/auto/diffable_table.py +++ b/piccolo/apps/migrations/auto/diffable_table.py @@ -5,14 +5,17 @@ from piccolo.apps.migrations.auto.operations import ( AddColumn, + AddConstraint, AlterColumn, DropColumn, + DropConstraint, ) from piccolo.apps.migrations.auto.serialisation import ( deserialise_params, serialise_params, ) from piccolo.columns.base import Column +from piccolo.constraint import Constraint from piccolo.table import Table, create_table_class @@ -55,6 +58,8 @@ class TableDelta: add_columns: t.List[AddColumn] = field(default_factory=list) drop_columns: t.List[DropColumn] = field(default_factory=list) alter_columns: t.List[AlterColumn] = field(default_factory=list) + add_constraints: t.List[AddConstraint] = field(default_factory=list) + drop_constraints: t.List[DropConstraint] = field(default_factory=list) def __eq__(self, value: TableDelta) -> bool: # type: ignore """ @@ -85,6 +90,19 @@ def __eq__(self, value) -> bool: return False +@dataclass +class ConstraintComparison: + constraint: Constraint + + def __hash__(self) -> int: + return self.constraint.__hash__() + + def __eq__(self, value) -> bool: + if isinstance(value, ConstraintComparison): + return self.constraint._meta.name == value.constraint._meta.name + return False + + @dataclass class DiffableTable: """ @@ -96,6 +114,7 @@ class DiffableTable: tablename: str schema: t.Optional[str] = None columns: t.List[Column] = field(default_factory=list) + constraints: t.List[Constraint] = field(default_factory=list) previous_class_name: t.Optional[str] = None def __post_init__(self) -> None: @@ -189,10 +208,54 @@ def __sub__(self, value: DiffableTable) -> TableDelta: ) ) + add_constraints = [ + AddConstraint( + table_class_name=self.class_name, + constraint_name=i.constraint._meta.name, + constraint_class_name=i.constraint.__class__.__name__, + constraint_class=i.constraint.__class__, + params=i.constraint._meta.params, + schema=self.schema, + ) + for i in sorted( + { + ConstraintComparison(constraint=constraint) + for constraint in self.constraints + } + - { + ConstraintComparison(constraint=constraint) + for constraint in value.constraints + }, + key=lambda x: x.constraint._meta.name, + ) + ] + + drop_constraints = [ + DropConstraint( + table_class_name=self.class_name, + constraint_name=i.constraint._meta.name, + tablename=value.tablename, + schema=self.schema, + ) + for i in sorted( + { + ConstraintComparison(constraint=constraint) + for constraint in value.constraints + } + - { + ConstraintComparison(constraint=constraint) + for constraint in self.constraints + }, + key=lambda x: x.constraint._meta.name, + ) + ] + return TableDelta( add_columns=add_columns, drop_columns=drop_columns, alter_columns=alter_columns, + add_constraints=add_constraints, + drop_constraints=drop_constraints, ) def __hash__(self) -> int: @@ -218,10 +281,14 @@ def to_table_class(self) -> t.Type[Table]: """ Converts the DiffableTable into a Table subclass. """ + class_members: t.Dict[str, t.Any] = {} + for column in self.columns: + class_members[column._meta.name] = column + for constraint in self.constraints: + class_members[constraint._meta.name] = constraint + return create_table_class( class_name=self.class_name, class_kwargs={"tablename": self.tablename, "schema": self.schema}, - class_members={ - column._meta.name: column for column in self.columns - }, + class_members=class_members, ) diff --git a/piccolo/apps/migrations/auto/migration_manager.py b/piccolo/apps/migrations/auto/migration_manager.py index fca36e8e7..b1744274a 100644 --- a/piccolo/apps/migrations/auto/migration_manager.py +++ b/piccolo/apps/migrations/auto/migration_manager.py @@ -10,12 +10,14 @@ AlterColumn, ChangeTableSchema, DropColumn, + DropConstraint, RenameColumn, RenameTable, ) from piccolo.apps.migrations.auto.serialisation import deserialise_params from piccolo.columns import Column, column_types from piccolo.columns.column_types import ForeignKey, Serial +from piccolo.constraint import Constraint, UniqueConstraint from piccolo.engine import engine_finder from piccolo.query import Query from piccolo.query.base import DDL @@ -127,6 +129,65 @@ def table_class_names(self) -> t.List[str]: return list({i.table_class_name for i in self.alter_columns}) +@dataclass +class AddConstraintClass: + constraint: Constraint + table_class_name: str + tablename: str + schema: t.Optional[str] + + +@dataclass +class AddConstraintCollection: + add_constraints: t.List[AddConstraintClass] = field(default_factory=list) + + def append(self, add_constraint: AddConstraintClass): + self.add_constraints.append(add_constraint) + + def for_table_class_name( + self, table_class_name: str + ) -> t.List[AddConstraintClass]: + return [ + i + for i in self.add_constraints + if i.table_class_name == table_class_name + ] + + def constraints_for_table_class_name( + self, table_class_name: str + ) -> t.List[Constraint]: + return [ + i.constraint + for i in self.add_constraints + if i.table_class_name == table_class_name + ] + + @property + def table_class_names(self) -> t.List[str]: + return list({i.table_class_name for i in self.add_constraints}) + + +@dataclass +class DropConstraintCollection: + drop_constraints: t.List[DropConstraint] = field(default_factory=list) + + def append(self, drop_constraint: DropConstraint): + self.drop_constraints.append(drop_constraint) + + def for_table_class_name( + self, table_class_name: str + ) -> t.List[DropConstraint]: + return [ + i + for i in self.drop_constraints + if i.table_class_name == table_class_name + ] + + @property + def table_class_names(self) -> t.List[str]: + return list({i.table_class_name for i in self.drop_constraints}) + + AsyncFunction = t.Callable[[], t.Coroutine] @@ -159,6 +220,12 @@ class MigrationManager: alter_columns: AlterColumnCollection = field( default_factory=AlterColumnCollection ) + add_constraints: AddConstraintCollection = field( + default_factory=AddConstraintCollection + ) + drop_constraints: DropConstraintCollection = field( + default_factory=DropConstraintCollection + ) raw: t.List[t.Union[t.Callable, AsyncFunction]] = field( default_factory=list ) @@ -345,6 +412,47 @@ def alter_column( ) ) + def add_constraint( + self, + table_class_name: str, + tablename: str, + constraint_name: str, + constraint_class: t.Type[Constraint], + params: t.Dict[str, t.Any], + schema: t.Optional[str] = None, + ): + if constraint_class is UniqueConstraint: + constraint = UniqueConstraint(**params) + else: + raise ValueError("Unrecognised constraint type") + + constraint._meta.name = constraint_name + + self.add_constraints.append( + AddConstraintClass( + constraint=constraint, + table_class_name=table_class_name, + tablename=tablename, + schema=schema, + ) + ) + + def drop_constraint( + self, + table_class_name: str, + tablename: str, + constraint_name: str, + schema: t.Optional[str] = None, + ): + self.drop_constraints.append( + DropConstraint( + table_class_name=table_class_name, + constraint_name=constraint_name, + tablename=tablename, + schema=schema, + ) + ) + def add_raw(self, raw: t.Union[t.Callable, AsyncFunction]): """ A migration manager can execute arbitrary functions or coroutines when @@ -740,16 +848,24 @@ async def _run_add_tables(self, backwards: bool = False): add_columns: t.List[AddColumnClass] = ( self.add_columns.for_table_class_name(add_table.class_name) ) + add_constraints: t.List[AddConstraintClass] = ( + self.add_constraints.for_table_class_name(add_table.class_name) + ) + class_members: t.Dict[str, t.Any] = {} + for add_column in add_columns: + class_members[add_column.column._meta.name] = add_column.column + for add_constraint in add_constraints: + class_members[add_constraint.constraint._meta.name] = ( + add_constraint.constraint + ) + _Table: t.Type[Table] = create_table_class( class_name=add_table.class_name, class_kwargs={ "tablename": add_table.tablename, "schema": add_table.schema, }, - class_members={ - add_column.column._meta.name: add_column.column - for add_column in add_columns - }, + class_members=class_members, ) table_classes.append(_Table) @@ -922,6 +1038,91 @@ async def _run_change_table_schema(self, backwards: bool = False): ) ) + async def _run_add_constraints(self, backwards: bool = False): + if backwards: + for add_constraint in self.add_constraints.add_constraints: + if add_constraint.table_class_name in [ + i.class_name for i in self.add_tables + ]: + # Don't reverse the add constraint as the table is going to + # be deleted. + continue + + _Table = create_table_class( + class_name=add_constraint.table_class_name, + class_kwargs={ + "tablename": add_constraint.tablename, + "schema": add_constraint.schema, + }, + ) + + await self._run_query( + _Table.alter().drop_constraint( + add_constraint.constraint._meta.name + ) + ) + else: + for table_class_name in self.add_constraints.table_class_names: + if table_class_name in [i.class_name for i in self.add_tables]: + continue # No need to add constraints to new tables + + add_constraints: t.List[AddConstraintClass] = ( + self.add_constraints.for_table_class_name(table_class_name) + ) + + _Table = create_table_class( + class_name=add_constraints[0].table_class_name, + class_kwargs={ + "tablename": add_constraints[0].tablename, + "schema": add_constraints[0].schema, + }, + ) + + for add_constraint in add_constraints: + await self._run_query( + _Table.alter().add_constraint( + add_constraint.constraint + ) + ) + + async def _run_drop_constraints(self, backwards: bool = False): + if backwards: + for drop_constraint in self.drop_constraints.drop_constraints: + _Table = await self.get_table_from_snapshot( + table_class_name=drop_constraint.table_class_name, + app_name=self.app_name, + offset=-1, + ) + constraint_to_restore = _Table._meta.get_constraint_by_name( + drop_constraint.constraint_name + ) + await self._run_query( + _Table.alter().add_constraint(constraint_to_restore) + ) + else: + for table_class_name in self.drop_constraints.table_class_names: + constraints = self.drop_constraints.for_table_class_name( + table_class_name + ) + + if not constraints: + continue + + _Table = create_table_class( + class_name=table_class_name, + class_kwargs={ + "tablename": constraints[0].tablename, + "schema": constraints[0].schema, + }, + ) + + for constraint in constraints: + await self._run_query( + _Table.alter().drop_constraint( + constraint_name=constraint.constraint_name + ) + ) + async def run(self, backwards: bool = False): direction = "backwards" if backwards else "forwards" if self.preview: @@ -958,6 +1159,10 @@ async def run(self, backwards: bool = False): # "ALTER COLUMN TYPE is not supported inside a transaction" if engine.engine_type != "cockroach": await self._run_alter_columns(backwards=backwards) + await self._run_add_constraints(backwards=backwards) + await self._run_drop_constraints(backwards=backwards) if engine.engine_type == "cockroach": await self._run_alter_columns(backwards=backwards) + await self._run_add_constraints(backwards=backwards) + await self._run_drop_constraints(backwards=backwards) diff --git a/piccolo/apps/migrations/auto/operations.py b/piccolo/apps/migrations/auto/operations.py index 0676bdbd4..8fc741e96 100644 --- a/piccolo/apps/migrations/auto/operations.py +++ b/piccolo/apps/migrations/auto/operations.py @@ -2,6 +2,7 @@ from dataclasses import dataclass from piccolo.columns.base import Column +from piccolo.constraint import Constraint @dataclass @@ -63,3 +64,21 @@ class AddColumn: column_class: t.Type[Column] params: t.Dict[str, t.Any] schema: t.Optional[str] = None + + +@dataclass +class AddConstraint: + table_class_name: str + constraint_name: str + constraint_class_name: str + constraint_class: t.Type[Constraint] + params: t.Dict[str, t.Any] + schema: t.Optional[str] = None + + +@dataclass +class DropConstraint: + table_class_name: str + constraint_name: str + tablename: str + schema: t.Optional[str] = None diff --git a/piccolo/apps/migrations/auto/schema_differ.py b/piccolo/apps/migrations/auto/schema_differ.py index 1d095b938..8fe1028d0 100644 --- a/piccolo/apps/migrations/auto/schema_differ.py +++ b/piccolo/apps/migrations/auto/schema_differ.py @@ -613,6 +613,69 @@ def add_columns(self) -> AlterStatements: extra_definitions=extra_definitions, ) + @property + def add_constraints(self) -> AlterStatements: + response: t.List[str] = [] + extra_imports: t.List[Import] = [] + extra_definitions: t.List[Definition] = [] + for table in self.schema: + snapshot_table = self._get_snapshot_table(table.class_name) + if snapshot_table: + delta: TableDelta = table - snapshot_table + else: + continue + + for add_constraint in delta.add_constraints: + constraint_class = add_constraint.constraint_class + extra_imports.append( + Import( + module=constraint_class.__module__, + target=constraint_class.__name__, + expect_conflict_with_global_name=getattr( + UniqueGlobalNames, + f"COLUMN_{constraint_class.__name__.upper()}", + None, + ), + ) + ) + + schema_str = ( + "None" + if add_constraint.schema is None + else f'"{add_constraint.schema}"' + ) + + response.append( + f"manager.add_constraint(table_class_name='{table.class_name}', tablename='{table.tablename}', constraint_name='{add_constraint.constraint_name}', constraint_class={constraint_class.__name__}, params={add_constraint.params}, schema={schema_str})" # noqa: E501 + ) + return AlterStatements( + statements=response, + extra_imports=extra_imports, + extra_definitions=extra_definitions, + ) + + @property + def drop_constraints(self) -> AlterStatements: + response = [] + for table in self.schema: + snapshot_table = self._get_snapshot_table(table.class_name) + if snapshot_table: + delta: TableDelta = table - snapshot_table + else: + continue + + for constraint in delta.drop_constraints: + schema_str = ( + "None" + if constraint.schema is None + else f'"{constraint.schema}"' + ) + + response.append( + f"manager.drop_constraint(table_class_name='{table.class_name}', tablename='{table.tablename}', constraint_name='{constraint.constraint_name}', schema={schema_str})" # noqa: E501 + ) + return AlterStatements(statements=response) + @property def rename_columns(self) -> AlterStatements: alter_statements = AlterStatements() @@ -679,6 +742,48 @@ def new_table_columns(self) -> AlterStatements: extra_definitions=extra_definitions, ) + @property + def new_table_constraints(self) -> AlterStatements: + new_tables: t.List[DiffableTable] = list( + set(self.schema) - set(self.schema_snapshot) + ) + + response: t.List[str] = [] + extra_imports: t.List[Import] = [] + extra_definitions: t.List[Definition] = [] + for table in new_tables: + if ( + table.class_name + in self.rename_tables_collection.new_class_names + ): + continue + + for constraint in table.constraints: + extra_imports.append( + Import( + module=constraint.__class__.__module__, + target=constraint.__class__.__name__, + expect_conflict_with_global_name=getattr( + UniqueGlobalNames, + f"COLUMN_{constraint.__class__.__name__.upper()}", + None, + ), + ) + ) + + schema_str = ( + "None" if table.schema is None else f'"{table.schema}"' + ) + + response.append( + f"manager.add_constraint(table_class_name='{table.class_name}', tablename='{table.tablename}', constraint_name='{constraint._meta.name}', constraint_class={constraint.__class__.__name__}, params={constraint._meta.params}, schema={schema_str})" # noqa: E501 + ) + return AlterStatements( + statements=response, + extra_imports=extra_imports, + extra_definitions=extra_definitions, + ) + ########################################################################### def get_alter_statements(self) -> t.List[AlterStatements]: @@ -691,10 +796,13 @@ def get_alter_statements(self) -> t.List[AlterStatements]: "Renamed tables": self.rename_tables, "Tables which changed schema": self.change_table_schemas, "Created table columns": self.new_table_columns, + "Created table constraints": self.new_table_constraints, "Dropped columns": self.drop_columns, "Columns added to existing tables": self.add_columns, "Renamed columns": self.rename_columns, "Altered columns": self.alter_columns, + "Dropped constraints": self.drop_constraints, + "Constraints added to existing tables": self.add_constraints, } for message, statements in alter_statements.items(): diff --git a/piccolo/apps/migrations/auto/schema_snapshot.py b/piccolo/apps/migrations/auto/schema_snapshot.py index 45963b717..50d8128d7 100644 --- a/piccolo/apps/migrations/auto/schema_snapshot.py +++ b/piccolo/apps/migrations/auto/schema_snapshot.py @@ -112,4 +112,23 @@ def get_snapshot(self) -> t.List[DiffableTable]: rename_column.new_db_column_name ) + add_constraints = ( + manager.add_constraints.constraints_for_table_class_name( + table.class_name + ) + ) + table.constraints.extend(add_constraints) + + drop_constraints = ( + manager.drop_constraints.for_table_class_name( + table.class_name + ) + ) + for drop_constraint in drop_constraints: + table.constraints = [ + i + for i in table.constraints + if i._meta.name != drop_constraint.constraint_name + ] + return tables diff --git a/piccolo/apps/migrations/commands/new.py b/piccolo/apps/migrations/commands/new.py index ff123aaa2..c18c36d31 100644 --- a/piccolo/apps/migrations/commands/new.py +++ b/piccolo/apps/migrations/commands/new.py @@ -191,6 +191,7 @@ async def get_alter_statements( class_name=i.__name__, tablename=i._meta.tablename, columns=i._meta.non_default_columns, + constraints=i._meta.constraints, schema=i._meta.schema, ) for i in app_config.table_classes diff --git a/piccolo/constraint.py b/piccolo/constraint.py new file mode 100644 index 000000000..6bb9e1c95 --- /dev/null +++ b/piccolo/constraint.py @@ -0,0 +1,55 @@ +from __future__ import annotations + +import typing as t +from abc import abstractmethod +from dataclasses import dataclass, field + + +class Constraint: + def __init__(self, **kwargs) -> None: + self._meta = ConstraintMeta(params=kwargs) + + def __hash__(self): + return hash(self._meta.name) + + @property + @abstractmethod + def ddl(self) -> str: + raise NotImplementedError + + +@dataclass +class ConstraintMeta: + # Used for representing the table in migrations. + params: t.Dict[str, t.Any] = field(default_factory=dict) + + # Set by the Table Metaclass: + _name: t.Optional[str] = None + + @property + def name(self) -> str: + if not self._name: + raise ValueError( + "`_name` isn't defined - the Table Metaclass should set it." + ) + return self._name + + @name.setter + def name(self, value: str): + self._name = value + + +class UniqueConstraint(Constraint): + def __init__( + self, + unique_columns: t.List[str], + **kwargs, + ) -> None: + self.unique_columns = unique_columns + kwargs.update({"unique_columns": unique_columns}) + super().__init__(**kwargs) + + @property + def ddl(self) -> str: + unique_columns_string = ", ".join(self.unique_columns) + return f"UNIQUE ({unique_columns_string})" diff --git a/piccolo/query/methods/alter.py b/piccolo/query/methods/alter.py index 040b2f883..10a2427b8 100644 --- a/piccolo/query/methods/alter.py +++ b/piccolo/query/methods/alter.py @@ -6,6 +6,7 @@ from piccolo.columns.base import Column from piccolo.columns.column_types import ForeignKey, Numeric, Varchar +from piccolo.constraint import Constraint from piccolo.query.base import DDL from piccolo.utils.warnings import Level, colored_warning @@ -177,6 +178,17 @@ def ddl(self) -> str: return f'ALTER COLUMN "{self.column_name}" TYPE VARCHAR({self.length})' +@dataclass +class AddConstraint(AlterStatement): + __slots__ = ("constraint",) + + constraint: Constraint + + @property + def ddl(self) -> str: + return f"ADD CONSTRAINT {self.constraint._meta.name} {self.constraint.ddl}" # noqa: E501 + + @dataclass class DropConstraint(AlterStatement): __slots__ = ("constraint_name",) @@ -275,6 +287,7 @@ class Alter(DDL): __slots__ = ( "_add_foreign_key_constraint", "_add", + "_add_constraint", "_drop_constraint", "_drop_default", "_drop_table", @@ -294,6 +307,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_constraint: t.List[AddConstraint] = [] self._drop_constraint: t.List[DropConstraint] = [] self._drop_default: t.List[DropDefault] = [] self._drop_table: t.Optional[DropTable] = None @@ -490,6 +504,10 @@ def _get_constraint_name(self, column: t.Union[str, ForeignKey]) -> str: tablename = self.table._meta.tablename return f"{tablename}_{column_name}_fk" + def add_constraint(self, constraint: Constraint) -> Alter: + self._add_constraint.append(AddConstraint(constraint=constraint)) + return self + def drop_constraint(self, constraint_name: str) -> Alter: self._drop_constraint.append( DropConstraint(constraint_name=constraint_name) @@ -590,6 +608,8 @@ def default_ddl(self) -> t.Sequence[str]: self._set_default, self._set_digits, self._set_schema, + self._add_constraint, + self._drop_constraint, ) ] diff --git a/piccolo/query/methods/create.py b/piccolo/query/methods/create.py index 68cccf6b2..418c45e5b 100644 --- a/piccolo/query/methods/create.py +++ b/piccolo/query/methods/create.py @@ -3,6 +3,7 @@ import typing as t from piccolo.query.base import DDL +from piccolo.query.methods.alter import AddConstraint from piccolo.query.methods.create_index import CreateIndex if t.TYPE_CHECKING: # pragma: no cover @@ -87,4 +88,9 @@ def default_ddl(self) -> t.Sequence[str]: ).ddl ) + for constraint in self.table._meta.constraints: + ddl.append( + f"ALTER TABLE {self.table._meta.get_formatted_tablename()} {AddConstraint(constraint=constraint).ddl}" # noqa: E501 + ) + return ddl diff --git a/piccolo/table.py b/piccolo/table.py index b4fcbf942..6a2e5926a 100644 --- a/piccolo/table.py +++ b/piccolo/table.py @@ -28,6 +28,7 @@ ) from piccolo.columns.readable import Readable from piccolo.columns.reference import LAZY_COLUMN_REFERENCES +from piccolo.constraint import Constraint from piccolo.custom_types import TableInstance from piccolo.engine import Engine, engine_finder from piccolo.query import ( @@ -84,6 +85,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) + constraints: t.List[Constraint] = field(default_factory=list) auto_update_columns: t.List[Column] = field(default_factory=list) tags: t.List[str] = field(default_factory=list) help_text: t.Optional[str] = None @@ -170,6 +172,15 @@ def get_column_by_name(self, name: str) -> Column: return column_object + def get_constraint_by_name(self, name: str) -> Constraint: + """ + Returns a constraint which matches the given name. + """ + for constraint in self.constraints: + if constraint._meta.name == name: + return constraint + raise ValueError(f"No matching constraint found with name == {name}") + def get_auto_update_values(self) -> t.Dict[Column, t.Any]: """ If columns have ``auto_update`` defined, then we retrieve these values. @@ -276,6 +287,7 @@ def __init_subclass__( auto_update_columns: t.List[Column] = [] primary_key: t.Optional[Column] = None m2m_relationships: t.List[M2M] = [] + constraints: t.List[Constraint] = [] attribute_names = itertools.chain( *[i.__dict__.keys() for i in reversed(cls.__mro__)] @@ -328,6 +340,10 @@ def __init_subclass__( attribute._meta._table = cls m2m_relationships.append(attribute) + if isinstance(attribute, Constraint): + attribute._meta._name = attribute_name + constraints.append(attribute) + if not primary_key: primary_key = cls._create_serial_primary_key() setattr(cls, "id", primary_key) @@ -351,6 +367,7 @@ def __init_subclass__( help_text=help_text, _db=db, m2m_relationships=m2m_relationships, + constraints=constraints, schema=schema, ) diff --git a/tests/apps/migrations/auto/test_migration_manager.py b/tests/apps/migrations/auto/test_migration_manager.py index a1988a029..b5d14ccb3 100644 --- a/tests/apps/migrations/auto/test_migration_manager.py +++ b/tests/apps/migrations/auto/test_migration_manager.py @@ -11,6 +11,7 @@ from piccolo.columns.base import OnDelete, OnUpdate from piccolo.columns.column_types import ForeignKey from piccolo.conf.apps import AppConfig +from piccolo.constraint import UniqueConstraint from piccolo.table import Table, sort_table_classes from piccolo.utils.lazy_loader import LazyLoader from tests.base import AsyncMock, DBTestCase, engine_is, engines_only @@ -336,6 +337,282 @@ def test_add_column(self) -> None: if engine_is("cockroach"): self.assertEqual(response, [{"id": row_id, "name": "Dave"}]) + @engines_only("postgres", "cockroach") + def test_add_table_with_unique_constraint(self): + """ + Test adding a table with a unique constraint to a MigrationManager. + """ + # Create table with unique constraint + manager = MigrationManager() + manager.add_table(class_name="Musician", tablename="musician") + manager.add_column( + table_class_name="Musician", + tablename="musician", + column_name="name", + column_class_name="Varchar", + ) + manager.add_column( + table_class_name="Musician", + tablename="musician", + column_name="label", + column_class_name="Varchar", + ) + manager.add_constraint( + table_class_name="Musician", + tablename="musician", + constraint_name="unique_name_label", + constraint_class=UniqueConstraint, + params={ + "unique_columns": ["name", "label"], + }, + ) + asyncio.run(manager.run()) + self.run_sync("INSERT INTO musician VALUES (default, 'a', 'a');") + with self.assertRaises(asyncpg.exceptions.UniqueViolationError): + self.run_sync("INSERT INTO musician VALUES (default, 'a', 'a');") + + # Reverse + asyncio.run(manager.run(backwards=True)) + self.assertTrue(not self.table_exists("musician")) + + @engines_only("postgres", "cockroach") + @patch.object( + BaseMigrationManager, "get_migration_managers", new_callable=AsyncMock + ) + @patch.object(BaseMigrationManager, "get_app_config") + def test_drop_table_with_unique_constraint( + self, get_app_config: MagicMock, get_migration_managers: MagicMock + ): + """ + Test dropping a table with a unique constraint to a MigrationManager. + """ + # Create table + manager_1 = MigrationManager() + manager_1.add_table(class_name="Musician", tablename="musician") + manager_1.add_column( + table_class_name="Musician", + tablename="musician", + column_name="name", + column_class_name="Varchar", + ) + manager_1.add_column( + table_class_name="Musician", + tablename="musician", + column_name="label", + column_class_name="Varchar", + ) + manager_1.add_constraint( + table_class_name="Musician", + tablename="musician", + constraint_name="unique_name_label", + constraint_class=UniqueConstraint, + params={ + "unique_columns": ["name", "label"], + }, + ) + asyncio.run(manager_1.run()) + + # Drop table + manager_2 = MigrationManager() + manager_2.drop_table( + class_name="Musician", + tablename="musician", + ) + asyncio.run(manager_2.run()) + self.assertTrue(not self.table_exists("musician")) + + # Reverse + get_migration_managers.return_value = [manager_1] + app_config = AppConfig(app_name="music", migrations_folder_path="") + get_app_config.return_value = app_config + asyncio.run(manager_2.run(backwards=True)) + self.run_sync("INSERT INTO musician VALUES (default, 'a', 'a');") + with self.assertRaises(asyncpg.exceptions.UniqueViolationError): + self.run_sync("INSERT INTO musician VALUES (default, 'a', 'a');") + + # Reverse + asyncio.run(manager_1.run(backwards=True)) + self.assertTrue(not self.table_exists("musician")) + + @engines_only("postgres", "cockroach") + @patch.object( + BaseMigrationManager, "get_migration_managers", new_callable=AsyncMock + ) + @patch.object(BaseMigrationManager, "get_app_config") + def test_rename_table_with_unique_constraint( + self, get_app_config: MagicMock, get_migration_managers: MagicMock + ): + """ + Test renaming a table with a unique constraint to a MigrationManager. + """ + # Create table + manager_1 = MigrationManager() + manager_1.add_table(class_name="Musician", tablename="musician") + manager_1.add_column( + table_class_name="Musician", + tablename="musician", + column_name="name", + column_class_name="Varchar", + ) + manager_1.add_column( + table_class_name="Musician", + tablename="musician", + column_name="label", + column_class_name="Varchar", + ) + manager_1.add_constraint( + table_class_name="Musician", + tablename="musician", + constraint_name="unique_name_label", + constraint_class=UniqueConstraint, + params={ + "unique_columns": ["name", "label"], + }, + ) + asyncio.run(manager_1.run()) + + # Rename table + manager_2 = MigrationManager() + manager_2.rename_table( + old_class_name="Musician", + old_tablename="musician", + new_class_name="Musician2", + new_tablename="musician2", + ) + asyncio.run(manager_2.run()) + self.assertTrue(not self.table_exists("musician")) + self.run_sync("INSERT INTO musician2 VALUES (default, 'a', 'a');") + with self.assertRaises(asyncpg.exceptions.UniqueViolationError): + self.run_sync("INSERT INTO musician2 VALUES (default, 'a', 'a');") + + # Reverse + get_migration_managers.return_value = [manager_1] + app_config = AppConfig(app_name="music", migrations_folder_path="") + get_app_config.return_value = app_config + asyncio.run(manager_2.run(backwards=True)) + self.assertTrue(not self.table_exists("musician2")) + with self.assertRaises(asyncpg.exceptions.UniqueViolationError): + self.run_sync("INSERT INTO musician VALUES (default, 'a', 'a');") + + # Reverse + asyncio.run(manager_1.run(backwards=True)) + self.assertTrue(not self.table_exists("musician")) + self.assertTrue(not self.table_exists("musician2")) + + @engines_only("postgres") + @patch.object( + BaseMigrationManager, "get_migration_managers", new_callable=AsyncMock + ) + @patch.object(BaseMigrationManager, "get_app_config") + def test_add_unique_constraint( + self, get_app_config: MagicMock, get_migration_managers: MagicMock + ): + """ + Test adding a unique constraint to a MigrationManager. + Cockroach DB doesn't support dropping unique constraints with ALTER TABLE DROP CONSTRAINT. + https://github.com/cockroachdb/cockroach/issues/42840 + """ # noqa: E501 + manager_1 = MigrationManager() + manager_1.add_table(class_name="Musician", tablename="musician") + manager_1.add_column( + table_class_name="Musician", + tablename="musician", + column_name="name", + column_class_name="Varchar", + ) + manager_1.add_column( + table_class_name="Musician", + tablename="musician", + column_name="label", + column_class_name="Varchar", + ) + asyncio.run(manager_1.run()) + + manager_2 = MigrationManager() + manager_2.add_constraint( + table_class_name="Musician", + tablename="musician", + constraint_name="musician_unique", + constraint_class=UniqueConstraint, + params={ + "unique_columns": ["name", "label"], + }, + ) + asyncio.run(manager_2.run()) + self.run_sync("INSERT INTO musician VALUES (default, 'a', 'a');") + with self.assertRaises(asyncpg.exceptions.UniqueViolationError): + self.run_sync("INSERT INTO musician VALUES (default, 'a', 'a');") + + # Reverse + get_migration_managers.return_value = [manager_1] + app_config = AppConfig(app_name="music", migrations_folder_path="") + get_app_config.return_value = app_config + asyncio.run(manager_2.run(backwards=True)) + self.run_sync("INSERT INTO musician VALUES (default, 'a', 'a');") + + # Reverse + asyncio.run(manager_1.run(backwards=True)) + self.assertTrue(not self.table_exists("musician")) + + @engines_only("postgres") + @patch.object( + BaseMigrationManager, "get_migration_managers", new_callable=AsyncMock + ) + @patch.object(BaseMigrationManager, "get_app_config") + def test_drop_unique_constraint( + self, get_app_config: MagicMock, get_migration_managers: MagicMock + ): + """ + Test dropping a unique constraint with a MigrationManager. + Cockroach DB doesn't support dropping unique constraints with ALTER TABLE DROP CONSTRAINT. + https://github.com/cockroachdb/cockroach/issues/42840 + """ # noqa: E501 + manager_1 = MigrationManager() + manager_1.add_table(class_name="Musician", tablename="musician") + manager_1.add_column( + table_class_name="Musician", + tablename="musician", + column_name="name", + column_class_name="Varchar", + ) + manager_1.add_column( + table_class_name="Musician", + tablename="musician", + column_name="label", + column_class_name="Varchar", + ) + manager_1.add_constraint( + table_class_name="Musician", + tablename="musician", + constraint_name="musician_unique", + constraint_class=UniqueConstraint, + params={ + "unique_columns": ["name", "label"], + }, + ) + asyncio.run(manager_1.run()) + + manager_2 = MigrationManager() + manager_2.drop_constraint( + table_class_name="Musician", + tablename="musician", + constraint_name="musician_unique", + ) + asyncio.run(manager_2.run()) + + # Reverse + get_migration_managers.return_value = [manager_1] + app_config = AppConfig(app_name="music", migrations_folder_path="") + get_app_config.return_value = app_config + asyncio.run(manager_2.run(backwards=True)) + self.run_sync("INSERT INTO musician VALUES (default, 'a', 'a');") + with self.assertRaises(asyncpg.exceptions.UniqueViolationError): + self.run_sync("INSERT INTO musician VALUES (default, 'a', 'a');") + + # Reverse + asyncio.run(manager_1.run(backwards=True)) + self.assertTrue(not self.table_exists("musician")) + @engines_only("postgres", "cockroach") def test_add_column_with_index(self): """ diff --git a/tests/apps/migrations/auto/test_schema_differ.py b/tests/apps/migrations/auto/test_schema_differ.py index 9cf6d26f2..b9fe5252c 100644 --- a/tests/apps/migrations/auto/test_schema_differ.py +++ b/tests/apps/migrations/auto/test_schema_differ.py @@ -13,6 +13,7 @@ SchemaDiffer, ) from piccolo.columns.column_types import Numeric, Varchar +from piccolo.constraint import UniqueConstraint class TestSchemaDiffer(TestCase): @@ -488,6 +489,92 @@ def test_db_column_name(self) -> None: "manager.alter_column(table_class_name='Ticket', tablename='ticket', column_name='price', db_column_name='custom', params={'digits': (4, 2)}, old_params={'digits': (5, 2)}, column_class=Numeric, old_column_class=Numeric, schema=None)", # noqa ) + def test_add_constraint(self) -> None: + """ + Test adding a constraint to an existing table. + """ + name_column = Varchar() + name_column._meta.name = "name" + + genre_column = Varchar() + genre_column._meta.name = "genre" + + name_unique_constraint = UniqueConstraint(unique_columns=["name"]) + name_unique_constraint._meta.name = "unique_name" + + name_genre_unique_constraint = UniqueConstraint(unique_columns=["name", "genre"]) + name_genre_unique_constraint._meta.name = "unique_name_genre" + + schema: t.List[DiffableTable] = [ + DiffableTable( + class_name="Band", + tablename="band", + columns=[name_column, genre_column], + constraints=[name_unique_constraint, name_genre_unique_constraint], + ) + ] + schema_snapshot: t.List[DiffableTable] = [ + DiffableTable( + class_name="Band", + tablename="band", + columns=[name_column, genre_column], + constraints=[name_unique_constraint], + ) + ] + + schema_differ = SchemaDiffer( + schema=schema, schema_snapshot=schema_snapshot, auto_input="y" + ) + + self.assertTrue(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='unique_name_genre', constraint_class=UniqueConstraint, params={'unique_columns': ['name', 'genre']}, schema=None)" # noqa: E501 + ) + + def test_drop_constraint(self) -> None: + """ + Test dropping a constraint from an existing table. + """ + name_column = Varchar() + name_column._meta.name = "name" + + genre_column = Varchar() + genre_column._meta.name = "genre" + + name_unique_constraint = UniqueConstraint(unique_columns=["name"]) + name_unique_constraint._meta.name = "unique_name" + + name_genre_unique_constraint = UniqueConstraint(unique_columns=["name", "genre"]) + name_genre_unique_constraint._meta.name = "unique_name_genre" + + schema: t.List[DiffableTable] = [ + DiffableTable( + class_name="Band", + tablename="band", + columns=[name_column, genre_column], + constraints=[name_unique_constraint], + ) + ] + schema_snapshot: t.List[DiffableTable] = [ + DiffableTable( + class_name="Band", + tablename="band", + columns=[name_column, genre_column], + constraints=[name_unique_constraint, name_genre_unique_constraint], + ) + ] + + schema_differ = SchemaDiffer( + schema=schema, schema_snapshot=schema_snapshot, auto_input="y" + ) + + self.assertTrue(len(schema_differ.drop_constraints.statements) == 1) + self.assertEqual( + schema_differ.drop_constraints.statements[0], + "manager.drop_constraint(table_class_name='Band', tablename='band', constraint_name='unique_name_genre', schema=None)" # noqa: E501 + ) + def test_alter_default(self): pass From 07f8b941c09fc43d378521a624a867b2e243633e Mon Sep 17 00:00:00 2001 From: Keita Ichihashi Date: Fri, 12 Apr 2024 19:21:01 +0900 Subject: [PATCH 02/29] Fix code format --- .../migrations/auto/test_migration_manager.py | 14 +----------- .../migrations/auto/test_schema_differ.py | 22 ++++++++++++++----- 2 files changed, 17 insertions(+), 19 deletions(-) diff --git a/tests/apps/migrations/auto/test_migration_manager.py b/tests/apps/migrations/auto/test_migration_manager.py index b5d14ccb3..8f8e6e628 100644 --- a/tests/apps/migrations/auto/test_migration_manager.py +++ b/tests/apps/migrations/auto/test_migration_manager.py @@ -339,10 +339,6 @@ def test_add_column(self) -> None: @engines_only("postgres", "cockroach") def test_add_table_with_unique_constraint(self): - """ - Test adding a table with a unique constraint to a MigrationManager. - """ - # Create table with unique constraint manager = MigrationManager() manager.add_table(class_name="Musician", tablename="musician") manager.add_column( @@ -383,10 +379,6 @@ def test_add_table_with_unique_constraint(self): def test_drop_table_with_unique_constraint( self, get_app_config: MagicMock, get_migration_managers: MagicMock ): - """ - Test dropping a table with a unique constraint to a MigrationManager. - """ - # Create table manager_1 = MigrationManager() manager_1.add_table(class_name="Musician", tablename="musician") manager_1.add_column( @@ -440,12 +432,8 @@ def test_drop_table_with_unique_constraint( ) @patch.object(BaseMigrationManager, "get_app_config") def test_rename_table_with_unique_constraint( - self, get_app_config: MagicMock, get_migration_managers: MagicMock + self, get_app_config: MagicMock, get_migration_managers: MagicMock ): - """ - Test renaming a table with a unique constraint to a MigrationManager. - """ - # Create table manager_1 = MigrationManager() manager_1.add_table(class_name="Musician", tablename="musician") manager_1.add_column( diff --git a/tests/apps/migrations/auto/test_schema_differ.py b/tests/apps/migrations/auto/test_schema_differ.py index b9fe5252c..d6f9f5de5 100644 --- a/tests/apps/migrations/auto/test_schema_differ.py +++ b/tests/apps/migrations/auto/test_schema_differ.py @@ -502,7 +502,9 @@ def test_add_constraint(self) -> None: name_unique_constraint = UniqueConstraint(unique_columns=["name"]) name_unique_constraint._meta.name = "unique_name" - name_genre_unique_constraint = UniqueConstraint(unique_columns=["name", "genre"]) + name_genre_unique_constraint = UniqueConstraint( + unique_columns=["name", "genre"] + ) name_genre_unique_constraint._meta.name = "unique_name_genre" schema: t.List[DiffableTable] = [ @@ -510,7 +512,10 @@ def test_add_constraint(self) -> None: class_name="Band", tablename="band", columns=[name_column, genre_column], - constraints=[name_unique_constraint, name_genre_unique_constraint], + constraints=[ + name_unique_constraint, + name_genre_unique_constraint, + ], ) ] schema_snapshot: t.List[DiffableTable] = [ @@ -529,7 +534,7 @@ def test_add_constraint(self) -> None: self.assertTrue(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='unique_name_genre', constraint_class=UniqueConstraint, params={'unique_columns': ['name', 'genre']}, schema=None)" # noqa: E501 + "manager.add_constraint(table_class_name='Band', tablename='band', constraint_name='unique_name_genre', constraint_class=UniqueConstraint, params={'unique_columns': ['name', 'genre']}, schema=None)", # noqa: E501 ) def test_drop_constraint(self) -> None: @@ -545,7 +550,9 @@ def test_drop_constraint(self) -> None: name_unique_constraint = UniqueConstraint(unique_columns=["name"]) name_unique_constraint._meta.name = "unique_name" - name_genre_unique_constraint = UniqueConstraint(unique_columns=["name", "genre"]) + name_genre_unique_constraint = UniqueConstraint( + unique_columns=["name", "genre"] + ) name_genre_unique_constraint._meta.name = "unique_name_genre" schema: t.List[DiffableTable] = [ @@ -561,7 +568,10 @@ def test_drop_constraint(self) -> None: class_name="Band", tablename="band", columns=[name_column, genre_column], - constraints=[name_unique_constraint, name_genre_unique_constraint], + constraints=[ + name_unique_constraint, + name_genre_unique_constraint, + ], ) ] @@ -572,7 +582,7 @@ def test_drop_constraint(self) -> None: self.assertTrue(len(schema_differ.drop_constraints.statements) == 1) self.assertEqual( schema_differ.drop_constraints.statements[0], - "manager.drop_constraint(table_class_name='Band', tablename='band', constraint_name='unique_name_genre', schema=None)" # noqa: E501 + "manager.drop_constraint(table_class_name='Band', tablename='band', constraint_name='unique_name_genre', schema=None)", # noqa: E501 ) def test_alter_default(self): From e4a5209964d232809ebf8b351e7f720f4c5794d6 Mon Sep 17 00:00:00 2001 From: Keita Ichihashi Date: Thu, 18 Apr 2024 18:39:03 +0900 Subject: [PATCH 03/29] Add tests --- .../migrations/auto/test_schema_differ.py | 57 ++++++++++++++++ .../migrations/auto/test_schema_snapshot.py | 65 +++++++++++++++++++ 2 files changed, 122 insertions(+) diff --git a/tests/apps/migrations/auto/test_schema_differ.py b/tests/apps/migrations/auto/test_schema_differ.py index d6f9f5de5..8768cf6b9 100644 --- a/tests/apps/migrations/auto/test_schema_differ.py +++ b/tests/apps/migrations/auto/test_schema_differ.py @@ -489,6 +489,63 @@ def test_db_column_name(self) -> None: "manager.alter_column(table_class_name='Ticket', tablename='ticket', column_name='price', db_column_name='custom', params={'digits': (4, 2)}, old_params={'digits': (5, 2)}, column_class=Numeric, old_column_class=Numeric, schema=None)", # noqa ) + def test_add_table_with_constraint(self) -> None: + """ + Test adding a new table with a constraint. + """ + name_column = Varchar() + name_column._meta.name = "name" + + genre_column = Varchar() + genre_column._meta.name = "genre" + + name_unique_constraint = UniqueConstraint(unique_columns=["name"]) + name_unique_constraint._meta.name = "unique_name" + + name_genre_unique_constraint = UniqueConstraint( + unique_columns=["name", "genre"] + ) + name_genre_unique_constraint._meta.name = "unique_name_genre" + + schema: t.List[DiffableTable] = [ + DiffableTable( + class_name="Band", + tablename="band", + columns=[name_column, genre_column], + constraints=[name_genre_unique_constraint], + ) + ] + schema_snapshot: t.List[DiffableTable] = [] + + schema_differ = SchemaDiffer( + schema=schema, schema_snapshot=schema_snapshot, auto_input="y" + ) + + create_tables = schema_differ.create_tables + self.assertTrue(len(create_tables.statements) == 1) + self.assertEqual( + create_tables.statements[0], + "manager.add_table(class_name='Band', tablename='band', schema=None, columns=None)", # noqa: E501 + ) + + new_table_columns = schema_differ.new_table_columns + self.assertTrue(len(new_table_columns.statements) == 2) + self.assertEqual( + new_table_columns.statements[0], + "manager.add_column(table_class_name='Band', tablename='band', column_name='name', db_column_name='name', column_class_name='Varchar', column_class=Varchar, params={'length': 255, 'default': '', 'null': False, 'primary_key': False, 'unique': False, 'index': False, 'index_method': IndexMethod.btree, 'choices': None, 'db_column_name': None, 'secret': False}, schema=None)", # noqa + ) + self.assertEqual( + new_table_columns.statements[1], + "manager.add_column(table_class_name='Band', tablename='band', column_name='genre', db_column_name='genre', column_class_name='Varchar', column_class=Varchar, params={'length': 255, 'default': '', 'null': False, 'primary_key': False, 'unique': False, 'index': False, 'index_method': IndexMethod.btree, 'choices': None, 'db_column_name': None, 'secret': False}, schema=None)", # noqa + ) + + new_table_constraints = schema_differ.new_table_constraints + self.assertTrue(len(new_table_constraints.statements) == 1) + self.assertEqual( + new_table_constraints.statements[0], + "manager.add_constraint(table_class_name='Band', tablename='band', constraint_name='unique_name_genre', constraint_class=UniqueConstraint, params={'unique_columns': ['name', 'genre']}, schema=None)", # noqa + ) + def test_add_constraint(self) -> None: """ Test adding a constraint to an existing table. diff --git a/tests/apps/migrations/auto/test_schema_snapshot.py b/tests/apps/migrations/auto/test_schema_snapshot.py index 834551f8a..901ddc02f 100644 --- a/tests/apps/migrations/auto/test_schema_snapshot.py +++ b/tests/apps/migrations/auto/test_schema_snapshot.py @@ -1,6 +1,7 @@ from unittest import TestCase from piccolo.apps.migrations.auto import MigrationManager, SchemaSnapshot +from piccolo.constraint import UniqueConstraint class TestSchemaSnaphot(TestCase): @@ -187,3 +188,67 @@ def test_get_table_from_snapshot(self): with self.assertRaises(ValueError): schema_snapshot.get_table_from_snapshot("Foo") + + def test_add_constraint(self): + """ + Test adding constraints. + """ + manager = MigrationManager() + manager.add_table(class_name="Manager", tablename="manager") + manager.add_column( + table_class_name="Manager", + tablename="manager", + column_name="name", + column_class_name="Varchar", + params={"length": 100}, + ) + manager.add_constraint( + table_class_name="Manager", + tablename="manager", + constraint_name="unique_name", + constraint_class=UniqueConstraint, + params={ + "unique_columns": ["name"], + }, + ) + + schema_snapshot = SchemaSnapshot(managers=[manager]) + snapshot = schema_snapshot.get_snapshot() + + self.assertTrue(len(snapshot) == 1) + self.assertTrue(len(snapshot[0].columns) == 1) + self.assertTrue(len(snapshot[0].constraints) == 1) + + def test_drop_constraint(self): + """ + Test dropping constraints. + """ + manager_1 = MigrationManager() + manager_1.add_table(class_name="Manager", tablename="manager") + manager_1.add_column( + table_class_name="Manager", + tablename="manager", + column_name="name", + column_class_name="Varchar", + params={"length": 100}, + ) + manager_1.add_constraint( + table_class_name="Manager", + tablename="manager", + constraint_name="unique_name", + constraint_class=UniqueConstraint, + params={ + "unique_columns": ["name"], + }, + ) + + manager_2 = MigrationManager() + manager_2.drop_constraint( + table_class_name="Manager", + tablename="manager", + constraint_name="unique_name", + ) + + schema_snapshot = SchemaSnapshot(managers=[manager_1, manager_2]) + snapshot = schema_snapshot.get_snapshot() + self.assertEqual(len(snapshot[0].constraints), 0) From 012c6088cae5f104c96d06a6dc2e913b780170ab Mon Sep 17 00:00:00 2001 From: Keita Ichihashi Date: Sat, 19 Oct 2024 16:33:23 +0900 Subject: [PATCH 04/29] Add comments --- piccolo/constraint.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/piccolo/constraint.py b/piccolo/constraint.py index 6bb9e1c95..93d8a163d 100644 --- a/piccolo/constraint.py +++ b/piccolo/constraint.py @@ -6,6 +6,10 @@ class Constraint: + """ + All other constraints inherit from ``Constraint``. Don't use it directly. + """ + def __init__(self, **kwargs) -> None: self._meta = ConstraintMeta(params=kwargs) @@ -20,6 +24,10 @@ def ddl(self) -> str: @dataclass class ConstraintMeta: + """ + This is used to store info about the constraint. + """ + # Used for representing the table in migrations. params: t.Dict[str, t.Any] = field(default_factory=dict) @@ -40,11 +48,19 @@ def name(self, value: str): class UniqueConstraint(Constraint): + """ + Unique constraint on the table columns. + """ + def __init__( self, unique_columns: t.List[str], **kwargs, ) -> None: + """ + :param unique_columns: + The table columns that should be unique together. + """ self.unique_columns = unique_columns kwargs.update({"unique_columns": unique_columns}) super().__init__(**kwargs) From fb02cd51cd5b89381e011b332de8c879ad0f2ab3 Mon Sep 17 00:00:00 2001 From: Daniel Townsend Date: Fri, 24 Jan 2025 13:53:26 +0000 Subject: [PATCH 05/29] `add_constraints` method --- docs/src/piccolo/schema/constraints.rst | 41 ++++++++++++++ docs/src/piccolo/schema/index.rst | 1 + piccolo/constraint.py | 71 ++++++++++++++++--------- piccolo/table.py | 22 ++++++++ 4 files changed, 109 insertions(+), 26 deletions(-) create mode 100644 docs/src/piccolo/schema/constraints.rst diff --git a/docs/src/piccolo/schema/constraints.rst b/docs/src/piccolo/schema/constraints.rst new file mode 100644 index 000000000..9ab5f6491 --- /dev/null +++ b/docs/src/piccolo/schema/constraints.rst @@ -0,0 +1,41 @@ +=========== +Constraints +=========== + +Single column +============= + +Unique constraints can be added to a single column using the ``unique=True`` +argument of ``Column``: + +.. code-block:: python + + class Band(Table): + name = Varchar(unique=True) + +------------------------------------------------------------------------------- + +Multi-column +============ + +``add_constraints`` +------------------- + +Use the ``add_constraints`` class method to add multi-column constraints to a +``Table``: + +.. currentmodule:: piccolo.table + +.. automethod:: Table.add_constraints + +------------------------------------------------------------------------------- + +Constraint types +================ + +``UniqueConstraint`` +-------------------- + +.. currentmodule:: piccolo.constraint + +.. autoclass:: UniqueConstraint diff --git a/docs/src/piccolo/schema/index.rst b/docs/src/piccolo/schema/index.rst index ec9b887e6..29396d7df 100644 --- a/docs/src/piccolo/schema/index.rst +++ b/docs/src/piccolo/schema/index.rst @@ -8,6 +8,7 @@ The schema is how you define your database tables, columns and relationships. ./defining ./column_types + ./constraints ./m2m ./one_to_one ./advanced diff --git a/piccolo/constraint.py b/piccolo/constraint.py index 93d8a163d..617d5b48c 100644 --- a/piccolo/constraint.py +++ b/piccolo/constraint.py @@ -4,14 +4,16 @@ from abc import abstractmethod from dataclasses import dataclass, field +from piccolo.columns import Column + class Constraint: """ All other constraints inherit from ``Constraint``. Don't use it directly. """ - def __init__(self, **kwargs) -> None: - self._meta = ConstraintMeta(params=kwargs) + def __init__(self, name: str, **kwargs) -> None: + self._meta = ConstraintMeta(name=name, params=kwargs) def __hash__(self): return hash(self._meta.name) @@ -28,24 +30,11 @@ class ConstraintMeta: This is used to store info about the constraint. """ + name: str + # Used for representing the table in migrations. params: t.Dict[str, t.Any] = field(default_factory=dict) - # Set by the Table Metaclass: - _name: t.Optional[str] = None - - @property - def name(self) -> str: - if not self._name: - raise ValueError( - "`_name` isn't defined - the Table Metaclass should set it." - ) - return self._name - - @name.setter - def name(self, value: str): - self._name = value - class UniqueConstraint(Constraint): """ @@ -54,18 +43,48 @@ class UniqueConstraint(Constraint): def __init__( self, - unique_columns: t.List[str], - **kwargs, + *columns: Column, + name: t.Optional[str] = None, + nulls_distinct: bool = True, ) -> None: """ - :param unique_columns: + :param columns: The table columns that should be unique together. - """ - self.unique_columns = unique_columns - kwargs.update({"unique_columns": unique_columns}) - super().__init__(**kwargs) + :param name: + The name of the constraint in the database. If not provided, we + generate a sensible default. + :param nulls_distinct: + See the `Postgres docs `_ + for more information. + + """ # noqa: E501 + if len(columns) < 1: + raise ValueError("At least 1 column must be specified.") + + tablenames = [column._meta.table._meta.tablename for column in columns] + + if len(set(tablenames)) > 1: + raise ValueError("The columns belong to different tables.") + + column_names = [column._meta.db_column_name for column in columns] + self.column_names = column_names + + if name is None: + name = "_".join([tablenames[0], "unique", *column_names]) + self.name = name + + self.nulls_distinct = nulls_distinct + + super().__init__( + name=name, + column_names=column_names, + nulls_distinct=nulls_distinct, + ) @property def ddl(self) -> str: - unique_columns_string = ", ".join(self.unique_columns) - return f"UNIQUE ({unique_columns_string})" + nulls_string = ( + "NULLS NOT DISTINCT " if self.nulls_distinct is False else "" + ) + columns_string = ", ".join(f'"{i}"' for i in self.column_names) + return f"UNIQUE {nulls_string}({columns_string})" diff --git a/piccolo/table.py b/piccolo/table.py index ccd935eb0..9114516ed 100644 --- a/piccolo/table.py +++ b/piccolo/table.py @@ -482,6 +482,28 @@ def from_dict( """ return cls(**data) + @classmethod + def add_constraints(cls, *constraint: Constraint): + """ + Add composite constraints to the table (e.g. a unique constraint across + multiple columns). + + You should wait for the ``Table`` to be initialised before calling + this method. For example:: + + from piccolo.constraint import UniqueConstraint + + class Album(Table): + name = Varchar() + band = ForeignKey(Band) + + Album.add_constraints( + UniqueConstraint(Album.name, Album.band) + ) + + """ + cls._meta.constraints.extend(constraint) + ########################################################################### def save( From 074e2734e4d0744b52fee91527af14f57527683a Mon Sep 17 00:00:00 2001 From: Daniel Townsend Date: Fri, 24 Jan 2025 21:12:42 +0000 Subject: [PATCH 06/29] change module name from `contraint` to `constraints` We're likely to have unique, and check constraints so I guess it makes sense to have it plural. --- docs/src/piccolo/schema/constraints.rst | 2 +- piccolo/apps/migrations/auto/diffable_table.py | 2 +- piccolo/apps/migrations/auto/migration_manager.py | 2 +- piccolo/apps/migrations/auto/operations.py | 2 +- piccolo/{constraint.py => constraints.py} | 0 piccolo/query/methods/alter.py | 2 +- piccolo/table.py | 4 ++-- tests/apps/migrations/auto/test_migration_manager.py | 2 +- tests/apps/migrations/auto/test_schema_differ.py | 2 +- tests/apps/migrations/auto/test_schema_snapshot.py | 2 +- 10 files changed, 10 insertions(+), 10 deletions(-) rename piccolo/{constraint.py => constraints.py} (100%) diff --git a/docs/src/piccolo/schema/constraints.rst b/docs/src/piccolo/schema/constraints.rst index 9ab5f6491..b7011148c 100644 --- a/docs/src/piccolo/schema/constraints.rst +++ b/docs/src/piccolo/schema/constraints.rst @@ -21,7 +21,7 @@ Multi-column ``add_constraints`` ------------------- -Use the ``add_constraints`` class method to add multi-column constraints to a +Use the ``add_constraints`` method to add multi-column constraints to a ``Table``: .. currentmodule:: piccolo.table diff --git a/piccolo/apps/migrations/auto/diffable_table.py b/piccolo/apps/migrations/auto/diffable_table.py index 516776b1f..07ac959e0 100644 --- a/piccolo/apps/migrations/auto/diffable_table.py +++ b/piccolo/apps/migrations/auto/diffable_table.py @@ -15,7 +15,7 @@ serialise_params, ) from piccolo.columns.base import Column -from piccolo.constraint import Constraint +from piccolo.constraints import Constraint from piccolo.table import Table, create_table_class diff --git a/piccolo/apps/migrations/auto/migration_manager.py b/piccolo/apps/migrations/auto/migration_manager.py index e5890af33..c525c7425 100644 --- a/piccolo/apps/migrations/auto/migration_manager.py +++ b/piccolo/apps/migrations/auto/migration_manager.py @@ -17,7 +17,7 @@ from piccolo.apps.migrations.auto.serialisation import deserialise_params from piccolo.columns import Column, column_types from piccolo.columns.column_types import ForeignKey, Serial -from piccolo.constraint import Constraint, UniqueConstraint +from piccolo.constraints import Constraint, UniqueConstraint from piccolo.engine import engine_finder from piccolo.query import Query from piccolo.query.base import DDL diff --git a/piccolo/apps/migrations/auto/operations.py b/piccolo/apps/migrations/auto/operations.py index 8fc741e96..a7d28aa20 100644 --- a/piccolo/apps/migrations/auto/operations.py +++ b/piccolo/apps/migrations/auto/operations.py @@ -2,7 +2,7 @@ from dataclasses import dataclass from piccolo.columns.base import Column -from piccolo.constraint import Constraint +from piccolo.constraints import Constraint @dataclass diff --git a/piccolo/constraint.py b/piccolo/constraints.py similarity index 100% rename from piccolo/constraint.py rename to piccolo/constraints.py diff --git a/piccolo/query/methods/alter.py b/piccolo/query/methods/alter.py index 10a2427b8..fbfdf5755 100644 --- a/piccolo/query/methods/alter.py +++ b/piccolo/query/methods/alter.py @@ -6,7 +6,7 @@ from piccolo.columns.base import Column from piccolo.columns.column_types import ForeignKey, Numeric, Varchar -from piccolo.constraint import Constraint +from piccolo.constraints import Constraint from piccolo.query.base import DDL from piccolo.utils.warnings import Level, colored_warning diff --git a/piccolo/table.py b/piccolo/table.py index 9114516ed..4c0523d29 100644 --- a/piccolo/table.py +++ b/piccolo/table.py @@ -28,7 +28,7 @@ ) from piccolo.columns.readable import Readable from piccolo.columns.reference import LAZY_COLUMN_REFERENCES -from piccolo.constraint import Constraint +from piccolo.constraints import Constraint from piccolo.custom_types import TableInstance from piccolo.engine import Engine, engine_finder from piccolo.query import ( @@ -491,7 +491,7 @@ def add_constraints(cls, *constraint: Constraint): You should wait for the ``Table`` to be initialised before calling this method. For example:: - from piccolo.constraint import UniqueConstraint + from piccolo.constraints import UniqueConstraint class Album(Table): name = Varchar() diff --git a/tests/apps/migrations/auto/test_migration_manager.py b/tests/apps/migrations/auto/test_migration_manager.py index 7055c5912..e72fa77e1 100644 --- a/tests/apps/migrations/auto/test_migration_manager.py +++ b/tests/apps/migrations/auto/test_migration_manager.py @@ -11,7 +11,7 @@ from piccolo.columns.base import OnDelete, OnUpdate from piccolo.columns.column_types import ForeignKey from piccolo.conf.apps import AppConfig -from piccolo.constraint import UniqueConstraint +from piccolo.constraints import UniqueConstraint from piccolo.engine import engine_finder from piccolo.table import Table, sort_table_classes from piccolo.utils.lazy_loader import LazyLoader diff --git a/tests/apps/migrations/auto/test_schema_differ.py b/tests/apps/migrations/auto/test_schema_differ.py index 8768cf6b9..92b68f59a 100644 --- a/tests/apps/migrations/auto/test_schema_differ.py +++ b/tests/apps/migrations/auto/test_schema_differ.py @@ -13,7 +13,7 @@ SchemaDiffer, ) from piccolo.columns.column_types import Numeric, Varchar -from piccolo.constraint import UniqueConstraint +from piccolo.constraints import UniqueConstraint class TestSchemaDiffer(TestCase): diff --git a/tests/apps/migrations/auto/test_schema_snapshot.py b/tests/apps/migrations/auto/test_schema_snapshot.py index 901ddc02f..ae1742c5b 100644 --- a/tests/apps/migrations/auto/test_schema_snapshot.py +++ b/tests/apps/migrations/auto/test_schema_snapshot.py @@ -1,7 +1,7 @@ from unittest import TestCase from piccolo.apps.migrations.auto import MigrationManager, SchemaSnapshot -from piccolo.constraint import UniqueConstraint +from piccolo.constraints import UniqueConstraint class TestSchemaSnaphot(TestCase): From a953b990691c4d16d3302bdc2a6c6e63f1f57687 Mon Sep 17 00:00:00 2001 From: Daniel Townsend Date: Sat, 25 Jan 2025 20:53:02 +0000 Subject: [PATCH 07/29] refactor -> `add_unique_constraint` --- .../apps/migrations/auto/migration_manager.py | 12 ++-- piccolo/constraints.py | 23 ++------ piccolo/table.py | 58 +++++++++++++++---- .../migrations/auto/test_schema_differ.py | 1 + 4 files changed, 58 insertions(+), 36 deletions(-) diff --git a/piccolo/apps/migrations/auto/migration_manager.py b/piccolo/apps/migrations/auto/migration_manager.py index c525c7425..21a74c1f8 100644 --- a/piccolo/apps/migrations/auto/migration_manager.py +++ b/piccolo/apps/migrations/auto/migration_manager.py @@ -867,16 +867,9 @@ async def _run_add_tables(self, backwards: bool = False): add_columns: t.List[AddColumnClass] = ( self.add_columns.for_table_class_name(add_table.class_name) ) - add_constraints: t.List[AddConstraintClass] = ( - self.add_constraints.for_table_class_name(add_table.class_name) - ) class_members: t.Dict[str, t.Any] = {} for add_column in add_columns: class_members[add_column.column._meta.name] = add_column.column - for add_constraint in add_constraints: - class_members[add_constraint.constraint._meta.name] = ( - add_constraint.constraint - ) _Table: t.Type[Table] = create_table_class( class_name=add_table.class_name, @@ -886,6 +879,11 @@ async def _run_add_tables(self, backwards: bool = False): }, class_members=class_members, ) + + _Table._meta.constraints.extend( + self.add_constraints.for_table_class_name(add_table.class_name) + ) + table_classes.append(_Table) # Sort by foreign key, so they're created in the right order. diff --git a/piccolo/constraints.py b/piccolo/constraints.py index 617d5b48c..84c326aa7 100644 --- a/piccolo/constraints.py +++ b/piccolo/constraints.py @@ -4,8 +4,6 @@ from abc import abstractmethod from dataclasses import dataclass, field -from piccolo.columns import Column - class Constraint: """ @@ -13,6 +11,7 @@ class Constraint: """ def __init__(self, name: str, **kwargs) -> None: + kwargs.update(name=name) self._meta = ConstraintMeta(name=name, params=kwargs) def __hash__(self): @@ -43,36 +42,24 @@ class UniqueConstraint(Constraint): def __init__( self, - *columns: Column, - name: t.Optional[str] = None, + column_names: t.Sequence[str], + name: str, nulls_distinct: bool = True, ) -> None: """ :param columns: The table columns that should be unique together. :param name: - The name of the constraint in the database. If not provided, we - generate a sensible default. + The name of the constraint in the database. :param nulls_distinct: See the `Postgres docs `_ for more information. """ # noqa: E501 - if len(columns) < 1: + if len(column_names) < 1: raise ValueError("At least 1 column must be specified.") - tablenames = [column._meta.table._meta.tablename for column in columns] - - if len(set(tablenames)) > 1: - raise ValueError("The columns belong to different tables.") - - column_names = [column._meta.db_column_name for column in columns] self.column_names = column_names - - if name is None: - name = "_".join([tablenames[0], "unique", *column_names]) - self.name = name - self.nulls_distinct = nulls_distinct super().__init__( diff --git a/piccolo/table.py b/piccolo/table.py index 4c0523d29..7a0bbdbd4 100644 --- a/piccolo/table.py +++ b/piccolo/table.py @@ -28,7 +28,7 @@ ) from piccolo.columns.readable import Readable from piccolo.columns.reference import LAZY_COLUMN_REFERENCES -from piccolo.constraints import Constraint +from piccolo.constraints import Constraint, UniqueConstraint from piccolo.custom_types import TableInstance from piccolo.engine import Engine, engine_finder from piccolo.query import ( @@ -343,10 +343,6 @@ def __init_subclass__( attribute._meta._table = cls m2m_relationships.append(attribute) - if isinstance(attribute, Constraint): - attribute._meta._name = attribute_name - constraints.append(attribute) - if not primary_key: primary_key = cls._create_serial_primary_key() setattr(cls, "id", primary_key) @@ -483,9 +479,14 @@ def from_dict( return cls(**data) @classmethod - def add_constraints(cls, *constraint: Constraint): + def add_unique_constraint( + cls, + *columns: Column, + name: t.Optional[str] = None, + nulls_distinct: bool = True, + ): """ - Add composite constraints to the table (e.g. a unique constraint across + Add a unique constraint to the table (e.g. a unique constraint across multiple columns). You should wait for the ``Table`` to be initialised before calling @@ -497,12 +498,47 @@ class Album(Table): name = Varchar() band = ForeignKey(Band) - Album.add_constraints( - UniqueConstraint(Album.name, Album.band) + Album.add_unique_constraint( + Album.name, + Album.band, ) - """ - cls._meta.constraints.extend(constraint) + Note - this method doesn't create the constraint in the database. That + is done either by creating a migration, or using ``create_table``. + + :param columns: + The table columns that should be unique together. + :param name: + The name of the constraint in the database. If not provided, we + generate a sensible default. + :param nulls_distinct: + See the `Postgres docs `_ + for more information. + + """ # noqa: E501 + + if len(columns) < 1: + raise ValueError("At least 1 column must be specified.") + + for column in columns: + if column._meta.table != cls: + raise ValueError( + f"The {column._meta.name} column doesn't belong to this " + "table." + ) + + column_names = [column._meta.db_column_name for column in columns] + + if name is None: + name = "_".join([cls._meta.tablename, *column_names, "unique"]) + + cls._meta.constraints.append( + UniqueConstraint( + column_names=column_names, + name=name, + nulls_distinct=nulls_distinct, + ) + ) ########################################################################### diff --git a/tests/apps/migrations/auto/test_schema_differ.py b/tests/apps/migrations/auto/test_schema_differ.py index 92b68f59a..28127ec32 100644 --- a/tests/apps/migrations/auto/test_schema_differ.py +++ b/tests/apps/migrations/auto/test_schema_differ.py @@ -499,6 +499,7 @@ def test_add_table_with_constraint(self) -> None: genre_column = Varchar() genre_column._meta.name = "genre" + # TODO - this doesn't seem to be used, remove it? name_unique_constraint = UniqueConstraint(unique_columns=["name"]) name_unique_constraint._meta.name = "unique_name" From f62205ff37ef3bfcb46600f39bfdf1ee2788e3b5 Mon Sep 17 00:00:00 2001 From: Daniel Townsend Date: Sat, 25 Jan 2025 21:04:49 +0000 Subject: [PATCH 08/29] update docs --- docs/src/piccolo/schema/constraints.rst | 25 +++++++------------------ piccolo/table.py | 12 ++++++------ 2 files changed, 13 insertions(+), 24 deletions(-) diff --git a/docs/src/piccolo/schema/constraints.rst b/docs/src/piccolo/schema/constraints.rst index b7011148c..72d1dc836 100644 --- a/docs/src/piccolo/schema/constraints.rst +++ b/docs/src/piccolo/schema/constraints.rst @@ -2,8 +2,11 @@ Constraints =========== +Unique constraints +================== + Single column -============= +------------- Unique constraints can be added to a single column using the ``unique=True`` argument of ``Column``: @@ -13,29 +16,15 @@ argument of ``Column``: class Band(Table): name = Varchar(unique=True) -------------------------------------------------------------------------------- - Multi-column -============ - -``add_constraints`` -------------------- +------------ -Use the ``add_constraints`` method to add multi-column constraints to a +Use the ``add_unique_constraint`` method to add a multi-column constraint to a ``Table``: .. currentmodule:: piccolo.table -.. automethod:: Table.add_constraints +.. automethod:: Table.add_unique_constraint ------------------------------------------------------------------------------- -Constraint types -================ - -``UniqueConstraint`` --------------------- - -.. currentmodule:: piccolo.constraint - -.. autoclass:: UniqueConstraint diff --git a/piccolo/table.py b/piccolo/table.py index 7a0bbdbd4..0ba574630 100644 --- a/piccolo/table.py +++ b/piccolo/table.py @@ -486,14 +486,11 @@ def add_unique_constraint( nulls_distinct: bool = True, ): """ - Add a unique constraint to the table (e.g. a unique constraint across - multiple columns). + Add a unique constraint across multiple columns. You should wait for the ``Table`` to be initialised before calling this method. For example:: - from piccolo.constraints import UniqueConstraint - class Album(Table): name = Varchar() band = ForeignKey(Band) @@ -503,8 +500,11 @@ class Album(Table): Album.band, ) - Note - this method doesn't create the constraint in the database. That - is done either by creating a migration, or using ``create_table``. + .. note:: + This method doesn't create the constraint in the database - it just + registers it with the ``Table``. To create it in the database, + either create a migration, or use ``create_table`` if it's a new + table. :param columns: The table columns that should be unique together. From 4e1b8e56b690698d37f0483471b5bc7c32b485cb Mon Sep 17 00:00:00 2001 From: Daniel Townsend Date: Sat, 25 Jan 2025 22:16:22 +0000 Subject: [PATCH 09/29] add check constraints --- docs/src/piccolo/schema/constraints.rst | 4 ++ piccolo/constraints.py | 26 +++++++++ piccolo/table.py | 73 ++++++++++++++++++++++++- 3 files changed, 100 insertions(+), 3 deletions(-) diff --git a/docs/src/piccolo/schema/constraints.rst b/docs/src/piccolo/schema/constraints.rst index 72d1dc836..efbdd13b5 100644 --- a/docs/src/piccolo/schema/constraints.rst +++ b/docs/src/piccolo/schema/constraints.rst @@ -28,3 +28,7 @@ Use the ``add_unique_constraint`` method to add a multi-column constraint to a ------------------------------------------------------------------------------- +Check constraints +================= + +.. automethod:: Table.add_check_constraint diff --git a/piccolo/constraints.py b/piccolo/constraints.py index 84c326aa7..29398f21c 100644 --- a/piccolo/constraints.py +++ b/piccolo/constraints.py @@ -75,3 +75,29 @@ def ddl(self) -> str: ) columns_string = ", ".join(f'"{i}"' for i in self.column_names) return f"UNIQUE {nulls_string}({columns_string})" + + +class CheckConstraint(Constraint): + """ + Check constraint on the table columns. + """ + + def __init__( + self, + condition: str, + name: str, + ) -> None: + """ + :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. + + """ + self.condition = condition + super().__init__(name=name, condition=condition) + + @property + def ddl(self) -> str: + return f"CHECK {self.condition})" diff --git a/piccolo/table.py b/piccolo/table.py index 0ba574630..26d01fff7 100644 --- a/piccolo/table.py +++ b/piccolo/table.py @@ -28,8 +28,8 @@ ) from piccolo.columns.readable import Readable from piccolo.columns.reference import LAZY_COLUMN_REFERENCES -from piccolo.constraints import Constraint, UniqueConstraint -from piccolo.custom_types import TableInstance +from piccolo.constraints import CheckConstraint, Constraint, UniqueConstraint +from piccolo.custom_types import Combinable, TableInstance from piccolo.engine import Engine, engine_finder from piccolo.query import ( Alter, @@ -478,6 +478,9 @@ def from_dict( """ return cls(**data) + ########################################################################### + # Constraints + @classmethod def add_unique_constraint( cls, @@ -486,7 +489,7 @@ def add_unique_constraint( nulls_distinct: bool = True, ): """ - Add a unique constraint across multiple columns. + Add a unique constraint to one or more columns. You should wait for the ``Table`` to be initialised before calling this method. For example:: @@ -500,6 +503,9 @@ class Album(Table): Album.band, ) + In the above example, the database will enforce that ``name`` and + ``band`` form a unique combination. + .. note:: This method doesn't create the constraint in the database - it just registers it with the ``Table``. To create it in the database, @@ -540,6 +546,67 @@ class Album(Table): ) ) + @classmethod + def add_check_constraint( + cls, + condition: Combinable, + name: t.Optional[str] = None, + ): + """ + Add a check constraint to the table. + + You should wait for the ``Table`` to be initialised before calling + this method. For example:: + + class Ticket(Table): + price = Decimal() + + Ticket.add_check_constraint( + Ticket.price >= 0 + ) + + .. note:: + This method doesn't create the constraint in the database - it just + registers it with the ``Table``. To create it in the database, + either create a migration, or use ``create_table`` if it's a new + table. + + You can have more complex conditions, for example:: + + Ticket.add_check_constraint( + (Ticket.price >= 0) & (Ticket.price < 100) + ) + + :param condition: + The syntax is the same as the ``where`` clause used by most + queries (e.g. ``select``). + :param name: + The name of the constraint in the database. If not provided, we + generate a sensible default. + + """ # noqa: E501 + from piccolo.query.mixins import WhereDelegate + + columns = WhereDelegate(_where=condition).get_where_columns() + + for column in columns: + if column._meta.table != cls: + raise ValueError( + f"The {column._meta.name} column doesn't belong to this " + "table." + ) + + if name is None: + column_names = [column._meta.db_column_name for column in columns] + name = "_".join([cls._meta.tablename, *column_names, "check"]) + + cls._meta.constraints.append( + CheckConstraint( + condition=condition.querystring.__str__(), + name=name, + ) + ) + ########################################################################### def save( From e92aca20c782458084db99ce05080a9e7151ef67 Mon Sep 17 00:00:00 2001 From: Daniel Townsend Date: Sat, 25 Jan 2025 22:24:31 +0000 Subject: [PATCH 10/29] make sure `CheckConstraint` works in migrations --- piccolo/apps/migrations/auto/migration_manager.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/piccolo/apps/migrations/auto/migration_manager.py b/piccolo/apps/migrations/auto/migration_manager.py index 21a74c1f8..34ceed117 100644 --- a/piccolo/apps/migrations/auto/migration_manager.py +++ b/piccolo/apps/migrations/auto/migration_manager.py @@ -440,11 +440,7 @@ def add_constraint( params: t.Dict[str, t.Any], schema: t.Optional[str] = None, ): - if constraint_class is UniqueConstraint: - constraint = UniqueConstraint(**params) - else: - raise ValueError("Unrecognised constraint type") - + constraint = constraint_class(**params) constraint._meta.name = constraint_name self.add_constraints.append( @@ -881,7 +877,12 @@ async def _run_add_tables(self, backwards: bool = False): ) _Table._meta.constraints.extend( - self.add_constraints.for_table_class_name(add_table.class_name) + [ + i.constraint + for i in self.add_constraints.for_table_class_name( + add_table.class_name + ) + ] ) table_classes.append(_Table) From 3f5d6eae3ba1b3696f4f05cf218294022c0088a4 Mon Sep 17 00:00:00 2001 From: Daniel Townsend Date: Sat, 25 Jan 2025 22:28:13 +0000 Subject: [PATCH 11/29] fix typo in check condition ddl statement --- piccolo/constraints.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/piccolo/constraints.py b/piccolo/constraints.py index 29398f21c..6325527c2 100644 --- a/piccolo/constraints.py +++ b/piccolo/constraints.py @@ -100,4 +100,4 @@ def __init__( @property def ddl(self) -> str: - return f"CHECK {self.condition})" + return f"CHECK ({self.condition})" From acd341d70376c387cf370fcc9d251f08a8cdcd67 Mon Sep 17 00:00:00 2001 From: Daniel Townsend Date: Sat, 25 Jan 2025 22:31:17 +0000 Subject: [PATCH 12/29] remove unused import --- piccolo/apps/migrations/auto/migration_manager.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/piccolo/apps/migrations/auto/migration_manager.py b/piccolo/apps/migrations/auto/migration_manager.py index 34ceed117..7519e3395 100644 --- a/piccolo/apps/migrations/auto/migration_manager.py +++ b/piccolo/apps/migrations/auto/migration_manager.py @@ -17,7 +17,7 @@ from piccolo.apps.migrations.auto.serialisation import deserialise_params from piccolo.columns import Column, column_types from piccolo.columns.column_types import ForeignKey, Serial -from piccolo.constraints import Constraint, UniqueConstraint +from piccolo.constraints import Constraint from piccolo.engine import engine_finder from piccolo.query import Query from piccolo.query.base import DDL From 5e7e201d7286ccbd617cb10b0d6e95c299f53ada Mon Sep 17 00:00:00 2001 From: Daniel Townsend Date: Sat, 25 Jan 2025 22:31:22 +0000 Subject: [PATCH 13/29] improve docstring --- piccolo/constraints.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/piccolo/constraints.py b/piccolo/constraints.py index 6325527c2..0ded9679d 100644 --- a/piccolo/constraints.py +++ b/piccolo/constraints.py @@ -79,7 +79,7 @@ def ddl(self) -> str: class CheckConstraint(Constraint): """ - Check constraint on the table columns. + Check constraint on the table. """ def __init__( From fae32e43a1048d44c777ea598a1bd4002650a730 Mon Sep 17 00:00:00 2001 From: Daniel Townsend Date: Sun, 26 Jan 2025 20:55:06 +0000 Subject: [PATCH 14/29] refactor to use `constraints` list on `Table` --- piccolo/apps/playground/commands/run.py | 10 ++ piccolo/constraints.py | 151 +++++++++++++++++++++++- piccolo/table.py | 147 +++-------------------- 3 files changed, 175 insertions(+), 133 deletions(-) diff --git a/piccolo/apps/playground/commands/run.py b/piccolo/apps/playground/commands/run.py index 343c02e38..e815035f0 100644 --- a/piccolo/apps/playground/commands/run.py +++ b/piccolo/apps/playground/commands/run.py @@ -95,6 +95,9 @@ def get_readable(cls) -> Readable: ) +from piccolo.constraints import Unique + + class Ticket(Table): class TicketType(Enum): sitting = "sitting" @@ -106,6 +109,13 @@ class TicketType(Enum): price = Numeric(digits=(5, 2)) ticket_type = Varchar(choices=TicketType, default=TicketType.standing) + # These should just be dataclasses ... and perform no logic. + # In the metaclass we do all of the logic to convert them to proper + # Constraints + # I think using a dataclass makes more sense ... because we want to pass + # args. + constraints = [Unique(columns=[concert, ticket_type])] + @classmethod def get_readable(cls) -> Readable: return Readable( diff --git a/piccolo/constraints.py b/piccolo/constraints.py index 0ded9679d..401c57e20 100644 --- a/piccolo/constraints.py +++ b/piccolo/constraints.py @@ -1,11 +1,152 @@ from __future__ import annotations import typing as t -from abc import abstractmethod +from abc import ABCMeta, abstractmethod from dataclasses import dataclass, field +if t.TYPE_CHECKING: + from piccolo.columns import Column + from piccolo.custom_types import Combinable -class Constraint: + +class ConstraintConfig(metaclass=ABCMeta): + + def validate_columns(self, columns: t.List[Column]): + """ + Make sure all the columns belong to the same table. + """ + if len({column._meta.table._meta.tablename for column in columns}) > 1: + raise ValueError("The columns don't all belong to the same table.") + + @abstractmethod + def to_constraint(self) -> Constraint: + """ + Override in subclasses. + """ + raise NotImplementedError() + + +class Unique(ConstraintConfig): + """ + Add a unique constraint to one or more columns. For example:: + + class Album(Table): + name = Varchar() + band = ForeignKey(Band) + + constraints = [ + Unique([Album.name, Album.band]) + ] + + In the above example, the database will enforce that ``name`` and + ``band`` form a unique combination. + + :param columns: + The table columns that should be unique together. + :param name: + The name of the constraint in the database. If not provided, we + generate a sensible default. + :param nulls_distinct: + See the `Postgres docs `_ + for more information. + + """ # noqa: E501 + + def __init__( + self, + columns: t.List[Column], + name: t.Optional[str] = None, + nulls_distinct: bool = True, + ): + if len(columns) < 1: + raise ValueError("At least 1 column must be specified.") + + self.validate_columns(columns) + + self.columns = columns + self.name = name + self.nulls_distinct = nulls_distinct + + def to_constraint(self) -> UniqueConstraint: + """ + You should wait for the ``Table`` metaclass to assign names all of the + columns before calling this method. + """ + column_names = [column._meta.db_column_name for column in self.columns] + + if self.name is None: + tablename = self.columns[0]._meta.table._meta.tablename + name = "_".join([tablename, *column_names, "unique"]) + + return UniqueConstraint( + column_names=column_names, + name=name, + nulls_distinct=self.nulls_distinct, + ) + + +class Check(ConstraintConfig): + """ + Add a check constraint to the table. For example:: + + class Ticket(Table): + price = Decimal() + + constraints = [ + Check(Ticket.price >= 0) + ] + + You can have more complex conditions. For example:: + + class Ticket(Table): + price = Decimal() + + constraints = [ + Check( + (Ticket.price >= 0) & (Ticket.price < 100) + ) + ] + + :param condition: + The syntax is the same as the ``where`` clause used by most + queries (e.g. ``select``). + :param name: + The name of the constraint in the database. If not provided, we + generate a sensible default. + + """ + + def __init__(self, condition: Combinable, name: t.Optional[str] = None): + from piccolo.query.mixins import WhereDelegate + + self.columns = WhereDelegate(_where=condition).get_where_columns() + self.validate_columns(self.columns) + + self.condition = condition + self.name = name + + def to_constraint(self) -> CheckConstraint: + """ + You should wait for the ``Table`` metaclass to assign names all of the + columns before calling this method. + """ + if self.name is None: + tablename = self.columns[0]._meta.table._meta.tablename + column_names = [ + column._meta.db_column_name for column in self.columns + ] + name = "_".join([tablename, *column_names, "check"]) + + return CheckConstraint( + condition=self.condition.querystring.__str__(), + name=name, + ) + + +############################################################################### + + +class Constraint(metaclass=ABCMeta): """ All other constraints inherit from ``Constraint``. Don't use it directly. """ @@ -38,6 +179,9 @@ class ConstraintMeta: class UniqueConstraint(Constraint): """ Unique constraint on the table columns. + + This is the internal representation that Piccolo uses for constraints - + the user just supplies ``Unique``. """ def __init__( @@ -80,6 +224,9 @@ def ddl(self) -> str: class CheckConstraint(Constraint): """ Check constraint on the table. + + This is the internal representation that Piccolo uses for constraints - + the user just supplies ``Check``. """ def __init__( diff --git a/piccolo/table.py b/piccolo/table.py index 26d01fff7..fe036029a 100644 --- a/piccolo/table.py +++ b/piccolo/table.py @@ -28,8 +28,8 @@ ) from piccolo.columns.readable import Readable from piccolo.columns.reference import LAZY_COLUMN_REFERENCES -from piccolo.constraints import CheckConstraint, Constraint, UniqueConstraint -from piccolo.custom_types import Combinable, TableInstance +from piccolo.constraints import Check, Constraint, ConstraintConfig, Unique +from piccolo.custom_types import TableInstance from piccolo.engine import Engine, engine_finder from piccolo.query import ( Alter, @@ -350,6 +350,20 @@ def __init_subclass__( columns.insert(0, primary_key) # PK should be added first default_columns.append(primary_key) + # Now the columns are all setup, we can do the constraints. + if t.cast( + t.List[ConstraintConfig], + constraint_configs := getattr(cls, "constraints"), + ): + for constraint_config in constraint_configs: + if isinstance(constraint_config, ConstraintConfig): + constraints.append(constraint_config.to_constraint()) + else: + raise ValueError( + "The `constraints` list should only contain `Unique`" + " or `Check`." + ) + cls._meta = TableMeta( tablename=tablename, columns=columns, @@ -478,135 +492,6 @@ def from_dict( """ return cls(**data) - ########################################################################### - # Constraints - - @classmethod - def add_unique_constraint( - cls, - *columns: Column, - name: t.Optional[str] = None, - nulls_distinct: bool = True, - ): - """ - Add a unique constraint to one or more columns. - - You should wait for the ``Table`` to be initialised before calling - this method. For example:: - - class Album(Table): - name = Varchar() - band = ForeignKey(Band) - - Album.add_unique_constraint( - Album.name, - Album.band, - ) - - In the above example, the database will enforce that ``name`` and - ``band`` form a unique combination. - - .. note:: - This method doesn't create the constraint in the database - it just - registers it with the ``Table``. To create it in the database, - either create a migration, or use ``create_table`` if it's a new - table. - - :param columns: - The table columns that should be unique together. - :param name: - The name of the constraint in the database. If not provided, we - generate a sensible default. - :param nulls_distinct: - See the `Postgres docs `_ - for more information. - - """ # noqa: E501 - - if len(columns) < 1: - raise ValueError("At least 1 column must be specified.") - - for column in columns: - if column._meta.table != cls: - raise ValueError( - f"The {column._meta.name} column doesn't belong to this " - "table." - ) - - column_names = [column._meta.db_column_name for column in columns] - - if name is None: - name = "_".join([cls._meta.tablename, *column_names, "unique"]) - - cls._meta.constraints.append( - UniqueConstraint( - column_names=column_names, - name=name, - nulls_distinct=nulls_distinct, - ) - ) - - @classmethod - def add_check_constraint( - cls, - condition: Combinable, - name: t.Optional[str] = None, - ): - """ - Add a check constraint to the table. - - You should wait for the ``Table`` to be initialised before calling - this method. For example:: - - class Ticket(Table): - price = Decimal() - - Ticket.add_check_constraint( - Ticket.price >= 0 - ) - - .. note:: - This method doesn't create the constraint in the database - it just - registers it with the ``Table``. To create it in the database, - either create a migration, or use ``create_table`` if it's a new - table. - - You can have more complex conditions, for example:: - - Ticket.add_check_constraint( - (Ticket.price >= 0) & (Ticket.price < 100) - ) - - :param condition: - The syntax is the same as the ``where`` clause used by most - queries (e.g. ``select``). - :param name: - The name of the constraint in the database. If not provided, we - generate a sensible default. - - """ # noqa: E501 - from piccolo.query.mixins import WhereDelegate - - columns = WhereDelegate(_where=condition).get_where_columns() - - for column in columns: - if column._meta.table != cls: - raise ValueError( - f"The {column._meta.name} column doesn't belong to this " - "table." - ) - - if name is None: - column_names = [column._meta.db_column_name for column in columns] - name = "_".join([cls._meta.tablename, *column_names, "check"]) - - cls._meta.constraints.append( - CheckConstraint( - condition=condition.querystring.__str__(), - name=name, - ) - ) - ########################################################################### def save( From f369b813a4633300431b44f82ba231dd10b76917 Mon Sep 17 00:00:00 2001 From: Daniel Townsend Date: Sun, 26 Jan 2025 21:06:01 +0000 Subject: [PATCH 15/29] update the docs --- docs/src/piccolo/schema/constraints.rst | 39 +++++++++++++++---------- piccolo/constraints.py | 12 ++++++-- 2 files changed, 33 insertions(+), 18 deletions(-) diff --git a/docs/src/piccolo/schema/constraints.rst b/docs/src/piccolo/schema/constraints.rst index efbdd13b5..56f3c2afb 100644 --- a/docs/src/piccolo/schema/constraints.rst +++ b/docs/src/piccolo/schema/constraints.rst @@ -2,11 +2,8 @@ Constraints =========== -Unique constraints -================== - -Single column -------------- +Simple unique constraints +========================= Unique constraints can be added to a single column using the ``unique=True`` argument of ``Column``: @@ -16,19 +13,31 @@ argument of ``Column``: class Band(Table): name = Varchar(unique=True) -Multi-column ------------- +------------------------------------------------------------------------------- + +``Table.constraints`` +===================== -Use the ``add_unique_constraint`` method to add a multi-column constraint to a -``Table``: +By adding a ``constraints`` list to your ``Table``, you can implement powerful +``UNIQUE`` and ``CHECK`` constraints. -.. currentmodule:: piccolo.table +``Unique`` +---------- -.. automethod:: Table.add_unique_constraint +.. currentmodule:: piccolo.constraints -------------------------------------------------------------------------------- +.. autoclass:: Unique + +``Check`` +---------- + +.. autoclass:: Check + +How are they created? +--------------------- -Check constraints -================= +If creating a new table using ``await MyTable.create_table()``, then the +constraints will also be created. -.. automethod:: Table.add_check_constraint +Also, if using auto migrations, they handle the creation and deletion of these +constraints for you. diff --git a/piccolo/constraints.py b/piccolo/constraints.py index 401c57e20..0e5138903 100644 --- a/piccolo/constraints.py +++ b/piccolo/constraints.py @@ -30,12 +30,14 @@ class Unique(ConstraintConfig): """ Add a unique constraint to one or more columns. For example:: + from piccolo.constraints import Unique + class Album(Table): name = Varchar() band = ForeignKey(Band) constraints = [ - Unique([Album.name, Album.band]) + Unique([name, band]) ] In the above example, the database will enforce that ``name`` and @@ -89,21 +91,25 @@ class Check(ConstraintConfig): """ Add a check constraint to the table. For example:: + from piccolo.constraints import Check + class Ticket(Table): price = Decimal() constraints = [ - Check(Ticket.price >= 0) + Check(price >= 0) ] You can have more complex conditions. For example:: + from piccolo.constraints import Check + class Ticket(Table): price = Decimal() constraints = [ Check( - (Ticket.price >= 0) & (Ticket.price < 100) + (price >= 0) & (price < 100) ) ] From 0c6d0cf9c6d8c17aad080ad8b3d3ab50ad676c82 Mon Sep 17 00:00:00 2001 From: Daniel Townsend Date: Sun, 26 Jan 2025 21:15:43 +0000 Subject: [PATCH 16/29] improve logic for getting `Table.constraints` --- piccolo/table.py | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/piccolo/table.py b/piccolo/table.py index fe036029a..9e8a14cbc 100644 --- a/piccolo/table.py +++ b/piccolo/table.py @@ -351,18 +351,18 @@ def __init_subclass__( default_columns.append(primary_key) # Now the columns are all setup, we can do the constraints. - if t.cast( + constraint_configs = t.cast( t.List[ConstraintConfig], - constraint_configs := getattr(cls, "constraints"), - ): - for constraint_config in constraint_configs: - if isinstance(constraint_config, ConstraintConfig): - constraints.append(constraint_config.to_constraint()) - else: - raise ValueError( - "The `constraints` list should only contain `Unique`" - " or `Check`." - ) + getattr(cls, "constraints", []), + ) + for constraint_config in constraint_configs: + if isinstance(constraint_config, ConstraintConfig): + constraints.append(constraint_config.to_constraint()) + else: + raise ValueError( + "The `constraints` list should only contain `Unique`" + " or `Check`." + ) cls._meta = TableMeta( tablename=tablename, From 400046fc57b31241e158c493f9fe954ac6e15726 Mon Sep 17 00:00:00 2001 From: Daniel Townsend Date: Thu, 30 Jan 2025 11:28:47 +0000 Subject: [PATCH 17/29] try lambda --- piccolo/apps/playground/commands/run.py | 9 +---- piccolo/constraints.py | 52 +++++++++++-------------- piccolo/table.py | 34 ++++++++-------- 3 files changed, 41 insertions(+), 54 deletions(-) diff --git a/piccolo/apps/playground/commands/run.py b/piccolo/apps/playground/commands/run.py index e815035f0..250f261b9 100644 --- a/piccolo/apps/playground/commands/run.py +++ b/piccolo/apps/playground/commands/run.py @@ -24,6 +24,7 @@ Varchar, ) from piccolo.columns.readable import Readable +from piccolo.constraints import Unique from piccolo.engine import PostgresEngine, SQLiteEngine from piccolo.engine.base import Engine from piccolo.table import Table @@ -95,9 +96,6 @@ def get_readable(cls) -> Readable: ) -from piccolo.constraints import Unique - - class Ticket(Table): class TicketType(Enum): sitting = "sitting" @@ -109,11 +107,6 @@ class TicketType(Enum): price = Numeric(digits=(5, 2)) ticket_type = Varchar(choices=TicketType, default=TicketType.standing) - # These should just be dataclasses ... and perform no logic. - # In the metaclass we do all of the logic to convert them to proper - # Constraints - # I think using a dataclass makes more sense ... because we want to pass - # args. constraints = [Unique(columns=[concert, ticket_type])] @classmethod diff --git a/piccolo/constraints.py b/piccolo/constraints.py index 0e5138903..215fac3a9 100644 --- a/piccolo/constraints.py +++ b/piccolo/constraints.py @@ -11,15 +11,8 @@ class ConstraintConfig(metaclass=ABCMeta): - def validate_columns(self, columns: t.List[Column]): - """ - Make sure all the columns belong to the same table. - """ - if len({column._meta.table._meta.tablename for column in columns}) > 1: - raise ValueError("The columns don't all belong to the same table.") - @abstractmethod - def to_constraint(self) -> Constraint: + def to_constraint(self, tablename: str) -> Constraint: """ Override in subclasses. """ @@ -36,7 +29,7 @@ class Album(Table): name = Varchar() band = ForeignKey(Band) - constraints = [ + constraints = lambda: [ Unique([name, band]) ] @@ -56,29 +49,27 @@ class Album(Table): def __init__( self, - columns: t.List[Column], + columns: t.Callable[[], [t.List[Column]]], name: t.Optional[str] = None, nulls_distinct: bool = True, ): if len(columns) < 1: raise ValueError("At least 1 column must be specified.") - self.validate_columns(columns) - self.columns = columns self.name = name self.nulls_distinct = nulls_distinct - def to_constraint(self) -> UniqueConstraint: + def to_constraint(self, tablename: str) -> UniqueConstraint: """ You should wait for the ``Table`` metaclass to assign names all of the columns before calling this method. """ - column_names = [column._meta.db_column_name for column in self.columns] + column_names = [ + column._meta.db_column_name for column in self.columns() + ] - if self.name is None: - tablename = self.columns[0]._meta.table._meta.tablename - name = "_".join([tablename, *column_names, "unique"]) + name = self.name or "_".join([tablename, *column_names, "unique"]) return UniqueConstraint( column_names=column_names, @@ -122,29 +113,30 @@ class Ticket(Table): """ - def __init__(self, condition: Combinable, name: t.Optional[str] = None): - from piccolo.query.mixins import WhereDelegate - - self.columns = WhereDelegate(_where=condition).get_where_columns() - self.validate_columns(self.columns) - + def __init__( + self, + condition: t.Callable[[], [Combinable]], + name: t.Optional[str] = None, + ): self.condition = condition self.name = name - def to_constraint(self) -> CheckConstraint: + def to_constraint(self, tablename: str) -> CheckConstraint: """ You should wait for the ``Table`` metaclass to assign names all of the columns before calling this method. """ - if self.name is None: - tablename = self.columns[0]._meta.table._meta.tablename - column_names = [ - column._meta.db_column_name for column in self.columns - ] + name = self.name + + if name is None: + from piccolo.query.mixins import WhereDelegate + + columns = WhereDelegate(_where=self.condition).get_where_columns() + column_names = [column._meta.db_column_name for column in columns] name = "_".join([tablename, *column_names, "check"]) return CheckConstraint( - condition=self.condition.querystring.__str__(), + condition=self.condition().querystring.__str__(), name=name, ) diff --git a/piccolo/table.py b/piccolo/table.py index 9e8a14cbc..9bdd4c0c2 100644 --- a/piccolo/table.py +++ b/piccolo/table.py @@ -290,7 +290,6 @@ def __init_subclass__( auto_update_columns: t.List[Column] = [] primary_key: t.Optional[Column] = None m2m_relationships: t.List[M2M] = [] - constraints: t.List[Constraint] = [] attribute_names = itertools.chain( *[i.__dict__.keys() for i in reversed(cls.__mro__)] @@ -350,20 +349,6 @@ def __init_subclass__( columns.insert(0, primary_key) # PK should be added first default_columns.append(primary_key) - # Now the columns are all setup, we can do the constraints. - constraint_configs = t.cast( - t.List[ConstraintConfig], - getattr(cls, "constraints", []), - ) - for constraint_config in constraint_configs: - if isinstance(constraint_config, ConstraintConfig): - constraints.append(constraint_config.to_constraint()) - else: - raise ValueError( - "The `constraints` list should only contain `Unique`" - " or `Check`." - ) - cls._meta = TableMeta( tablename=tablename, columns=columns, @@ -380,7 +365,6 @@ def __init_subclass__( help_text=help_text, _db=db, m2m_relationships=m2m_relationships, - constraints=constraints, schema=schema, ) @@ -395,6 +379,24 @@ def __init_subclass__( foreign_key_column ) + # Now the table and columns are all setup, we can do the constraints. + constraints: t.List[Constraint] = [] + constraint_configs = t.cast( + t.List[ConstraintConfig], + getattr(cls, "constraints", []), + ) + for constraint_config in constraint_configs: + if isinstance(constraint_config, ConstraintConfig): + constraints.append( + constraint_config.to_constraint(tablename=tablename) + ) + else: + raise ValueError( + "The `constraints` list should only contain `Unique`" + " or `Check`." + ) + cls._meta.constraints = constraints + TABLE_REGISTRY.append(cls) def __init__( From 7bc9a2e745a8f0e8d467fcb1b849099f9aa825d1 Mon Sep 17 00:00:00 2001 From: Daniel Townsend Date: Fri, 31 Jan 2025 17:09:55 +0000 Subject: [PATCH 18/29] strings --- piccolo/constraints.py | 207 ++++++++++++----------------------------- piccolo/table.py | 22 +---- 2 files changed, 62 insertions(+), 167 deletions(-) diff --git a/piccolo/constraints.py b/piccolo/constraints.py index 215fac3a9..d02cf0c7b 100644 --- a/piccolo/constraints.py +++ b/piccolo/constraints.py @@ -4,145 +4,6 @@ from abc import ABCMeta, abstractmethod from dataclasses import dataclass, field -if t.TYPE_CHECKING: - from piccolo.columns import Column - from piccolo.custom_types import Combinable - - -class ConstraintConfig(metaclass=ABCMeta): - - @abstractmethod - def to_constraint(self, tablename: str) -> Constraint: - """ - Override in subclasses. - """ - raise NotImplementedError() - - -class Unique(ConstraintConfig): - """ - Add a unique constraint to one or more columns. For example:: - - from piccolo.constraints import Unique - - class Album(Table): - name = Varchar() - band = ForeignKey(Band) - - constraints = lambda: [ - Unique([name, band]) - ] - - In the above example, the database will enforce that ``name`` and - ``band`` form a unique combination. - - :param columns: - The table columns that should be unique together. - :param name: - The name of the constraint in the database. If not provided, we - generate a sensible default. - :param nulls_distinct: - See the `Postgres docs `_ - for more information. - - """ # noqa: E501 - - def __init__( - self, - columns: t.Callable[[], [t.List[Column]]], - name: t.Optional[str] = None, - nulls_distinct: bool = True, - ): - if len(columns) < 1: - raise ValueError("At least 1 column must be specified.") - - self.columns = columns - self.name = name - self.nulls_distinct = nulls_distinct - - def to_constraint(self, tablename: str) -> UniqueConstraint: - """ - You should wait for the ``Table`` metaclass to assign names all of the - columns before calling this method. - """ - column_names = [ - column._meta.db_column_name for column in self.columns() - ] - - name = self.name or "_".join([tablename, *column_names, "unique"]) - - return UniqueConstraint( - column_names=column_names, - name=name, - nulls_distinct=self.nulls_distinct, - ) - - -class Check(ConstraintConfig): - """ - Add a check constraint to the table. For example:: - - from piccolo.constraints import Check - - class Ticket(Table): - price = Decimal() - - constraints = [ - Check(price >= 0) - ] - - You can have more complex conditions. For example:: - - from piccolo.constraints import Check - - class Ticket(Table): - price = Decimal() - - constraints = [ - Check( - (price >= 0) & (price < 100) - ) - ] - - :param condition: - The syntax is the same as the ``where`` clause used by most - queries (e.g. ``select``). - :param name: - The name of the constraint in the database. If not provided, we - generate a sensible default. - - """ - - def __init__( - self, - condition: t.Callable[[], [Combinable]], - name: t.Optional[str] = None, - ): - self.condition = condition - self.name = name - - def to_constraint(self, tablename: str) -> CheckConstraint: - """ - You should wait for the ``Table`` metaclass to assign names all of the - columns before calling this method. - """ - name = self.name - - if name is None: - from piccolo.query.mixins import WhereDelegate - - columns = WhereDelegate(_where=self.condition).get_where_columns() - column_names = [column._meta.db_column_name for column in columns] - name = "_".join([tablename, *column_names, "check"]) - - return CheckConstraint( - condition=self.condition().querystring.__str__(), - name=name, - ) - - -############################################################################### - class Constraint(metaclass=ABCMeta): """ @@ -174,13 +35,34 @@ class ConstraintMeta: params: t.Dict[str, t.Any] = field(default_factory=dict) -class UniqueConstraint(Constraint): +class Unique(Constraint): """ - Unique constraint on the table columns. + Add a unique constraint to one or more columns. For example:: - This is the internal representation that Piccolo uses for constraints - - the user just supplies ``Unique``. - """ + from piccolo.constraints import Unique + + class Album( + Table, + constraints=[ + Unique(['name', 'band']) + ] + ): + name = Varchar() + band = ForeignKey(Band) + + In the above example, the database will enforce that ``name`` and + ``band`` form a unique combination. + + :param columns: + The table columns that should be unique together. + :param name: + The name of the constraint in the database. If not provided, we + generate a sensible default. + :param nulls_distinct: + See the `Postgres docs `_ + for more information. + + """ # noqa: E501 def __init__( self, @@ -219,12 +101,41 @@ def ddl(self) -> str: return f"UNIQUE {nulls_string}({columns_string})" -class CheckConstraint(Constraint): +class Check(Constraint): """ - Check constraint on the table. + Add a check constraint to the table. For example:: + + from piccolo.constraints import Check + + class Ticket( + Table, + constraints=[ + Check('price >= 0') + ] + ): + price = Decimal() + + You can have more complex conditions. For example:: + + from piccolo.constraints import Check + + class Ticket( + Table, + constraints=[ + Check( + "price >= 0 AND price < 100" + ) + ] + ): + price = Decimal() + + :param condition: + The syntax is the same as the ``where`` clause used by most + queries (e.g. ``select``). + :param name: + The name of the constraint in the database. If not provided, we + generate a sensible default. - This is the internal representation that Piccolo uses for constraints - - the user just supplies ``Check``. """ def __init__( diff --git a/piccolo/table.py b/piccolo/table.py index 9bdd4c0c2..8155e0239 100644 --- a/piccolo/table.py +++ b/piccolo/table.py @@ -28,7 +28,7 @@ ) from piccolo.columns.readable import Readable from piccolo.columns.reference import LAZY_COLUMN_REFERENCES -from piccolo.constraints import Check, Constraint, ConstraintConfig, Unique +from piccolo.constraints import Constraint from piccolo.custom_types import TableInstance from piccolo.engine import Engine, engine_finder from piccolo.query import ( @@ -243,6 +243,7 @@ def __init_subclass__( tags: t.Optional[t.List[str]] = None, help_text: t.Optional[str] = None, schema: t.Optional[str] = None, + constraints: t.Optional[t.List[Constraint]] = None, ): # sourcery no-metrics """ Automatically populate the _meta, which includes the tablename, and @@ -366,6 +367,7 @@ def __init_subclass__( _db=db, m2m_relationships=m2m_relationships, schema=schema, + constraints=constraints or [], ) for foreign_key_column in foreign_key_columns: @@ -379,24 +381,6 @@ def __init_subclass__( foreign_key_column ) - # Now the table and columns are all setup, we can do the constraints. - constraints: t.List[Constraint] = [] - constraint_configs = t.cast( - t.List[ConstraintConfig], - getattr(cls, "constraints", []), - ) - for constraint_config in constraint_configs: - if isinstance(constraint_config, ConstraintConfig): - constraints.append( - constraint_config.to_constraint(tablename=tablename) - ) - else: - raise ValueError( - "The `constraints` list should only contain `Unique`" - " or `Check`." - ) - cls._meta.constraints = constraints - TABLE_REGISTRY.append(cls) def __init__( From 388e89f5b8935ff43ed35d2d5a45b16389bacda0 Mon Sep 17 00:00:00 2001 From: Daniel Townsend Date: Fri, 31 Jan 2025 17:10:02 +0000 Subject: [PATCH 19/29] Revert "strings" This reverts commit 7bc9a2e745a8f0e8d467fcb1b849099f9aa825d1. --- piccolo/constraints.py | 207 +++++++++++++++++++++++++++++------------ piccolo/table.py | 22 ++++- 2 files changed, 167 insertions(+), 62 deletions(-) diff --git a/piccolo/constraints.py b/piccolo/constraints.py index d02cf0c7b..215fac3a9 100644 --- a/piccolo/constraints.py +++ b/piccolo/constraints.py @@ -4,6 +4,145 @@ from abc import ABCMeta, abstractmethod from dataclasses import dataclass, field +if t.TYPE_CHECKING: + from piccolo.columns import Column + from piccolo.custom_types import Combinable + + +class ConstraintConfig(metaclass=ABCMeta): + + @abstractmethod + def to_constraint(self, tablename: str) -> Constraint: + """ + Override in subclasses. + """ + raise NotImplementedError() + + +class Unique(ConstraintConfig): + """ + Add a unique constraint to one or more columns. For example:: + + from piccolo.constraints import Unique + + class Album(Table): + name = Varchar() + band = ForeignKey(Band) + + constraints = lambda: [ + Unique([name, band]) + ] + + In the above example, the database will enforce that ``name`` and + ``band`` form a unique combination. + + :param columns: + The table columns that should be unique together. + :param name: + The name of the constraint in the database. If not provided, we + generate a sensible default. + :param nulls_distinct: + See the `Postgres docs `_ + for more information. + + """ # noqa: E501 + + def __init__( + self, + columns: t.Callable[[], [t.List[Column]]], + name: t.Optional[str] = None, + nulls_distinct: bool = True, + ): + if len(columns) < 1: + raise ValueError("At least 1 column must be specified.") + + self.columns = columns + self.name = name + self.nulls_distinct = nulls_distinct + + def to_constraint(self, tablename: str) -> UniqueConstraint: + """ + You should wait for the ``Table`` metaclass to assign names all of the + columns before calling this method. + """ + column_names = [ + column._meta.db_column_name for column in self.columns() + ] + + name = self.name or "_".join([tablename, *column_names, "unique"]) + + return UniqueConstraint( + column_names=column_names, + name=name, + nulls_distinct=self.nulls_distinct, + ) + + +class Check(ConstraintConfig): + """ + Add a check constraint to the table. For example:: + + from piccolo.constraints import Check + + class Ticket(Table): + price = Decimal() + + constraints = [ + Check(price >= 0) + ] + + You can have more complex conditions. For example:: + + from piccolo.constraints import Check + + class Ticket(Table): + price = Decimal() + + constraints = [ + Check( + (price >= 0) & (price < 100) + ) + ] + + :param condition: + The syntax is the same as the ``where`` clause used by most + queries (e.g. ``select``). + :param name: + The name of the constraint in the database. If not provided, we + generate a sensible default. + + """ + + def __init__( + self, + condition: t.Callable[[], [Combinable]], + name: t.Optional[str] = None, + ): + self.condition = condition + self.name = name + + def to_constraint(self, tablename: str) -> CheckConstraint: + """ + You should wait for the ``Table`` metaclass to assign names all of the + columns before calling this method. + """ + name = self.name + + if name is None: + from piccolo.query.mixins import WhereDelegate + + columns = WhereDelegate(_where=self.condition).get_where_columns() + column_names = [column._meta.db_column_name for column in columns] + name = "_".join([tablename, *column_names, "check"]) + + return CheckConstraint( + condition=self.condition().querystring.__str__(), + name=name, + ) + + +############################################################################### + class Constraint(metaclass=ABCMeta): """ @@ -35,34 +174,13 @@ class ConstraintMeta: params: t.Dict[str, t.Any] = field(default_factory=dict) -class Unique(Constraint): +class UniqueConstraint(Constraint): """ - Add a unique constraint to one or more columns. For example:: - - from piccolo.constraints import Unique - - class Album( - Table, - constraints=[ - Unique(['name', 'band']) - ] - ): - name = Varchar() - band = ForeignKey(Band) - - In the above example, the database will enforce that ``name`` and - ``band`` form a unique combination. + Unique constraint on the table columns. - :param columns: - The table columns that should be unique together. - :param name: - The name of the constraint in the database. If not provided, we - generate a sensible default. - :param nulls_distinct: - See the `Postgres docs `_ - for more information. - - """ # noqa: E501 + This is the internal representation that Piccolo uses for constraints - + the user just supplies ``Unique``. + """ def __init__( self, @@ -101,41 +219,12 @@ def ddl(self) -> str: return f"UNIQUE {nulls_string}({columns_string})" -class Check(Constraint): +class CheckConstraint(Constraint): """ - Add a check constraint to the table. For example:: - - from piccolo.constraints import Check - - class Ticket( - Table, - constraints=[ - Check('price >= 0') - ] - ): - price = Decimal() - - You can have more complex conditions. For example:: - - from piccolo.constraints import Check - - class Ticket( - Table, - constraints=[ - Check( - "price >= 0 AND price < 100" - ) - ] - ): - price = Decimal() - - :param condition: - The syntax is the same as the ``where`` clause used by most - queries (e.g. ``select``). - :param name: - The name of the constraint in the database. If not provided, we - generate a sensible default. + Check constraint on the table. + This is the internal representation that Piccolo uses for constraints - + the user just supplies ``Check``. """ def __init__( diff --git a/piccolo/table.py b/piccolo/table.py index 8155e0239..9bdd4c0c2 100644 --- a/piccolo/table.py +++ b/piccolo/table.py @@ -28,7 +28,7 @@ ) from piccolo.columns.readable import Readable from piccolo.columns.reference import LAZY_COLUMN_REFERENCES -from piccolo.constraints import Constraint +from piccolo.constraints import Check, Constraint, ConstraintConfig, Unique from piccolo.custom_types import TableInstance from piccolo.engine import Engine, engine_finder from piccolo.query import ( @@ -243,7 +243,6 @@ def __init_subclass__( tags: t.Optional[t.List[str]] = None, help_text: t.Optional[str] = None, schema: t.Optional[str] = None, - constraints: t.Optional[t.List[Constraint]] = None, ): # sourcery no-metrics """ Automatically populate the _meta, which includes the tablename, and @@ -367,7 +366,6 @@ def __init_subclass__( _db=db, m2m_relationships=m2m_relationships, schema=schema, - constraints=constraints or [], ) for foreign_key_column in foreign_key_columns: @@ -381,6 +379,24 @@ def __init_subclass__( foreign_key_column ) + # Now the table and columns are all setup, we can do the constraints. + constraints: t.List[Constraint] = [] + constraint_configs = t.cast( + t.List[ConstraintConfig], + getattr(cls, "constraints", []), + ) + for constraint_config in constraint_configs: + if isinstance(constraint_config, ConstraintConfig): + constraints.append( + constraint_config.to_constraint(tablename=tablename) + ) + else: + raise ValueError( + "The `constraints` list should only contain `Unique`" + " or `Check`." + ) + cls._meta.constraints = constraints + TABLE_REGISTRY.append(cls) def __init__( From 83dad231fa3fb8e2cd8e960bdc9e9dc047f658be Mon Sep 17 00:00:00 2001 From: Daniel Townsend Date: Fri, 31 Jan 2025 21:23:36 +0000 Subject: [PATCH 20/29] added `querystring_for_constraint` --- piccolo/columns/combination.py | 29 ++++++++++++++++++++++ piccolo/constraints.py | 45 +++++++++++++++++++++++++--------- piccolo/table.py | 6 +++-- 3 files changed, 67 insertions(+), 13 deletions(-) diff --git a/piccolo/columns/combination.py b/piccolo/columns/combination.py index e080cced2..97876e5c3 100644 --- a/piccolo/columns/combination.py +++ b/piccolo/columns/combination.py @@ -52,6 +52,14 @@ def querystring_for_update(self) -> QueryString: self.second.querystring_for_update, ) + @property + def querystring_for_constraint(self) -> QueryString: + return QueryString( + "({} " + self.operator + " {})", + self.first.querystring_for_constraint, + self.second.querystring_for_constraint, + ) + def __str__(self): return self.querystring.__str__() @@ -134,6 +142,10 @@ def __init__(self, sql: str, *args: t.Any) -> None: def querystring_for_update(self) -> QueryString: return self.querystring + @property + def querystring_for_constraint(self) -> QueryString: + return self.querystring + def __str__(self): return self.querystring.__str__() @@ -249,5 +261,22 @@ def querystring_for_update(self) -> QueryString: return QueryString(template, *args) + @property + def querystring_for_constraint(self) -> QueryString: + args: t.List[t.Any] = [] + if self.value != UNDEFINED: + args.append(self.value) + + if self.values != UNDEFINED: + args.append(self.values_querystring) + + template = self.operator.template.format( + name=f'"{self.column._meta.db_column_name}"', + value="{}", + values="{}", + ) + + return QueryString(template, *args) + def __str__(self): return self.querystring.__str__() diff --git a/piccolo/constraints.py b/piccolo/constraints.py index 215fac3a9..88d86eeed 100644 --- a/piccolo/constraints.py +++ b/piccolo/constraints.py @@ -29,7 +29,7 @@ class Album(Table): name = Varchar() band = ForeignKey(Band) - constraints = lambda: [ + constraints = [ Unique([name, band]) ] @@ -49,7 +49,7 @@ class Album(Table): def __init__( self, - columns: t.Callable[[], [t.List[Column]]], + columns: t.List[t.Union[Column, str]], name: t.Optional[str] = None, nulls_distinct: bool = True, ): @@ -65,11 +65,18 @@ def to_constraint(self, tablename: str) -> UniqueConstraint: You should wait for the ``Table`` metaclass to assign names all of the columns before calling this method. """ + from piccolo.columns import Column + column_names = [ - column._meta.db_column_name for column in self.columns() + ( + column._meta.db_column_name + if isinstance(column, Column) + else column + ) + for column in self.columns ] - name = self.name or "_".join([tablename, *column_names, "unique"]) + name = self.name or "_".join(["unique", tablename, *column_names]) return UniqueConstraint( column_names=column_names, @@ -115,7 +122,7 @@ class Ticket(Table): def __init__( self, - condition: t.Callable[[], [Combinable]], + condition: t.Union[Combinable, str], name: t.Optional[str] = None, ): self.condition = condition @@ -126,17 +133,33 @@ def to_constraint(self, tablename: str) -> CheckConstraint: You should wait for the ``Table`` metaclass to assign names all of the columns before calling this method. """ + from piccolo.columns.combination import CombinableMixin + name = self.name if name is None: - from piccolo.query.mixins import WhereDelegate - - columns = WhereDelegate(_where=self.condition).get_where_columns() - column_names = [column._meta.db_column_name for column in columns] - name = "_".join([tablename, *column_names, "check"]) + if isinstance(self.condition, str): + name = "_".join( + ["check", tablename, str(hash(self.condition))] + ) + else: + from piccolo.query.mixins import WhereDelegate + + columns = WhereDelegate( + _where=self.condition + ).get_where_columns() + column_names = [ + column._meta.db_column_name for column in columns + ] + name = "_".join(["check", tablename, *column_names]) + + if isinstance(self.condition, CombinableMixin): + condition_str = self.condition.querystring_for_constraint.__str__() + else: + condition_str = self.condition return CheckConstraint( - condition=self.condition().querystring.__str__(), + condition=condition_str, name=name, ) diff --git a/piccolo/table.py b/piccolo/table.py index 9bdd4c0c2..c96e11217 100644 --- a/piccolo/table.py +++ b/piccolo/table.py @@ -302,11 +302,14 @@ def __init_subclass__( attribute = getattr(cls, attribute_name) if isinstance(attribute, Column): + column = attribute + column._meta._name = attribute_name + # We have to copy, then override the existing column # definition, in case this column is inheritted from a mixin. # Otherwise, when we set attributes on that column, it will # effect all other users of that mixin. - column = attribute.copy() + column = column.copy() setattr(cls, attribute_name, column) if column._meta.primary_key: @@ -315,7 +318,6 @@ def __init_subclass__( non_default_columns.append(column) columns.append(column) - column._meta._name = attribute_name column._meta._table = cls if isinstance(column, Array): From dc561046b7b008c33d4c5eeb435e8dcf74ffee5a Mon Sep 17 00:00:00 2001 From: Daniel Townsend Date: Sat, 1 Feb 2025 10:06:10 +0000 Subject: [PATCH 21/29] stop using `constraints` list This doesn't work well with mixins and inheritance --- piccolo/apps/playground/commands/run.py | 2 +- piccolo/constraints.py | 66 ++++++++++--------------- piccolo/table.py | 32 ++++++------ 3 files changed, 42 insertions(+), 58 deletions(-) diff --git a/piccolo/apps/playground/commands/run.py b/piccolo/apps/playground/commands/run.py index 250f261b9..5b82bc3d0 100644 --- a/piccolo/apps/playground/commands/run.py +++ b/piccolo/apps/playground/commands/run.py @@ -107,7 +107,7 @@ class TicketType(Enum): price = Numeric(digits=(5, 2)) ticket_type = Varchar(choices=TicketType, default=TicketType.standing) - constraints = [Unique(columns=[concert, ticket_type])] + unique_concert_ticket_type = Unique(columns=[concert, ticket_type]) @classmethod def get_readable(cls) -> Readable: diff --git a/piccolo/constraints.py b/piccolo/constraints.py index 88d86eeed..f33424afb 100644 --- a/piccolo/constraints.py +++ b/piccolo/constraints.py @@ -11,8 +11,10 @@ class ConstraintConfig(metaclass=ABCMeta): + name: str + @abstractmethod - def to_constraint(self, tablename: str) -> Constraint: + def to_constraint(self) -> Constraint: """ Override in subclasses. """ @@ -29,9 +31,8 @@ class Album(Table): name = Varchar() band = ForeignKey(Band) - constraints = [ - Unique([name, band]) - ] + unique_name_band = Unique([name, band]) + In the above example, the database will enforce that ``name`` and ``band`` form a unique combination. @@ -50,17 +51,15 @@ class Album(Table): def __init__( self, columns: t.List[t.Union[Column, str]], - name: t.Optional[str] = None, nulls_distinct: bool = True, ): if len(columns) < 1: raise ValueError("At least 1 column must be specified.") self.columns = columns - self.name = name self.nulls_distinct = nulls_distinct - def to_constraint(self, tablename: str) -> UniqueConstraint: + def to_constraint(self) -> UniqueConstraint: """ You should wait for the ``Table`` metaclass to assign names all of the columns before calling this method. @@ -76,11 +75,9 @@ def to_constraint(self, tablename: str) -> UniqueConstraint: for column in self.columns ] - name = self.name or "_".join(["unique", tablename, *column_names]) - return UniqueConstraint( column_names=column_names, - name=name, + name=self.name, nulls_distinct=self.nulls_distinct, ) @@ -94,9 +91,7 @@ class Check(ConstraintConfig): class Ticket(Table): price = Decimal() - constraints = [ - Check(price >= 0) - ] + check_price_positive = Check(price >= 0) You can have more complex conditions. For example:: @@ -105,11 +100,9 @@ class Ticket(Table): class Ticket(Table): price = Decimal() - constraints = [ - Check( - (price >= 0) & (price < 100) - ) - ] + check_price_range = Check( + (price >= 0) & (price < 100) + ) :param condition: The syntax is the same as the ``where`` clause used by most @@ -123,36 +116,16 @@ class Ticket(Table): def __init__( self, condition: t.Union[Combinable, str], - name: t.Optional[str] = None, ): self.condition = condition - self.name = name - def to_constraint(self, tablename: str) -> CheckConstraint: + def to_constraint(self) -> CheckConstraint: """ You should wait for the ``Table`` metaclass to assign names all of the columns before calling this method. """ from piccolo.columns.combination import CombinableMixin - name = self.name - - if name is None: - if isinstance(self.condition, str): - name = "_".join( - ["check", tablename, str(hash(self.condition))] - ) - else: - from piccolo.query.mixins import WhereDelegate - - columns = WhereDelegate( - _where=self.condition - ).get_where_columns() - column_names = [ - column._meta.db_column_name for column in columns - ] - name = "_".join(["check", tablename, *column_names]) - if isinstance(self.condition, CombinableMixin): condition_str = self.condition.querystring_for_constraint.__str__() else: @@ -160,7 +133,7 @@ def to_constraint(self, tablename: str) -> CheckConstraint: return CheckConstraint( condition=condition_str, - name=name, + name=self.name, ) @@ -184,6 +157,10 @@ def __hash__(self): def ddl(self) -> str: raise NotImplementedError + @abstractmethod + def _table_str(self) -> str: + raise NotImplementedError + @dataclass class ConstraintMeta: @@ -241,6 +218,12 @@ def ddl(self) -> str: columns_string = ", ".join(f'"{i}"' for i in self.column_names) return f"UNIQUE {nulls_string}({columns_string})" + def _table_str(self) -> str: + columns_string = ", ".join( + [f'"{column_name}"' for column_name in self.column_names] + ) + return f'{self._meta.name} = Unique([{columns_string}])")' + class CheckConstraint(Constraint): """ @@ -269,3 +252,6 @@ def __init__( @property def ddl(self) -> str: return f"CHECK ({self.condition})" + + def _table_str(self) -> str: + return f'{self._meta.name} = Check("{self.condition}"))' diff --git a/piccolo/table.py b/piccolo/table.py index c96e11217..08db66fe8 100644 --- a/piccolo/table.py +++ b/piccolo/table.py @@ -290,6 +290,7 @@ def __init_subclass__( auto_update_columns: t.List[Column] = [] primary_key: t.Optional[Column] = None m2m_relationships: t.List[M2M] = [] + constraint_configs: t.List[ConstraintConfig] = [] attribute_names = itertools.chain( *[i.__dict__.keys() for i in reversed(cls.__mro__)] @@ -303,7 +304,7 @@ def __init_subclass__( attribute = getattr(cls, attribute_name) if isinstance(attribute, Column): column = attribute - column._meta._name = attribute_name + column._meta._name = attribute_name # We have to copy, then override the existing column # definition, in case this column is inheritted from a mixin. @@ -344,6 +345,10 @@ def __init_subclass__( attribute._meta._table = cls m2m_relationships.append(attribute) + if isinstance(attribute, ConstraintConfig): + attribute.name = attribute_name + constraint_configs.append(attribute) + if not primary_key: primary_key = cls._create_serial_primary_key() setattr(cls, "id", primary_key) @@ -383,20 +388,8 @@ def __init_subclass__( # Now the table and columns are all setup, we can do the constraints. constraints: t.List[Constraint] = [] - constraint_configs = t.cast( - t.List[ConstraintConfig], - getattr(cls, "constraints", []), - ) for constraint_config in constraint_configs: - if isinstance(constraint_config, ConstraintConfig): - constraints.append( - constraint_config.to_constraint(tablename=tablename) - ) - else: - raise ValueError( - "The `constraints` list should only contain `Unique`" - " or `Check`." - ) + constraints.append(constraint_config.to_constraint()) cls._meta.constraints = constraints TABLE_REGISTRY.append(cls) @@ -1384,7 +1377,7 @@ def _table_str( if excluded_params is None: excluded_params = [] spacer = "\n " - columns = [] + column_strings: t.List[str] = [] for col in cls._meta.columns: params: t.List[str] = [] for key, value in col._meta.params.items(): @@ -1400,10 +1393,15 @@ def _table_str( if not abbreviated: params.append(f"{key}={_value}") params_string = ", ".join(params) - columns.append( + column_strings.append( f"{col._meta.name} = {col.__class__.__name__}({params_string})" ) - columns_string = spacer.join(columns) + columns_string = spacer.join(column_strings) + + constraint_strings: t.List[str] = [] + for constraint in cls._meta.constraints: + constraint_strings.append(constraint._table_str()) + tablename = repr(cls._meta.tablename) parent_class_name = cls.mro()[1].__name__ From b3039d5e3ab2ac3437c1d27ed1c907aa398fa041 Mon Sep 17 00:00:00 2001 From: Daniel Townsend Date: Sat, 1 Feb 2025 17:21:14 +0000 Subject: [PATCH 22/29] update `_table_str` --- piccolo/apps/playground/commands/run.py | 3 ++- piccolo/constraints.py | 7 +++++-- piccolo/table.py | 5 ++++- 3 files changed, 11 insertions(+), 4 deletions(-) diff --git a/piccolo/apps/playground/commands/run.py b/piccolo/apps/playground/commands/run.py index 5b82bc3d0..80c0f7f33 100644 --- a/piccolo/apps/playground/commands/run.py +++ b/piccolo/apps/playground/commands/run.py @@ -24,7 +24,7 @@ Varchar, ) from piccolo.columns.readable import Readable -from piccolo.constraints import Unique +from piccolo.constraints import Check, Unique from piccolo.engine import PostgresEngine, SQLiteEngine from piccolo.engine.base import Engine from piccolo.table import Table @@ -108,6 +108,7 @@ class TicketType(Enum): ticket_type = Varchar(choices=TicketType, default=TicketType.standing) unique_concert_ticket_type = Unique(columns=[concert, ticket_type]) + check_price = Check((price > 0) & (price < 200)) @classmethod def get_readable(cls) -> Readable: diff --git a/piccolo/constraints.py b/piccolo/constraints.py index f33424afb..0c21efd26 100644 --- a/piccolo/constraints.py +++ b/piccolo/constraints.py @@ -222,7 +222,10 @@ def _table_str(self) -> str: columns_string = ", ".join( [f'"{column_name}"' for column_name in self.column_names] ) - return f'{self._meta.name} = Unique([{columns_string}])")' + return ( + f"{self._meta.name} = Unique([{columns_string}], " + f"nulls_distinct={self.nulls_distinct})" + ) class CheckConstraint(Constraint): @@ -254,4 +257,4 @@ def ddl(self) -> str: return f"CHECK ({self.condition})" def _table_str(self) -> str: - return f'{self._meta.name} = Check("{self.condition}"))' + return f'{self._meta.name} = Check("{self.condition}")' diff --git a/piccolo/table.py b/piccolo/table.py index 08db66fe8..a8987c782 100644 --- a/piccolo/table.py +++ b/piccolo/table.py @@ -1401,6 +1401,7 @@ def _table_str( constraint_strings: t.List[str] = [] for constraint in cls._meta.constraints: constraint_strings.append(constraint._table_str()) + constraints_string = spacer.join(constraint_strings) tablename = repr(cls._meta.tablename) @@ -1413,7 +1414,9 @@ def _table_str( ) return ( - f"class {cls.__name__}({class_args}):\n" f" {columns_string}\n" + f"class {cls.__name__}({class_args}):\n" + f" {columns_string}\n" + f" {constraints_string}\n" ) From 0963b584901cfd64aec422be3b30f5484b5e0f9c Mon Sep 17 00:00:00 2001 From: Daniel Townsend Date: Sat, 1 Feb 2025 17:32:06 +0000 Subject: [PATCH 23/29] update tests Call constraints with correct args --- .../migrations/auto/test_schema_differ.py | 28 +++++++++---------- 1 file changed, 14 insertions(+), 14 deletions(-) diff --git a/tests/apps/migrations/auto/test_schema_differ.py b/tests/apps/migrations/auto/test_schema_differ.py index 28127ec32..3c0a71b39 100644 --- a/tests/apps/migrations/auto/test_schema_differ.py +++ b/tests/apps/migrations/auto/test_schema_differ.py @@ -499,14 +499,10 @@ def test_add_table_with_constraint(self) -> None: genre_column = Varchar() genre_column._meta.name = "genre" - # TODO - this doesn't seem to be used, remove it? - name_unique_constraint = UniqueConstraint(unique_columns=["name"]) - name_unique_constraint._meta.name = "unique_name" - name_genre_unique_constraint = UniqueConstraint( - unique_columns=["name", "genre"] + column_names=["name", "genre"], + name="unique_name_genre", ) - name_genre_unique_constraint._meta.name = "unique_name_genre" schema: t.List[DiffableTable] = [ DiffableTable( @@ -557,13 +553,15 @@ def test_add_constraint(self) -> None: genre_column = Varchar() genre_column._meta.name = "genre" - name_unique_constraint = UniqueConstraint(unique_columns=["name"]) - name_unique_constraint._meta.name = "unique_name" + name_unique_constraint = UniqueConstraint( + column_names=["name"], + name="unique_name", + ) name_genre_unique_constraint = UniqueConstraint( - unique_columns=["name", "genre"] + column_names=["name", "genre"], + name="unique_name_genre", ) - name_genre_unique_constraint._meta.name = "unique_name_genre" schema: t.List[DiffableTable] = [ DiffableTable( @@ -605,13 +603,15 @@ def test_drop_constraint(self) -> None: genre_column = Varchar() genre_column._meta.name = "genre" - name_unique_constraint = UniqueConstraint(unique_columns=["name"]) - name_unique_constraint._meta.name = "unique_name" + name_unique_constraint = UniqueConstraint( + column_names=["name"], + name="unique_name", + ) name_genre_unique_constraint = UniqueConstraint( - unique_columns=["name", "genre"] + column_names=["name", "genre"], + name="unique_name_genre", ) - name_genre_unique_constraint._meta.name = "unique_name_genre" schema: t.List[DiffableTable] = [ DiffableTable( From cc107b935ecc9a33328a8d84c9795529d6732f98 Mon Sep 17 00:00:00 2001 From: Daniel Townsend Date: Sat, 1 Feb 2025 17:34:50 +0000 Subject: [PATCH 24/29] update tests --- tests/apps/migrations/auto/test_migration_manager.py | 10 +++++----- tests/apps/migrations/auto/test_schema_differ.py | 4 ++-- tests/apps/migrations/auto/test_schema_snapshot.py | 4 ++-- 3 files changed, 9 insertions(+), 9 deletions(-) diff --git a/tests/apps/migrations/auto/test_migration_manager.py b/tests/apps/migrations/auto/test_migration_manager.py index e72fa77e1..ffbe77b9b 100644 --- a/tests/apps/migrations/auto/test_migration_manager.py +++ b/tests/apps/migrations/auto/test_migration_manager.py @@ -360,7 +360,7 @@ def test_add_table_with_unique_constraint(self): constraint_name="unique_name_label", constraint_class=UniqueConstraint, params={ - "unique_columns": ["name", "label"], + "column_names": ["name", "label"], }, ) asyncio.run(manager.run()) @@ -400,7 +400,7 @@ def test_drop_table_with_unique_constraint( constraint_name="unique_name_label", constraint_class=UniqueConstraint, params={ - "unique_columns": ["name", "label"], + "column_names": ["name", "label"], }, ) asyncio.run(manager_1.run()) @@ -455,7 +455,7 @@ def test_rename_table_with_unique_constraint( constraint_name="unique_name_label", constraint_class=UniqueConstraint, params={ - "unique_columns": ["name", "label"], + "column_names": ["name", "label"], }, ) asyncio.run(manager_1.run()) @@ -524,7 +524,7 @@ def test_add_unique_constraint( constraint_name="musician_unique", constraint_class=UniqueConstraint, params={ - "unique_columns": ["name", "label"], + "column_names": ["name", "label"], }, ) asyncio.run(manager_2.run()) @@ -576,7 +576,7 @@ def test_drop_unique_constraint( constraint_name="musician_unique", constraint_class=UniqueConstraint, params={ - "unique_columns": ["name", "label"], + "column_names": ["name", "label"], }, ) asyncio.run(manager_1.run()) diff --git a/tests/apps/migrations/auto/test_schema_differ.py b/tests/apps/migrations/auto/test_schema_differ.py index 3c0a71b39..376ddc975 100644 --- a/tests/apps/migrations/auto/test_schema_differ.py +++ b/tests/apps/migrations/auto/test_schema_differ.py @@ -540,7 +540,7 @@ def test_add_table_with_constraint(self) -> None: self.assertTrue(len(new_table_constraints.statements) == 1) self.assertEqual( new_table_constraints.statements[0], - "manager.add_constraint(table_class_name='Band', tablename='band', constraint_name='unique_name_genre', constraint_class=UniqueConstraint, params={'unique_columns': ['name', 'genre']}, schema=None)", # noqa + "manager.add_constraint(table_class_name='Band', tablename='band', constraint_name='unique_name_genre', constraint_class=UniqueConstraint, params={'column_names': ['name', 'genre']}, schema=None)", # noqa ) def test_add_constraint(self) -> None: @@ -590,7 +590,7 @@ def test_add_constraint(self) -> None: self.assertTrue(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='unique_name_genre', constraint_class=UniqueConstraint, params={'unique_columns': ['name', 'genre']}, schema=None)", # noqa: E501 + "manager.add_constraint(table_class_name='Band', tablename='band', constraint_name='unique_name_genre', constraint_class=UniqueConstraint, params={'column_names': ['name', 'genre']}, schema=None)", # noqa: E501 ) def test_drop_constraint(self) -> None: diff --git a/tests/apps/migrations/auto/test_schema_snapshot.py b/tests/apps/migrations/auto/test_schema_snapshot.py index ae1742c5b..dbb0a761d 100644 --- a/tests/apps/migrations/auto/test_schema_snapshot.py +++ b/tests/apps/migrations/auto/test_schema_snapshot.py @@ -208,7 +208,7 @@ def test_add_constraint(self): constraint_name="unique_name", constraint_class=UniqueConstraint, params={ - "unique_columns": ["name"], + "column_names": ["name"], }, ) @@ -238,7 +238,7 @@ def test_drop_constraint(self): constraint_name="unique_name", constraint_class=UniqueConstraint, params={ - "unique_columns": ["name"], + "column_names": ["name"], }, ) From 423c40cb66c797dcfd601b8f142f5be4dcacb2d4 Mon Sep 17 00:00:00 2001 From: Daniel Townsend Date: Sat, 1 Feb 2025 17:41:57 +0000 Subject: [PATCH 25/29] fix linter errors --- piccolo/table.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/piccolo/table.py b/piccolo/table.py index a8987c782..f63793de4 100644 --- a/piccolo/table.py +++ b/piccolo/table.py @@ -28,7 +28,7 @@ ) from piccolo.columns.readable import Readable from piccolo.columns.reference import LAZY_COLUMN_REFERENCES -from piccolo.constraints import Check, Constraint, ConstraintConfig, Unique +from piccolo.constraints import Constraint, ConstraintConfig from piccolo.custom_types import TableInstance from piccolo.engine import Engine, engine_finder from piccolo.query import ( From 7995249c3983411b0ad150fde1e44d1adbe93965 Mon Sep 17 00:00:00 2001 From: Daniel Townsend Date: Sat, 1 Feb 2025 20:06:18 +0000 Subject: [PATCH 26/29] update docstrings --- piccolo/constraints.py | 14 ++++---------- 1 file changed, 4 insertions(+), 10 deletions(-) diff --git a/piccolo/constraints.py b/piccolo/constraints.py index 0c21efd26..424dbd3ed 100644 --- a/piccolo/constraints.py +++ b/piccolo/constraints.py @@ -39,9 +39,6 @@ class Album(Table): :param columns: The table columns that should be unique together. - :param name: - The name of the constraint in the database. If not provided, we - generate a sensible default. :param nulls_distinct: See the `Postgres docs `_ for more information. @@ -61,8 +58,8 @@ def __init__( def to_constraint(self) -> UniqueConstraint: """ - You should wait for the ``Table`` metaclass to assign names all of the - columns before calling this method. + You should wait for the ``Table`` metaclass to assign names to all of + the columns before calling this method. """ from piccolo.columns import Column @@ -107,9 +104,6 @@ class Ticket(Table): :param condition: The syntax is the same as the ``where`` clause used by most queries (e.g. ``select``). - :param name: - The name of the constraint in the database. If not provided, we - generate a sensible default. """ @@ -121,8 +115,8 @@ def __init__( def to_constraint(self) -> CheckConstraint: """ - You should wait for the ``Table`` metaclass to assign names all of the - columns before calling this method. + You should wait for the ``Table`` metaclass to assign names to all of + the columns before calling this method. """ from piccolo.columns.combination import CombinableMixin From 1b9e178861c33a5016609e7b8702a362c01d3cd1 Mon Sep 17 00:00:00 2001 From: Daniel Townsend Date: Sat, 1 Feb 2025 20:09:38 +0000 Subject: [PATCH 27/29] Update constraints.rst --- docs/src/piccolo/schema/constraints.rst | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/src/piccolo/schema/constraints.rst b/docs/src/piccolo/schema/constraints.rst index 56f3c2afb..5dc0b2445 100644 --- a/docs/src/piccolo/schema/constraints.rst +++ b/docs/src/piccolo/schema/constraints.rst @@ -15,11 +15,11 @@ argument of ``Column``: ------------------------------------------------------------------------------- -``Table.constraints`` +Advanced constraints ===================== -By adding a ``constraints`` list to your ``Table``, you can implement powerful -``UNIQUE`` and ``CHECK`` constraints. +You can add you can implement powerful ``UNIQUE`` and ``CHECK`` constraints +on your ``Table``. ``Unique`` ---------- From b712fb17e8573d20c2016514db9005663a2be60f Mon Sep 17 00:00:00 2001 From: Daniel Townsend Date: Sat, 1 Feb 2025 20:22:03 +0000 Subject: [PATCH 28/29] add docs for using constraints with mixins --- docs/src/piccolo/schema/advanced.rst | 32 +++++++++++++++++++++---- docs/src/piccolo/schema/constraints.rst | 7 ++++++ 2 files changed, 35 insertions(+), 4 deletions(-) diff --git a/docs/src/piccolo/schema/advanced.rst b/docs/src/piccolo/schema/advanced.rst index 0e4f3db0f..9cb1cd1e1 100644 --- a/docs/src/piccolo/schema/advanced.rst +++ b/docs/src/piccolo/schema/advanced.rst @@ -119,6 +119,8 @@ for example with :ref:`table_finder `. ------------------------------------------------------------------------------- +.. _Mixins: + Mixins ------ @@ -127,15 +129,37 @@ use mixins to reduce the amount of repetition. .. code-block:: python - from piccolo.columns import Varchar, Boolean + from piccolo.columns import Date, Varchar from piccolo.table import Table - class FavouriteMixin: - favourite = Boolean(default=False) + class DateOfBirthMixin: + date_of_birth = Date() + + + class Manager(DateOfBirthMixin, Table): + name = Varchar() + +You can also add :ref:`constraints ` to your mixin classes. + +.. code-block:: python + + import datetime + + from piccolo.columns import Varchar, Date + from piccolo.constraints import Check + from piccolo.table import Table + + + class DateOfBirthMixin: + date_of_birth = Date() + + min_date_of_birth = Check( + date_of_birth >= datetime.date(year=1920, month=1, day=1) + ) - class Manager(FavouriteMixin, Table): + class Manager(DateOfBirthMixin, Table): name = Varchar() ------------------------------------------------------------------------------- diff --git a/docs/src/piccolo/schema/constraints.rst b/docs/src/piccolo/schema/constraints.rst index 5dc0b2445..bdaa1eb3e 100644 --- a/docs/src/piccolo/schema/constraints.rst +++ b/docs/src/piccolo/schema/constraints.rst @@ -15,6 +15,8 @@ argument of ``Column``: ------------------------------------------------------------------------------- +.. _AdvancedConstraints: + Advanced constraints ===================== @@ -41,3 +43,8 @@ constraints will also be created. Also, if using auto migrations, they handle the creation and deletion of these constraints for you. + +Mixins +------ + +Constraints can be added to :ref:`mixin classes ` for reusability. From f49edf0359c52c68778213e8462c72d1eaab8794 Mon Sep 17 00:00:00 2001 From: Daniel Townsend Date: Sat, 1 Feb 2025 20:31:34 +0000 Subject: [PATCH 29/29] make sure dates can be used in check constraints --- piccolo/querystring.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/piccolo/querystring.py b/piccolo/querystring.py index 7dec758a8..8bc080dcd 100644 --- a/piccolo/querystring.py +++ b/piccolo/querystring.py @@ -3,7 +3,7 @@ import typing as t from abc import ABCMeta, abstractmethod from dataclasses import dataclass -from datetime import datetime +from datetime import date, datetime from importlib.util import find_spec from string import Formatter @@ -138,6 +138,10 @@ def __str__(self): """ The SQL returned by the ``__str__`` method isn't used directly in queries - it's just a usability feature. + + The only exception to this is CHECK constraints, where we use this to + convert simple querystrings into strings. + """ _, bundled, combined_args = self.bundle( start_index=1, bundled=[], combined_args=[] @@ -153,7 +157,7 @@ def __str__(self): _type = type(arg) if _type == str: converted_args.append(f"'{arg}'") - elif _type == datetime: + elif _type == datetime or _type == date: dt_string = arg.isoformat() converted_args.append(f"'{dt_string}'") elif _type == UUID or _type == apgUUID: