datahub/metadata-ingestion/examples/library/form_create_with_dynamic_assignment.py

103 lines
3.1 KiB
Python

import logging
import os
from datahub.emitter.mcp import MetadataChangeProposalWrapper
from datahub.emitter.rest_emitter import DatahubRestEmitter
from datahub.metadata.schema_classes import (
CriterionClass,
DynamicFormAssignmentClass,
FilterClass,
FormActorAssignmentClass,
FormInfoClass,
FormPromptClass,
FormPromptTypeClass,
FormTypeClass,
StructuredPropertyParamsClass,
)
from datahub.metadata.urns import FormUrn
log = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)
# Create a form that will be dynamically assigned to all Snowflake datasets
# in the "Finance" domain
# Define the form metadata
form_urn = FormUrn("snowflake_finance_compliance")
form_info = FormInfoClass(
name="Snowflake Finance Data Compliance",
description="Compliance form for all Snowflake datasets in the Finance domain",
type=FormTypeClass.VERIFICATION,
actors=FormActorAssignmentClass(owners=True),
prompts=[
FormPromptClass(
id="retention_time_prompt",
title="Data Retention Period",
description="Specify how long this data should be retained",
type=FormPromptTypeClass.STRUCTURED_PROPERTY,
structuredPropertyParams=StructuredPropertyParamsClass(
urn="urn:li:structuredProperty:io.acryl.dataRetentionTime"
),
required=True,
),
FormPromptClass(
id="pii_classification_prompt",
title="PII Classification",
description="Classify whether this dataset contains PII",
type=FormPromptTypeClass.STRUCTURED_PROPERTY,
structuredPropertyParams=StructuredPropertyParamsClass(
urn="urn:li:structuredProperty:io.acryl.piiClassification"
),
required=True,
),
],
)
# Define dynamic assignment filter
# This form will be assigned to all entities matching these criteria
dynamic_assignment = DynamicFormAssignmentClass(
filter=FilterClass(
criteria=[
CriterionClass(
field="platform",
value="urn:li:dataPlatform:snowflake",
condition="EQUAL",
),
CriterionClass(
field="domains",
value="urn:li:domain:finance",
condition="EQUAL",
),
CriterionClass(
field="_entityType",
value="urn:li:entityType:dataset",
condition="EQUAL",
),
]
)
)
# Create rest emitter
rest_emitter = DatahubRestEmitter(
gms_server=os.getenv("DATAHUB_GMS_URL", "http://localhost:8080"),
token=os.getenv("DATAHUB_GMS_TOKEN"),
)
# Emit the form info
form_info_event = MetadataChangeProposalWrapper(
entityUrn=str(form_urn),
aspect=form_info,
)
rest_emitter.emit(form_info_event)
# Emit the dynamic assignment
dynamic_assignment_event = MetadataChangeProposalWrapper(
entityUrn=str(form_urn),
aspect=dynamic_assignment,
)
rest_emitter.emit(dynamic_assignment_event)
log.info(
f"Created form {form_urn} with dynamic assignment to Snowflake datasets in Finance domain"
)