Skip to content

Commit

Permalink
add to and to_column(s) to ColumnLevelConstraint and ModelLevelConstr…
Browse files Browse the repository at this point in the history
…aint contracts
  • Loading branch information
MichelleArk committed Jul 5, 2024
1 parent a28ff8a commit 3eaf7f7
Showing 1 changed file with 23 additions and 1 deletion.
24 changes: 23 additions & 1 deletion dbt_common/contracts/constraints.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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.")

0 comments on commit 3eaf7f7

Please sign in to comment.