From 3eaf7f7d30e5b571ede7bb4ffced32501c86bc3e Mon Sep 17 00:00:00 2001 From: Michelle Ark Date: Fri, 5 Jul 2024 12:28:44 -0400 Subject: [PATCH] add to and to_column(s) to ColumnLevelConstraint and ModelLevelConstraint contracts --- dbt_common/contracts/constraints.py | 24 +++++++++++++++++++++++- 1 file changed, 23 insertions(+), 1 deletion(-) diff --git a/dbt_common/contracts/constraints.py b/dbt_common/contracts/constraints.py index c01ee6f8..7cb11afc 100644 --- a/dbt_common/contracts/constraints.py +++ b/dbt_common/contracts/constraints.py @@ -2,7 +2,7 @@ from enum import Enum from typing import Optional, List -from dbt_common.dataclass_schema import dbtClassMixin +from dbt_common.dataclass_schema import dbtClassMixin, ValidationError class ConstraintType(str, Enum): @@ -36,8 +36,30 @@ class ColumnLevelConstraint(dbtClassMixin): warn_unsupported: bool = ( True # Warn if constraint is not supported by the platform and won't be in DDL ) + to: Optional[str] = None + to_column: Optional[str] = None + + @classmethod + def validate(cls, data): + super().validate(data) + if data.get("type") is not ConstraintType.foreign_key: + if data.get("to") is not None: + raise ValidationError(f"Only column-level constraint of type {ConstraintType.foreign_key} can specify a 'to' field.") + if data.get("to_column") is not None: + raise ValidationError(f"Only column-level constraint of type {ConstraintType.foreign_key} can specify a 'to_column' field.") @dataclass class ModelLevelConstraint(ColumnLevelConstraint): columns: List[str] = field(default_factory=list) + to: Optional[str] = None + to_columns: List[str] = field(default_factory=list) + + @classmethod + def validate(cls, data): + super().validate(data) + if data.get("type") is not ConstraintType.foreign_key: + if data.get("to") is not None: + raise ValidationError(f"Only model-level constraint of type {ConstraintType.foreign_key} can specify a 'to' field.") + if data.get("to_columns") is not None: + raise ValidationError(f"Only model-level constraint of type {ConstraintType.foreign_key} can specify a 'to_columns' field.")