def test_kafka_source_workunits_topic_pattern(self, mock_kafka):
mock_kafka_instance = mock_kafka.return_value
mock_cluster_metadata = MagicMock()
mock_cluster_metadata.topics = ["test", "foobar", "bazbaz"]
mock_kafka_instance.list_topics.return_value = mock_cluster_metadata
ctx = PipelineContext(run_id="test1")
kafka_source = KafkaSource.create(
{
"topic_patterns": {"allow": ["test"]},
"connection": {"bootstrap": "localhost:9092"},
},
ctx,
)
workunits = [w for w in kafka_source.get_workunits()]
mock_kafka.assert_called_once()
mock_kafka_instance.list_topics.assert_called_once()
assert len(workunits) == 2
mock_cluster_metadata.topics = ["test", "test2", "bazbaz"]
ctx = PipelineContext(run_id="test2")
kafka_source = KafkaSource.create(
{
"topic_patterns": {"allow": ["test.*"]},
"connection": {"bootstrap": "localhost:9092"},
},
ctx,
)
workunits = [w for w in kafka_source.get_workunits()]
assert len(workunits) == 4
def __init__(self,
config: PipelineConfig,
dry_run: bool = False,
preview_mode: bool = False):
self.config = config
self.dry_run = dry_run
self.preview_mode = preview_mode
self.ctx = PipelineContext(
run_id=self.config.run_id,
datahub_api=self.config.datahub_api,
pipeline_name=self.config.pipeline_name,
dry_run=dry_run,
preview_mode=preview_mode,
)
source_type = self.config.source.type
source_class = source_registry.get(source_type)
self.source: Source = source_class.create(
self.config.source.dict().get("config", {}), self.ctx)
logger.debug(f"Source type:{source_type},{source_class} configured")
sink_type = self.config.sink.type
sink_class = sink_registry.get(sink_type)
sink_config = self.config.sink.dict().get("config", {})
self.sink: Sink = sink_class.create(sink_config, self.ctx)
logger.debug(
f"Sink type:{self.config.sink.type},{sink_class} configured")
self.extractor_class = extractor_registry.get(
self.config.source.extractor)
self._configure_transforms()
self._configure_reporting()
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 create_mocked_dbt_source() -> DBTSource:
ctx = PipelineContext("test-run-id")
graph = mock.MagicMock()
graph.get_ownership.return_value = mce_builder.make_ownership_aspect_from_urn_list(
["urn:li:corpuser:test_user"], "AUDIT")
graph.get_glossary_terms.return_value = (
mce_builder.make_glossary_terms_aspect_from_urn_list(
["urn:li:glossaryTerm:old", "urn:li:glossaryTerm:old2"]))
graph.get_tags.return_value = mce_builder.make_global_tag_aspect_with_tag_list(
["non_dbt_existing", "dbt:existing"])
ctx.graph = graph
return DBTSource(DBTConfig(**create_base_dbt_config()), ctx, "dbt")
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_underlying_platform_cannot_be_other_than_athena():
source = GlueSource(
ctx=PipelineContext(run_id="glue-source-test"),
config=GlueSourceConfig(aws_region="us-west-2",
underlying_platform="data-warehouse"),
)
assert source.get_underlying_platform() == "glue"
def test_kafka_source_stateful_ingestion_requires_platform_instance(
self, ):
class StatefulProviderMock:
def __init__(self, config, ctx):
self.ctx = ctx
self.config = config
def is_stateful_ingestion_configured(self):
return self.config.stateful_ingestion.enabled
kafka_source_patcher = unittest.mock.patch.object(
KafkaSource, "__bases__", (StatefulProviderMock, ))
ctx = PipelineContext(run_id="test", pipeline_name="test")
with pytest.raises(ConfigurationError):
with kafka_source_patcher:
# prevent delattr on __bases__ on context __exit__
kafka_source_patcher.is_local = True
KafkaSource.create(
{
"stateful_ingestion": {
"enabled": "true"
},
"connection": {
"bootstrap": "localhost:9092"
},
},
ctx,
)
def test_upstream_table_generation_with_temporary_table_with_multiple_temp_upstream():
from datahub.ingestion.api.common import PipelineContext
from datahub.ingestion.source.sql.bigquery import BigQueryConfig, BigQuerySource
from datahub.ingestion.source.usage.bigquery_usage import BigQueryTableRef
a: BigQueryTableRef = BigQueryTableRef(
project="test-project", dataset="test-dataset", table="a"
)
b: BigQueryTableRef = BigQueryTableRef(
project="test-project", dataset="_temp-dataset", table="b"
)
c: BigQueryTableRef = BigQueryTableRef(
project="test-project", dataset="test-dataset", table="c"
)
d: BigQueryTableRef = BigQueryTableRef(
project="test-project", dataset="_test-dataset", table="d"
)
e: BigQueryTableRef = BigQueryTableRef(
project="test-project", dataset="test-dataset", table="e"
)
config = BigQueryConfig.parse_obj(
{
"project_id": "test-project",
}
)
source = BigQuerySource(config=config, ctx=PipelineContext(run_id="test"))
source.lineage_metadata = {
str(a): set([str(b)]),
str(b): set([str(c), str(d)]),
str(d): set([str(e)]),
}
upstreams = source.get_upstream_tables(str(a), [])
assert list(upstreams).sort() == [c, e].sort()
def test_oracle_config():
base_config = {
"username": "******",
"password": "******",
"host_port": "host:1521",
}
config = OracleConfig.parse_obj({
**base_config,
"service_name": "svc01",
})
assert (config.get_sql_alchemy_url() ==
"oracle+cx_oracle://user:password@host:1521/?service_name=svc01")
with pytest.raises(ValueError):
config = OracleConfig.parse_obj({
**base_config,
"database": "db",
"service_name": "svc01",
})
with unittest.mock.patch(
"datahub.ingestion.source.sql.sql_common.SQLAlchemySource.get_workunits"
):
OracleSource.create(
{
**base_config,
"service_name": "svc01",
},
PipelineContext("test-oracle-config"),
).get_workunits()
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_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_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_kafka_source_configuration(self, mock_kafka):
ctx = PipelineContext(run_id="test")
kafka_source = KafkaSource.create(
{"connection": {"bootstrap": "foobar:9092"}}, ctx
)
kafka_source.close()
assert mock_kafka.call_count == 1
def test_platform_must_be_valid():
with pytest.raises(ConfigurationError):
GlueSource(
ctx=PipelineContext(run_id="glue-source-test"),
config=GlueSourceConfig(aws_region="us-west-2",
platform="data-warehouse"),
)
def test_underlying_platform_takes_precendence():
source = GlueSource(
ctx=PipelineContext(run_id="glue-source-test"),
config=GlueSourceConfig(aws_region="us-west-2",
underlying_platform="athena"),
)
assert source.platform == "athena"
def test_bq_usage_source(pytestconfig, tmp_path):
# from google.cloud.logging_v2 import ProtobufEntry
test_resources_dir: pathlib.Path = (pytestconfig.rootpath /
"tests/integration/bigquery-usage")
bigquery_reference_logs_path = test_resources_dir / "bigquery_logs.json"
if WRITE_REFERENCE_FILE:
source = BigQueryUsageSource.create(
dict(
projects=[
"harshal-playground-306419",
],
start_time=datetime.now(tz=timezone.utc) - timedelta(days=25),
),
PipelineContext(run_id="bq-usage-test"),
)
entries = list(
source._get_bigquery_log_entries_via_gcp_logging(
source._make_bigquery_logging_clients()))
entries = [entry._replace(logger=None) for entry in entries]
log_entries = jsonpickle.encode(entries, indent=4)
with bigquery_reference_logs_path.open("w") as logs:
logs.write(log_entries)
with unittest.mock.patch(
"datahub.ingestion.source.usage.bigquery_usage.GCPLoggingClient",
autospec=True) as MockClient:
# Add mock BigQuery API responses.
with bigquery_reference_logs_path.open() as logs:
reference_logs = jsonpickle.decode(logs.read())
MockClient().list_entries.return_value = reference_logs
# Run a BigQuery usage ingestion run.
pipeline = Pipeline.create({
"run_id": "test-bigquery-usage",
"source": {
"type": "bigquery-usage",
"config": {
"projects": ["sample-bigquery-project-1234"],
"start_time": "2021-01-01T00:00Z",
"end_time": "2021-07-01T00:00Z",
},
},
"sink": {
"type": "file",
"config": {
"filename": f"{tmp_path}/bigquery_usages.json",
},
},
})
pipeline.run()
pipeline.raise_from_status()
mce_helpers.check_golden_file(
pytestconfig,
output_path=tmp_path / "bigquery_usages.json",
golden_path=test_resources_dir / "bigquery_usages_golden.json",
)
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 __init__(self, config: PipelineConfig):
self.config = config
self.ctx = PipelineContext(run_id=self.config.run_id)
source_type = self.config.source.type
try:
source_class = source_class_mapping[source_type]
except KeyError as e:
raise ValueError(
f"Did not find a registered source class for {source_type}"
) from e
self.source: Source = source_class.create(
self.config.source.dict().get("config", {}), self.ctx)
logger.debug(f"Source type:{source_type},{source_class} configured")
sink_type = self.config.sink.type
try:
sink_class = sink_class_mapping[sink_type]
except KeyError as e:
raise ValueError(
f"Did not find a registered sink class for {sink_type}") from e
sink_config = self.config.sink.dict().get("config", {})
self.sink: Sink = sink_class.create(sink_config, self.ctx)
logger.debug(
f"Sink type:{self.config.sink.type},{sink_class} configured")
# Ensure extractor can be constructed, even though we use them later
self.extractor_class = self.get_class_from_name(
self.config.source.extractor)
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_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_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={})]))
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_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 redash_source() -> RedashSource:
return RedashSource(
ctx=PipelineContext(run_id="redash-source-test"),
config=RedashConfig(
connect_uri="http://localhost:5000",
api_key="REDASH_API_KEY",
),
)
def test_simple_datahub_extract(load_source):
catalog, conf, source = load_source
datahub_source = CatalogSource.create(
yaml.safe_load(conf)["catalog"],
PipelineContext(run_id="test_extract"))
datahub_source.config.source_names = [source.name]
assert len(list(datahub_source.get_workunits())) == 3
def _create_provider(self) -> IngestionReportingProviderBase:
ctx: PipelineContext = PipelineContext(
run_id=self.run_id, pipeline_name=self.pipeline_name
)
ctx.graph = self.mock_graph
return DatahubIngestionReportingProvider.create(
{}, ctx, name=DatahubIngestionReportingProvider.__name__
)
def redash_source_parse_table_names_from_sql() -> RedashSource:
return RedashSource(
ctx=PipelineContext(run_id="redash-source-test"),
config=RedashConfig(
connect_uri="http://localhost:5000",
api_key="REDASH_API_KEY",
parse_table_names_from_sql=True,
),
)
def glue_source(platform_instance: Optional[str] = None) -> GlueSource:
return GlueSource(
ctx=PipelineContext(run_id="glue-source-test"),
config=GlueSourceConfig(
aws_region="us-west-2",
extract_transforms=True,
platform_instance=platform_instance,
),
)
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_simple_dataset_ownership_with_invalid_type_transformation(mock_time):
with pytest.raises(ValueError):
SimpleAddDatasetOwnership.create(
{
"owner_urns": [
builder.make_user_urn("person1"),
],
"ownership_type": "INVALID_TYPE",
},
PipelineContext(run_id="test"),
)