diff --git a/datahub-actions/src/datahub_actions/plugin/action/propagation/propagation_utils.py b/datahub-actions/src/datahub_actions/plugin/action/propagation/propagation_utils.py index 1cf07777..c27ec602 100644 --- a/datahub-actions/src/datahub_actions/plugin/action/propagation/propagation_utils.py +++ b/datahub-actions/src/datahub_actions/plugin/action/propagation/propagation_utils.py @@ -24,6 +24,7 @@ from datahub.emitter.mce_builder import make_schema_field_urn from datahub.ingestion.graph.client import DataHubGraph, SearchFilterRule from datahub.metadata.schema_classes import MetadataAttributionClass +from datahub.utilities.str_enum import StrEnum from datahub.utilities.urns.urn import Urn, guess_entity_type from pydantic import validator from pydantic.fields import Field @@ -35,13 +36,13 @@ SYSTEM_ACTOR = "urn:li:corpuser:__datahub_system" -class RelationshipType(str, Enum): +class RelationshipType(StrEnum): LINEAGE = "lineage" # signifies all types of lineage HIERARCHY = "hierarchy" # signifies all types of hierarchy SIBLING = "sibling" # signifies all types of sibling -class DirectionType(str, Enum): +class DirectionType(StrEnum): UP = "up" # signifies upstream or parent (depending on relationship type) DOWN = "down" # signifies downstream or child (depending on relationship type) ALL = "all" # signifies all directions