def test_datahub_lineage_operator(mock_hook):
task = DatahubEmitterOperator(
task_id="emit_lineage",
datahub_rest_conn_id=datahub_rest_connection_config.conn_id,
mces=[
builder.make_lineage_mce(
[
builder.make_dataset_urn("snowflake", "mydb.schema.tableA"),
builder.make_dataset_urn("snowflake", "mydb.schema.tableB"),
],
builder.make_dataset_urn("snowflake", "mydb.schema.tableC"),
)
],
)
task.execute(None)
mock_hook.assert_called()
mock_hook.return_value.emit_mces.assert_called_once()
def test_datahub_lineage_operator(mock_emit):
with patch_airflow_connection(datahub_rest_connection_config) as config:
task = DatahubEmitterOperator(
task_id="emit_lineage",
datahub_conn_id=config.conn_id,
mces=[
builder.make_lineage_mce(
[
builder.make_dataset_urn("snowflake", "mydb.schema.tableA"),
builder.make_dataset_urn("snowflake", "mydb.schema.tableB"),
],
builder.make_dataset_urn("snowflake", "mydb.schema.tableC"),
)
],
)
task.execute(None)
mock_emit.assert_called()
),
some_other_table AS (
SELECT id, some_column FROM `mydb.schema.tableB`
)
SELECT * FROM some_table
LEFT JOIN some_other_table ON some_table.unique_id=some_other_table.id"""
transformation_task = SnowflakeOperator(
task_id="snowflake_transformation",
dag=dag,
snowflake_conn_id="snowflake_default",
sql=sql,
)
emit_lineage_task = DatahubEmitterOperator(
task_id="emit_lineage",
datahub_rest_conn_id="datahub_rest_default",
mces=[
builder.make_lineage_mce(
[
builder.make_dataset_urn("snowflake",
"mydb.schema.tableA"),
builder.make_dataset_urn("snowflake",
"mydb.schema.tableB"),
],
builder.make_dataset_urn("snowflake", "mydb.schema.tableC"),
)
],
)
transformation_task >> emit_lineage_task