def test_add_dataset_browse_paths():
dataset = make_generic_dataset()
transformer = AddDatasetBrowsePathTransformer.create(
{"path_templates": ["/abc"]},
PipelineContext(run_id="test"),
)
transformed = list(transformer.transform([RecordEnvelope(dataset, metadata={})]))
browse_path_aspect = builder.get_aspect_if_available(
transformed[0].record, models.BrowsePathsClass
)
assert browse_path_aspect
assert browse_path_aspect.paths == ["/abc"]
transformer = AddDatasetBrowsePathTransformer.create(
{
"path_templates": [
"/PLATFORM/foo/DATASET_PARTS/ENV",
"/ENV/PLATFORM/bar/DATASET_PARTS/",
]
},
PipelineContext(run_id="test"),
)
transformed = list(transformer.transform([RecordEnvelope(dataset, metadata={})]))
browse_path_aspect = builder.get_aspect_if_available(
transformed[0].record, models.BrowsePathsClass
)
assert browse_path_aspect
assert browse_path_aspect.paths == [
"/abc",
"/bigquery/foo/example1/prod",
"/prod/bigquery/bar/example1/",
]
def test_mcp_add_tags_missing(mock_time):
dataset_mcp = make_generic_dataset_mcp()
transformer = SimpleAddDatasetTags.create(
{
"tag_urns": [
builder.make_tag_urn("NeedsDocumentation"),
builder.make_tag_urn("Legacy"),
]
},
PipelineContext(run_id="test-tags"),
)
input_stream: List[RecordEnvelope] = [
RecordEnvelope(input, metadata={}) for input in [dataset_mcp]
]
input_stream.append(RecordEnvelope(record=EndOfStream(), metadata={}))
outputs = list(transformer.transform(input_stream))
assert len(outputs) == 3
assert outputs[0].record == dataset_mcp
# Check that tags were added, this will be the second result
tags_aspect = outputs[1].record.aspect
assert tags_aspect
assert len(tags_aspect.tags) == 2
assert tags_aspect.tags[0].tag == builder.make_tag_urn(
"NeedsDocumentation")
assert isinstance(outputs[-1].record, EndOfStream)
def test_pattern_dataset_ownership_with_type_transformation(mock_time):
input = make_generic_dataset()
transformer = PatternAddDatasetOwnership.create(
{
"owner_pattern": {
"rules": {
".*example1.*": [builder.make_user_urn("person1")],
}
},
"ownership_type": "PRODUCER",
},
PipelineContext(run_id="test"),
)
output = list(
transformer.transform([
RecordEnvelope(input, metadata={}),
RecordEnvelope(EndOfStream(), metadata={}),
]))
assert len(output) == 3
ownership_aspect = output[1].record.aspect
assert ownership_aspect
assert len(ownership_aspect.owners) == 1
assert ownership_aspect.owners[
0].type == models.OwnershipTypeClass.PRODUCER
def test_simple_dataset_ownership_with_type_transformation(mock_time):
input = make_generic_dataset()
transformer = SimpleAddDatasetOwnership.create(
{
"owner_urns": [
builder.make_user_urn("person1"),
],
"ownership_type": "PRODUCER",
},
PipelineContext(run_id="test"),
)
output = list(
transformer.transform([
RecordEnvelope(input, metadata={}),
RecordEnvelope(EndOfStream(), metadata={}),
]))
assert len(output) == 3
# original MCE is unchanged
assert input == output[0].record
ownership_aspect = output[1].record.aspect
assert isinstance(ownership_aspect, OwnershipClass)
assert len(ownership_aspect.owners) == 1
assert ownership_aspect.owners[
0].type == models.OwnershipTypeClass.PRODUCER
def test_old_transformers_working_as_before(mock_time):
dataset_mce = make_generic_dataset()
dataset_mcp = make_generic_dataset_mcp()
transformer = OldMCETransformer.create(
{},
PipelineContext(run_id="test-old-transformer"),
)
outputs = list(
transformer.transform([
RecordEnvelope(input, metadata={})
for input in [dataset_mce, dataset_mcp,
EndOfStream()]
]))
assert len(outputs) == 3 # MCP will come back untouched
assert outputs[0].record == dataset_mce
# Check that glossary terms were added.
props_aspect = builder.get_aspect_if_available(outputs[0].record,
DatasetPropertiesClass)
assert props_aspect
assert props_aspect.description == "Old Transformer was here"
assert outputs[1].record == dataset_mcp
assert isinstance(outputs[-1].record, EndOfStream)
# MCP only stream
dataset_mcps = [
make_generic_dataset_mcp(),
make_generic_dataset_mcp(aspect=DatasetPropertiesClass(
description="Another test MCP")),
EndOfStream(),
]
transformer = OldMCETransformer.create(
{},
PipelineContext(run_id="test-old-transformer"),
)
outputs = list(
transformer.transform(
[RecordEnvelope(input, metadata={}) for input in dataset_mcps]))
assert len(outputs) == 3 # MCP-s will come back untouched
assert outputs[0].record == dataset_mcps[0]
assert outputs[1].record == dataset_mcps[1]
assert isinstance(outputs[-1].record, EndOfStream)
def test_pattern_dataset_tags_transformation(mock_time):
dataset_mce = make_generic_dataset()
transformer = PatternAddDatasetTags.create(
{
"tag_pattern": {
"rules": {
".*example1.*": [
builder.make_tag_urn("Private"),
builder.make_tag_urn("Legacy"),
],
".*example2.*":
[builder.make_term_urn("Needs Documentation")],
}
},
},
PipelineContext(run_id="test-tags"),
)
outputs = list(
transformer.transform(
[RecordEnvelope(input, metadata={}) for input in [dataset_mce]]))
assert len(outputs) == 1
# Check that glossary terms were added.
tags_aspect = builder.get_aspect_if_available(outputs[0].record,
models.GlobalTagsClass)
assert tags_aspect
assert len(tags_aspect.tags) == 2
assert tags_aspect.tags[0].tag == builder.make_tag_urn("Private")
assert builder.make_tag_urn("Needs Documentation") not in tags_aspect.tags
def test_simple_dataset_terms_transformation(mock_time):
dataset_mce = make_generic_dataset()
transformer = SimpleAddDatasetTerms.create(
{
"term_urns": [
builder.make_term_urn("Test"),
builder.make_term_urn("Needs Review"),
]
},
PipelineContext(run_id="test-terms"),
)
outputs = list(
transformer.transform([
RecordEnvelope(input, metadata={})
for input in [dataset_mce, EndOfStream()]
]))
assert len(outputs) == 3
# Check that glossary terms were added.
terms_aspect = outputs[1].record.aspect
assert terms_aspect
assert len(terms_aspect.terms) == 2
assert terms_aspect.terms[0].urn == builder.make_term_urn("Test")
def test_pattern_dataset_tags_transformation(mock_time):
dataset_mce = make_generic_dataset()
transformer = PatternAddDatasetTags.create(
{
"tag_pattern": {
"rules": {
".*example1.*": [
builder.make_tag_urn("Private"),
builder.make_tag_urn("Legacy"),
],
".*example2.*":
[builder.make_term_urn("Needs Documentation")],
}
},
},
PipelineContext(run_id="test-tags"),
)
outputs = list(
transformer.transform([
RecordEnvelope(input, metadata={})
for input in [dataset_mce, EndOfStream()]
]))
assert len(outputs) == 3
tags_aspect = outputs[1].record.aspect
assert tags_aspect
assert len(tags_aspect.tags) == 2
assert tags_aspect.tags[0].tag == builder.make_tag_urn("Private")
assert builder.make_tag_urn("Needs Documentation") not in tags_aspect.tags
def test_simple_add_dataset_properties(mock_time):
dataset_mce = make_dataset_with_properties()
new_properties = {"new-simple-property": "new-value"}
transformer = SimpleAddDatasetProperties.create(
{
"properties": new_properties,
},
PipelineContext(run_id="test-simple-properties"),
)
outputs = list(
transformer.transform(
[RecordEnvelope(input, metadata={}) for input in [dataset_mce]]))
assert len(outputs) == 1
custom_properties = builder.get_aspect_if_available(
outputs[0].record, models.DatasetPropertiesClass)
print(str(custom_properties))
assert custom_properties is not None
assert custom_properties.customProperties == {
**EXISTING_PROPERTIES,
**new_properties,
}
def test_pattern_dataset_terms_transformation(mock_time):
dataset_mce = make_generic_dataset()
transformer = PatternAddDatasetTerms.create(
{
"term_pattern": {
"rules": {
".*example1.*": [
builder.make_term_urn("AccountBalance"),
builder.make_term_urn("Email"),
],
".*example2.*": [builder.make_term_urn("Address")],
}
},
},
PipelineContext(run_id="test-terms"),
)
outputs = list(
transformer.transform(
[RecordEnvelope(input, metadata={}) for input in [dataset_mce]]))
assert len(outputs) == 1
# Check that glossary terms were added.
terms_aspect = builder.get_aspect_if_available(outputs[0].record,
models.GlossaryTermsClass)
assert terms_aspect
assert len(terms_aspect.terms) == 2
assert terms_aspect.terms[0].urn == builder.make_term_urn("AccountBalance")
assert builder.make_term_urn("AccountBalance") not in terms_aspect.terms
def test_kafka_sink_write(self, mock_k_callback, mock_producer,
mock_context):
mock_producer_instance = mock_producer.return_value
mock_k_callback_instance = mock_k_callback.return_value
callback = MagicMock(spec=WriteCallback)
kafka_sink = DatahubKafkaSink.create(
{"connection": {
"bootstrap": "foobar:9092"
}}, mock_context)
mce = builder.make_lineage_mce(
[
builder.make_dataset_urn("bigquery", "upstream1"),
builder.make_dataset_urn("bigquery", "upstream2"),
],
builder.make_dataset_urn("bigquery", "downstream1"),
)
re = RecordEnvelope(record=mce, metadata={})
kafka_sink.write_record_async(re, callback)
mock_producer_instance.poll.assert_called_once(
) # producer should call poll() first
self.validate_kafka_callback(
mock_k_callback, re,
callback) # validate kafka callback was constructed appropriately
# validate that confluent_kafka.Producer.produce was called with the right arguments
mock_producer_instance.produce.assert_called_once()
args, kwargs = mock_producer_instance.produce.call_args
assert kwargs["value"] == mce
assert kwargs["key"] # produce call should include a Kafka key
created_callback = kwargs["on_delivery"]
assert created_callback == mock_k_callback_instance.kafka_callback
def test_simple_dataset_tags_transformation(mock_time):
dataset_mce = make_generic_dataset()
transformer = SimpleAddDatasetTags.create(
{
"tag_urns": [
builder.make_tag_urn("NeedsDocumentation"),
builder.make_tag_urn("Legacy"),
]
},
PipelineContext(run_id="test-tags"),
)
outputs = list(
transformer.transform(
[RecordEnvelope(input, metadata={}) for input in [dataset_mce]]
)
)
assert len(outputs) == 1
# Check that tags were added.
tags_aspect = builder.get_aspect_if_available(
outputs[0].record, models.GlobalTagsClass
)
assert tags_aspect
assert len(tags_aspect.tags) == 2
assert tags_aspect.tags[0].tag == builder.make_tag_urn("NeedsDocumentation")
def run(self) -> None:
callback = LoggingCallback()
extractor: Extractor = self.extractor_class()
for wu in itertools.islice(self.source.get_workunits(),
10 if self.preview_mode else None):
# TODO: change extractor interface
extractor.configure({}, self.ctx)
if not self.dry_run:
self.sink.handle_work_unit_start(wu)
record_envelopes = extractor.get_records(wu)
for record_envelope in self.transform(record_envelopes):
if not self.dry_run:
self.sink.write_record_async(record_envelope, callback)
extractor.close()
if not self.dry_run:
self.sink.handle_work_unit_end(wu)
self.source.close()
# no more data is coming, we need to let the transformers produce any additional records if they are holding on to state
for record_envelope in self.transform([
RecordEnvelope(record=EndOfStream(),
metadata={"workunit_id": "end-of-stream"})
]):
if not self.dry_run and not isinstance(record_envelope.record,
EndOfStream):
# TODO: propagate EndOfStream and other control events to sinks, to allow them to flush etc.
self.sink.write_record_async(record_envelope, callback)
self.sink.close()
self.process_commits()
def transform(
self, record_envelopes: Iterable[RecordEnvelope]
) -> Iterable[RecordEnvelope]:
for envelope in record_envelopes:
if not self._should_process(envelope.record):
# early exit
pass
elif isinstance(envelope.record, MetadataChangeEventClass):
envelope = self._transform_or_record_mce(envelope)
elif isinstance(envelope.record,
MetadataChangeProposalWrapper) and isinstance(
self, SingleAspectTransformer):
return_envelope = self._transform_or_record_mcp(envelope)
if return_envelope is None:
continue
else:
envelope = return_envelope
elif isinstance(envelope.record, EndOfStream) and isinstance(
self, SingleAspectTransformer):
# walk through state and call transform for any unprocessed entities
for urn, state in self.entity_map.items():
if "seen" in state:
# call transform on this entity_urn
last_seen_mcp = state["seen"].get("mcp")
last_seen_mce_system_metadata = state["seen"].get(
"mce")
transformed_aspect = self.transform_aspect(
entity_urn=urn,
aspect_name=self.aspect_name(),
aspect=last_seen_mcp.aspect if last_seen_mcp
and last_seen_mcp.aspectName == self.aspect_name()
else None,
)
if transformed_aspect:
# for end of stream records, we modify the workunit-id
structured_urn = Urn.create_from_string(urn)
simple_name = "-".join(
structured_urn.get_entity_id())
record_metadata = envelope.metadata.copy()
record_metadata.update({
"workunit_id":
f"txform-{simple_name}-{self.aspect_name()}"
})
yield RecordEnvelope(
record=MetadataChangeProposalWrapper(
entityUrn=urn,
entityType=structured_urn.get_type(),
changeType=ChangeTypeClass.UPSERT,
systemMetadata=last_seen_mcp.systemMetadata
if last_seen_mcp else
last_seen_mce_system_metadata,
aspectName=self.aspect_name(),
aspect=transformed_aspect,
),
metadata=record_metadata,
)
self._mark_processed(urn)
yield envelope
def test_simple_dataset_ownership_tranformation(mock_time):
no_owner_aspect = make_generic_dataset()
with_owner_aspect = make_dataset_with_owner()
not_a_dataset = models.MetadataChangeEventClass(
proposedSnapshot=models.DataJobSnapshotClass(
urn="urn:li:dataJob:(urn:li:dataFlow:(airflow,dag_abc,PROD),task_456)",
aspects=[
models.DataJobInfoClass(
name="User Deletions",
description="Constructs the fct_users_deleted from logging_events",
type=models.AzkabanJobTypeClass.SQL,
)
],
)
)
inputs = [
no_owner_aspect,
with_owner_aspect,
not_a_dataset,
]
transformer = SimpleAddDatasetOwnership.create(
{
"owner_urns": [
builder.make_user_urn("person1"),
builder.make_user_urn("person2"),
]
},
PipelineContext(run_id="test"),
)
outputs = list(
transformer.transform([RecordEnvelope(input, metadata={}) for input in inputs])
)
assert len(outputs) == len(inputs)
# Check the first entry.
first_ownership_aspect = builder.get_aspect_if_available(
outputs[0].record, models.OwnershipClass
)
assert first_ownership_aspect
assert len(first_ownership_aspect.owners) == 2
# Check the second entry.
second_ownership_aspect = builder.get_aspect_if_available(
outputs[1].record, models.OwnershipClass
)
assert second_ownership_aspect
assert len(second_ownership_aspect.owners) == 3
# Verify that the third entry is unchanged.
assert inputs[2] == outputs[2].record
def test_import_resolver():
transformer = AddDatasetTags.create(
{
"get_tags_to_add": "tests.unit.test_transform_dataset.dummy_tag_resolver_method"
},
PipelineContext(run_id="test-tags"),
)
output = list(
transformer.transform(
[RecordEnvelope(input, metadata={}) for input in [make_generic_dataset()]]
)
)
assert output
def test_simple_remove_dataset_ownership():
with_owner_aspect = make_dataset_with_owner()
transformer = SimpleRemoveDatasetOwnership.create(
{},
PipelineContext(run_id="test"),
)
outputs = list(
transformer.transform([RecordEnvelope(with_owner_aspect,
metadata={})]))
ownership_aspect = builder.get_aspect_if_available(outputs[0].record,
models.OwnershipClass)
assert ownership_aspect
assert len(ownership_aspect.owners) == 0
def test_mark_status_dataset():
dataset = make_generic_dataset()
transformer = MarkDatasetStatus.create(
{"removed": True},
PipelineContext(run_id="test"),
)
removed = list(transformer.transform([RecordEnvelope(dataset, metadata={})]))
status_aspect = builder.get_aspect_if_available(
removed[0].record, models.StatusClass
)
assert status_aspect
assert status_aspect.removed is True
transformer = MarkDatasetStatus.create(
{"removed": False},
PipelineContext(run_id="test"),
)
not_removed = list(transformer.transform([RecordEnvelope(dataset, metadata={})]))
status_aspect = builder.get_aspect_if_available(
not_removed[0].record, models.StatusClass
)
assert status_aspect
assert status_aspect.removed is False
def test_kafka_sink_write(self, mock_k_callback, mock_producer, mock_context):
mock_producer_instance = mock_producer.return_value
mock_k_callback_instance = mock_k_callback.return_value
callback = MagicMock(spec=WriteCallback)
kafka_sink = DatahubKafkaSink.create(
{"connection": {"bootstrap": "foobar:9092"}}, mock_context
)
re = RecordEnvelope(record=sentinel, metadata={})
kafka_sink.write_record_async(re, callback)
assert mock_producer_instance.poll.call_count == 1 # poll() called once
self.validate_kafka_callback(
mock_k_callback, re, callback
) # validate kafka callback was constructed appropriately
# validate that confluent_kafka.Producer.produce was called with the right arguments
args, kwargs = mock_producer_instance.produce.call_args
created_callback = kwargs["on_delivery"]
assert created_callback == mock_k_callback_instance.kafka_callback
def test_supression_works():
dataset_mce = make_generic_dataset()
dataset_mcp = make_generic_dataset_mcp(
aspect_name="datasetProperties",
aspect=DatasetPropertiesClass(description="supressable description"),
)
transformer = SuppressingTransformer.create(
{},
PipelineContext(run_id="test-suppress-transformer"),
)
outputs = list(
transformer.transform([
RecordEnvelope(input, metadata={})
for input in [dataset_mce, dataset_mcp,
EndOfStream()]
]))
assert len(outputs) == 2 # MCP will be dropped
def test_simple_dataset_ownership_with_type_transformation(mock_time):
input = make_generic_dataset()
transformer = SimpleAddDatasetOwnership.create(
{
"owner_urns": [
builder.make_user_urn("person1"),
],
"ownership_type": "PRODUCER",
},
PipelineContext(run_id="test"),
)
output = list(transformer.transform([RecordEnvelope(input, metadata={})]))
assert len(output) == 1
ownership_aspect = builder.get_aspect_if_available(output[0].record,
models.OwnershipClass)
assert ownership_aspect
assert len(ownership_aspect.owners) == 1
assert ownership_aspect.owners[
0].type == models.OwnershipTypeClass.PRODUCER
def test_kafka_sink_mcp(self, mock_producer, mock_callback):
from datahub.emitter.mcp import MetadataChangeProposalWrapper
mcp = MetadataChangeProposalWrapper(
entityType="dataset",
entityUrn="urn:li:dataset:(urn:li:dataPlatform:mysql,User.UserAccount,PROD)",
changeType=models.ChangeTypeClass.UPSERT,
aspectName="datasetProfile",
aspect=models.DatasetProfileClass(
rowCount=2000,
columnCount=15,
timestampMillis=1626995099686,
),
)
kafka_sink = DatahubKafkaSink.create(
{"connection": {"bootstrap": "localhost:9092"}},
PipelineContext(run_id="test"),
)
kafka_sink.write_record_async(
RecordEnvelope(record=mcp, metadata={}), mock_callback
)
kafka_sink.close()
assert mock_producer.call_count == 2 # constructor should be called
def test_add_dataset_properties(mock_time):
dataset_mce = make_dataset_with_properties()
transformer = AddDatasetProperties.create(
{
"add_properties_resolver_class":
"tests.unit.test_transform_dataset.DummyPropertiesResolverClass"
},
PipelineContext(run_id="test-properties"),
)
outputs = list(
transformer.transform(
[RecordEnvelope(input, metadata={}) for input in [dataset_mce]]))
assert len(outputs) == 1
custom_properties = builder.get_aspect_if_available(
outputs[0].record, models.DatasetPropertiesClass)
assert custom_properties is not None
assert custom_properties.customProperties == {
**EXISTING_PROPERTIES,
**PROPERTIES_TO_ADD,
}
def create_test_data(test_file):
assertion_urn = "urn:li:assertion:2d3b06a6e77e1f24adc9860a05ea089b"
dataset_urn = make_dataset_urn(platform="postgres", name="foo")
assertion_info = AssertionInfoClass(
type=AssertionTypeClass.DATASET,
customProperties={"suite_name": "demo_suite"},
datasetAssertion=DatasetAssertionInfoClass(
fields=[make_schema_field_urn(dataset_urn, "col1")],
dataset=dataset_urn,
scope=DatasetAssertionScopeClass.DATASET_COLUMN,
operator=AssertionStdOperatorClass.LESS_THAN,
nativeType="column_value_is_less_than",
aggregation=AssertionStdAggregation.IDENTITY,
nativeParameters={"max_value": "99"},
),
)
# The assertion definition
mcp1 = MetadataChangeProposalWrapper(
entityType="assertion",
changeType="UPSERT",
entityUrn=assertion_urn,
aspectName="assertionInfo",
aspect=assertion_info,
)
timestamps = [
1643794280350,
1643794280352,
1643794280354,
1643880726872,
1643880726874,
1643880726875,
]
msg_ids = []
# The assertion run event attached to the dataset
mcp2 = MetadataChangeProposalWrapper(
entityType="assertion",
entityUrn=assertion_urn,
changeType="UPSERT",
aspectName="assertionRunEvent",
aspect=AssertionRunEventClass(
timestampMillis=timestamps[0],
partitionSpec=PartitionSpecClass(
partition="[{'country': 'IN'}]",
type=PartitionTypeClass.PARTITION,
),
messageId=str(timestamps[0]),
assertionUrn=assertion_urn,
asserteeUrn=dataset_urn,
result=AssertionResultClass(
type=AssertionResultTypeClass.SUCCESS,
actualAggValue=90,
externalUrl="http://example.com/uuid1",
),
runId="uuid1",
status=AssertionRunStatusClass.COMPLETE,
),
)
mcp3 = MetadataChangeProposalWrapper(
entityType="assertion",
entityUrn=assertion_urn,
changeType="UPSERT",
aspectName="assertionRunEvent",
aspect=AssertionRunEventClass(
timestampMillis=timestamps[1],
partitionSpec=PartitionSpecClass(
partition="[{'country': 'US'}]",
type=PartitionTypeClass.PARTITION,
),
messageId=str(timestamps[1]),
assertionUrn=assertion_urn,
asserteeUrn=dataset_urn,
result=AssertionResultClass(
type=AssertionResultTypeClass.FAILURE,
actualAggValue=101,
externalUrl="http://example.com/uuid1",
),
runId="uuid1",
status=AssertionRunStatusClass.COMPLETE,
),
)
# Result of evaluating this assertion on the whole dataset
mcp4 = MetadataChangeProposalWrapper(
entityType="assertion",
entityUrn=assertion_urn,
changeType="UPSERT",
aspectName="assertionRunEvent",
aspect=AssertionRunEventClass(
timestampMillis=timestamps[2],
partitionSpec=PartitionSpecClass(
partition="FULL_TABLE_SNAPSHOT",
type=PartitionTypeClass.FULL_TABLE,
),
messageId=str(timestamps[2]),
assertionUrn=assertion_urn,
asserteeUrn=dataset_urn,
result=AssertionResultClass(
type=AssertionResultTypeClass.SUCCESS,
actualAggValue=93,
externalUrl="http://example.com/uuid1",
),
runId="uuid1",
status=AssertionRunStatusClass.COMPLETE,
),
)
mcp5 = MetadataChangeProposalWrapper(
entityType="assertion",
entityUrn=assertion_urn,
changeType="UPSERT",
aspectName="assertionRunEvent",
aspect=AssertionRunEventClass(
timestampMillis=timestamps[3],
partitionSpec=PartitionSpecClass(
partition="[{'country': 'IN'}]",
type=PartitionTypeClass.PARTITION,
),
messageId=str(timestamps[3]),
assertionUrn=assertion_urn,
asserteeUrn=dataset_urn,
result=AssertionResultClass(
type=AssertionResultTypeClass.SUCCESS,
actualAggValue=90,
externalUrl="http://example.com/uuid1",
),
runId="uuid1",
status=AssertionRunStatusClass.COMPLETE,
),
)
mcp6 = MetadataChangeProposalWrapper(
entityType="assertion",
entityUrn=assertion_urn,
changeType="UPSERT",
aspectName="assertionRunEvent",
aspect=AssertionRunEventClass(
timestampMillis=timestamps[4],
partitionSpec=PartitionSpecClass(
partition="[{'country': 'US'}]",
type=PartitionTypeClass.PARTITION,
),
messageId=str(timestamps[4]),
assertionUrn=assertion_urn,
asserteeUrn=dataset_urn,
result=AssertionResultClass(
type=AssertionResultTypeClass.FAILURE,
actualAggValue=101,
externalUrl="http://example.com/uuid1",
),
runId="uuid1",
status=AssertionRunStatusClass.COMPLETE,
),
)
# Result of evaluating this assertion on the whole dataset
mcp7 = MetadataChangeProposalWrapper(
entityType="assertion",
entityUrn=assertion_urn,
changeType="UPSERT",
aspectName="assertionRunEvent",
aspect=AssertionRunEventClass(
timestampMillis=timestamps[5],
partitionSpec=PartitionSpecClass(
partition="FULL_TABLE_SNAPSHOT",
type=PartitionTypeClass.FULL_TABLE,
),
messageId=str(timestamps[5]),
assertionUrn=assertion_urn,
asserteeUrn=dataset_urn,
result=AssertionResultClass(
type=AssertionResultTypeClass.SUCCESS,
actualAggValue=93,
externalUrl="http://example.com/uuid1",
),
runId="uuid1",
status=AssertionRunStatusClass.COMPLETE,
),
)
fileSink: FileSink = FileSink.create(
FileSinkConfig(filename=test_file),
ctx=PipelineContext(run_id="test-file"))
for mcp in [mcp1, mcp2, mcp3, mcp4, mcp5, mcp6, mcp7]:
fileSink.write_record_async(RecordEnvelope(record=mcp, metadata={}),
write_callback=NoopWriteCallback())
fileSink.close()
def test_pattern_dataset_ownership_tranformation(mock_time):
no_owner_aspect = make_generic_dataset()
with_owner_aspect = models.MetadataChangeEventClass(
proposedSnapshot=models.DatasetSnapshotClass(
urn="urn:li:dataset:(urn:li:dataPlatform:bigquery,example2,PROD)",
aspects=[
models.OwnershipClass(
owners=[
models.OwnerClass(
owner=builder.make_user_urn("fake_owner"),
type=models.OwnershipTypeClass.DATAOWNER,
),
],
lastModified=models.AuditStampClass(
time=1625266033123, actor="urn:li:corpuser:datahub"
),
)
],
),
)
not_a_dataset = models.MetadataChangeEventClass(
proposedSnapshot=models.DataJobSnapshotClass(
urn="urn:li:dataJob:(urn:li:dataFlow:(airflow,dag_abc,PROD),task_456)",
aspects=[
models.DataJobInfoClass(
name="User Deletions",
description="Constructs the fct_users_deleted from logging_events",
type=models.AzkabanJobTypeClass.SQL,
)
],
)
)
inputs = [
no_owner_aspect,
with_owner_aspect,
not_a_dataset,
]
transformer = PatternAddDatasetOwnership.create(
{
"owner_pattern": {
"rules": {
".*example1.*": [builder.make_user_urn("person1")],
".*example2.*": [builder.make_user_urn("person2")],
}
},
},
PipelineContext(run_id="test"),
)
outputs = list(
transformer.transform([RecordEnvelope(input, metadata={}) for input in inputs])
)
assert len(outputs) == len(inputs)
# Check the first entry.
first_ownership_aspect = builder.get_aspect_if_available(
outputs[0].record, models.OwnershipClass
)
assert first_ownership_aspect
assert len(first_ownership_aspect.owners) == 1
# Check the second entry.
second_ownership_aspect = builder.get_aspect_if_available(
outputs[1].record, models.OwnershipClass
)
assert second_ownership_aspect
assert len(second_ownership_aspect.owners) == 2
# Verify that the third entry is unchanged.
assert inputs[2] == outputs[2].record
def generate(
schema_files: List[str],
server: Optional[str],
file: Optional[str],
dot: Optional[str],
png: Optional[str],
extra_docs: Optional[str],
) -> None:
logger.info(f"server = {server}")
logger.info(f"file = {file}")
logger.info(f"dot = {dot}")
logger.info(f"png = {png}")
entity_extra_docs = {}
if extra_docs:
for path in glob.glob(f"{extra_docs}/**/*.md", recursive=True):
m = re.search("/docs/entities/(.*)/*.md", path)
if m:
entity_name = m.group(1)
with open(path, "r") as doc_file:
file_contents = doc_file.read()
final_markdown = preprocess_markdown(file_contents)
entity_extra_docs[entity_name] = final_markdown
for schema_file in schema_files:
if schema_file.endswith(".yml") or schema_file.endswith(".yaml"):
# registry file
load_registry_file(schema_file)
else:
# schema file
load_schema_file(schema_file)
if entity_extra_docs:
for entity_name in entity_extra_docs:
entity_registry.get(
entity_name).doc_file_contents = entity_extra_docs[entity_name]
relationship_graph = RelationshipGraph()
events = generate_stitched_record(relationship_graph)
generated_docs_dir = "../docs/generated/metamodel"
import shutil
shutil.rmtree(f"{generated_docs_dir}/entities", ignore_errors=True)
entity_names = [(x, entity_registry.get(x))
for x in generated_documentation]
sorted_entity_names = get_sorted_entity_names(entity_names)
index = 0
for category, sorted_entities in sorted_entity_names:
for entity_name in sorted_entities:
entity_def = entity_registry.get(entity_name)
entity_category = entity_def.category
entity_dir = f"{generated_docs_dir}/entities/"
import os
os.makedirs(entity_dir, exist_ok=True)
with open(f"{entity_dir}/{entity_name}.md", "w") as fp:
fp.write("---\n")
fp.write(f"sidebar_position: {index}\n")
fp.write("---\n")
fp.write(generated_documentation[entity_name])
index += 1
if file:
logger.info(f"Will write events to {file}")
Path(file).parent.mkdir(parents=True, exist_ok=True)
fileSink = FileSink(
PipelineContext(run_id="generated-metaModel"),
FileSinkConfig(filename=file),
)
for e in events:
fileSink.write_record_async(RecordEnvelope(e, metadata={}),
write_callback=NoopWriteCallback())
fileSink.close()
pipeline_config = {
"source": {
"type": "file",
"config": {
"filename": file
},
},
"sink": {
"type": "datahub-rest",
"config": {
"server": "${DATAHUB_SERVER:-http://localhost:8080}",
"token": "${DATAHUB_TOKEN:-}",
},
},
"run_id": "modeldoc-generated",
}
pipeline_file = Path(file).parent.absolute() / "pipeline.yml"
with open(pipeline_file, "w") as f:
json.dump(pipeline_config, f, indent=2)
logger.info(f"Wrote pipeline to {pipeline_file}")
if server:
logger.info(f"Will send events to {server}")
assert server.startswith(
"http://"), "server address must start with http://"
emitter = DatahubRestEmitter(gms_server=server)
emitter.test_connection()
for e in events:
emitter.emit(e)
if dot:
logger.info(f"Will write dot file to {dot}")
import pydot
graph = pydot.Dot("my_graph", graph_type="graph")
for node, adjacency in relationship_graph.map.items():
my_node = pydot.Node(
node,
label=node,
shape="box",
)
graph.add_node(my_node)
if adjacency.self_loop:
for relnship in adjacency.self_loop:
graph.add_edge(
pydot.Edge(src=relnship.src,
dst=relnship.dst,
label=relnship.name))
if adjacency.outgoing:
for relnship in adjacency.outgoing:
graph.add_edge(
pydot.Edge(src=relnship.src,
dst=relnship.dst,
label=relnship.name))
Path(dot).parent.mkdir(parents=True, exist_ok=True)
graph.write_raw(dot)
if png:
try:
graph.write_png(png)
except Exception as e:
logger.error(
"Failed to create png file. Do you have graphviz installed?"
)
raise e
def test_mark_status_dataset(tmp_path):
dataset = make_generic_dataset()
transformer = MarkDatasetStatus.create(
{"removed": True},
PipelineContext(run_id="test"),
)
removed = list(
transformer.transform([
RecordEnvelope(dataset, metadata={}),
]))
assert len(removed) == 1
status_aspect = builder.get_aspect_if_available(removed[0].record,
models.StatusClass)
assert status_aspect
assert status_aspect.removed is True
transformer = MarkDatasetStatus.create(
{"removed": False},
PipelineContext(run_id="test"),
)
not_removed = list(
transformer.transform([
RecordEnvelope(dataset, metadata={}),
]))
assert len(not_removed) == 1
status_aspect = builder.get_aspect_if_available(not_removed[0].record,
models.StatusClass)
assert status_aspect
assert status_aspect.removed is False
mcp = make_generic_dataset_mcp(
aspect_name="datasetProperties",
aspect=DatasetPropertiesClass(description="Test dataset"),
)
events_file = create_and_run_test_pipeline(
events=[mcp],
transformers=[{
"type": "mark_dataset_status",
"config": {
"removed": True
}
}],
path=tmp_path,
)
# assert dataset properties aspect was preserved
assert (tests.test_helpers.mce_helpers.assert_for_each_entity(
entity_type="dataset",
aspect_name="datasetProperties",
aspect_field_matcher={"description": "Test dataset"},
file=events_file,
) == 1)
# assert Status aspect was generated
assert (tests.test_helpers.mce_helpers.assert_for_each_entity(
entity_type="dataset",
aspect_name="status",
aspect_field_matcher={"removed": True},
file=events_file,
) == 1)
# MCE only
test_aspect = DatasetPropertiesClass(description="Test dataset")
events_file = create_and_run_test_pipeline(
events=[make_generic_dataset(aspects=[test_aspect])],
transformers=[{
"type": "mark_dataset_status",
"config": {
"removed": True
}
}],
path=tmp_path,
)
# assert dataset properties aspect was preserved
assert (tests.test_helpers.mce_helpers.assert_entity_mce_aspect(
entity_urn=mcp.entityUrn or "",
aspect=test_aspect,
aspect_type=DatasetPropertiesClass,
file=events_file,
) == 1)
# assert Status aspect was generated
assert (tests.test_helpers.mce_helpers.assert_for_each_entity(
entity_type="dataset",
aspect_name="status",
aspect_field_matcher={"removed": True},
file=events_file,
) == 1)
# MCE (non-matching) + MCP (matching)
test_aspect = DatasetPropertiesClass(description="Test dataset")
events_file = create_and_run_test_pipeline(
events=[
make_generic_dataset(aspects=[test_aspect]),
make_generic_dataset_mcp(),
],
transformers=[{
"type": "mark_dataset_status",
"config": {
"removed": True
}
}],
path=tmp_path,
)
# assert dataset properties aspect was preserved
assert (tests.test_helpers.mce_helpers.assert_entity_mce_aspect(
entity_urn=mcp.entityUrn or "",
aspect=test_aspect,
aspect_type=DatasetPropertiesClass,
file=events_file,
) == 1)
# assert Status aspect was generated
assert (tests.test_helpers.mce_helpers.assert_for_each_entity(
entity_type="dataset",
aspect_name="status",
aspect_field_matcher={"removed": True},
file=events_file,
) == 1)
# MCE (matching) + MCP (non-matching)
test_status_aspect = StatusClass(removed=False)
events_file = create_and_run_test_pipeline(
events=[
make_generic_dataset(aspects=[test_status_aspect]),
make_generic_dataset_mcp(
aspect_name="datasetProperties",
aspect=DatasetPropertiesClass(description="test dataset"),
),
],
transformers=[{
"type": "mark_dataset_status",
"config": {
"removed": True
}
}],
path=tmp_path,
)
# assert MCE was transformed
assert (tests.test_helpers.mce_helpers.assert_entity_mce_aspect(
entity_urn=mcp.entityUrn or "",
aspect=StatusClass(removed=True),
aspect_type=StatusClass,
file=events_file,
) == 1)
# assert MCP aspect was preserved
assert (tests.test_helpers.mce_helpers.assert_for_each_entity(
entity_type="dataset",
aspect_name="datasetProperties",
aspect_field_matcher={"description": "test dataset"},
file=events_file,
) == 1)
# MCE (non-matching) + MCP (non-matching)
test_mcp_aspect = GlobalTagsClass(
tags=[TagAssociationClass(tag="urn:li:tag:test")])
test_dataset_props_aspect = DatasetPropertiesClass(
description="Test dataset")
events_file = create_and_run_test_pipeline(
events=[
make_generic_dataset(aspects=[test_dataset_props_aspect]),
make_generic_dataset_mcp(aspect_name="globalTags",
aspect=test_mcp_aspect),
],
transformers=[{
"type": "mark_dataset_status",
"config": {
"removed": True
}
}],
path=tmp_path,
)
# assert MCE was preserved
assert (tests.test_helpers.mce_helpers.assert_entity_mce_aspect(
entity_urn=mcp.entityUrn or "",
aspect=test_dataset_props_aspect,
aspect_type=DatasetPropertiesClass,
file=events_file,
) == 1)
# assert MCP aspect was preserved
assert (tests.test_helpers.mce_helpers.assert_for_each_entity(
entity_type="dataset",
aspect_name="globalTags",
aspect_field_matcher={"tags": [{
"tag": "urn:li:tag:test"
}]},
file=events_file,
) == 1)
# assert MCP Status aspect was generated
assert (tests.test_helpers.mce_helpers.assert_for_each_entity(
entity_type="dataset",
aspect_name="status",
aspect_field_matcher={"removed": True},
file=events_file,
) == 1)
def test_add_dataset_browse_paths():
dataset = make_generic_dataset()
transformer = AddDatasetBrowsePathTransformer.create(
{"path_templates": ["/abc"]},
PipelineContext(run_id="test"),
)
transformed = list(
transformer.transform([
RecordEnvelope(dataset, metadata={}),
RecordEnvelope(EndOfStream(), metadata={}),
]))
browse_path_aspect = transformed[1].record.aspect
assert browse_path_aspect
assert browse_path_aspect.paths == ["/abc"]
# use an mce with a pre-existing browse path
dataset_mce = make_generic_dataset(
aspects=[StatusClass(removed=False), browse_path_aspect])
transformer = AddDatasetBrowsePathTransformer.create(
{
"path_templates": [
"/PLATFORM/foo/DATASET_PARTS/ENV",
"/ENV/PLATFORM/bar/DATASET_PARTS/",
]
},
PipelineContext(run_id="test"),
)
transformed = list(
transformer.transform([
RecordEnvelope(dataset_mce, metadata={}),
RecordEnvelope(EndOfStream(), metadata={}),
]))
assert len(transformed) == 2
browse_path_aspect = builder.get_aspect_if_available(
transformed[0].record, BrowsePathsClass)
assert browse_path_aspect
assert browse_path_aspect.paths == [
"/abc",
"/bigquery/foo/example1/prod",
"/prod/bigquery/bar/example1/",
]
transformer = AddDatasetBrowsePathTransformer.create(
{
"path_templates": [
"/xyz",
],
"replace_existing": True,
},
PipelineContext(run_id="test"),
)
transformed = list(
transformer.transform([
RecordEnvelope(dataset_mce, metadata={}),
RecordEnvelope(EndOfStream(), metadata={}),
]))
assert len(transformed) == 2
browse_path_aspect = builder.get_aspect_if_available(
transformed[0].record, BrowsePathsClass)
assert browse_path_aspect
assert browse_path_aspect.paths == [
"/xyz",
]
def generate(schema_files: List[str], server: Optional[str],
file: Optional[str], dot: Optional[str],
png: Optional[str]) -> None:
logger.info(f"server = {server}")
logger.info(f"file = {file}")
logger.info(f"dot = {dot}")
logger.info(f"png = {png}")
for schema_file in schema_files:
if schema_file.endswith(".yml") or schema_file.endswith(".yaml"):
# registry file
load_registry_file(schema_file)
else:
# schema file
load_schema_file(schema_file)
relationship_graph = RelationshipGraph()
events = generate_stitched_record(relationship_graph)
if file:
logger.info(f"Will write events to {file}")
Path(file).parent.mkdir(parents=True, exist_ok=True)
fileSink = FileSink(
PipelineContext(run_id="generated-metaModel"),
FileSinkConfig(filename=file),
)
for e in events:
fileSink.write_record_async(RecordEnvelope(e, metadata={}),
write_callback=NoopWriteCallback())
fileSink.close()
pipeline_config = {
"source": {
"type": "file",
"config": {
"filename": file
},
},
"sink": {
"type": "datahub-rest",
"config": {
"server": "${DATAHUB_SERVER:-http://localhost:8080}",
"token": "${DATAHUB_TOKEN:-}",
},
},
"run_id": "modeldoc-generated",
}
pipeline_file = Path(file).parent.absolute() / "pipeline.yml"
with open(pipeline_file, "w") as f:
json.dump(pipeline_config, f, indent=2)
logger.info(f"Wrote pipeline to {pipeline_file}")
if server:
logger.info(f"Will send events to {server}")
assert server.startswith(
"http://"), "server address must start with http://"
emitter = DatahubRestEmitter(gms_server=server)
emitter.test_connection()
for e in events:
emitter.emit(e)
if dot:
logger.info(f"Will write dot file to {dot}")
import pydot
graph = pydot.Dot("my_graph", graph_type="graph")
for node, adjacency in relationship_graph.map.items():
my_node = pydot.Node(
node,
label=node,
shape="box",
)
graph.add_node(my_node)
if adjacency.self_loop:
for relnship in adjacency.self_loop:
graph.add_edge(
pydot.Edge(src=relnship.src,
dst=relnship.dst,
label=relnship.name))
if adjacency.outgoing:
for relnship in adjacency.outgoing:
graph.add_edge(
pydot.Edge(src=relnship.src,
dst=relnship.dst,
label=relnship.name))
Path(dot).parent.mkdir(parents=True, exist_ok=True)
graph.write_raw(dot)
if png:
try:
graph.write_png(png)
except Exception as e:
logger.error(
"Failed to create png file. Do you have graphviz installed?"
)
raise e