2022-11-11 15:04:36 -05:00
|
|
|
from datahub_provider._airflow_compat import MARKUPSAFE_PATCHED
|
|
|
|
|
2022-04-13 08:19:39 +02:00
|
|
|
import datetime
|
2021-04-05 19:11:28 -07:00
|
|
|
import json
|
2021-04-12 17:40:15 -07:00
|
|
|
import os
|
2022-04-13 08:19:39 +02:00
|
|
|
import sys
|
2021-04-05 19:11:28 -07:00
|
|
|
from contextlib import contextmanager
|
|
|
|
from typing import Iterator
|
|
|
|
from unittest import mock
|
2022-04-13 08:19:39 +02:00
|
|
|
from unittest.mock import Mock
|
2021-04-05 19:11:28 -07:00
|
|
|
|
2021-04-12 17:40:15 -07:00
|
|
|
import airflow.configuration
|
2021-04-26 23:08:43 -07:00
|
|
|
import airflow.version
|
2022-02-19 21:13:01 -05:00
|
|
|
import packaging.version
|
2021-04-12 17:40:15 -07:00
|
|
|
import pytest
|
2021-04-26 23:08:43 -07:00
|
|
|
from airflow.lineage import apply_lineage, prepare_lineage
|
2022-04-13 08:19:39 +02:00
|
|
|
from airflow.models import DAG, Connection, DagBag, DagRun, TaskInstance
|
2022-11-11 15:04:36 -05:00
|
|
|
from airflow.operators.dummy import DummyOperator
|
2021-04-12 17:40:15 -07:00
|
|
|
from airflow.utils.dates import days_ago
|
|
|
|
|
2021-04-05 19:11:28 -07:00
|
|
|
import datahub.emitter.mce_builder as builder
|
2021-05-12 15:01:11 -07:00
|
|
|
from datahub_provider import get_provider_info
|
|
|
|
from datahub_provider.entities import Dataset
|
2022-11-01 21:12:34 -07:00
|
|
|
from datahub_provider.hooks.datahub import DatahubKafkaHook, DatahubRestHook
|
2021-05-12 15:01:11 -07:00
|
|
|
from datahub_provider.operators.datahub import DatahubEmitterOperator
|
2021-04-05 19:11:28 -07:00
|
|
|
|
2022-11-11 15:04:36 -05:00
|
|
|
assert MARKUPSAFE_PATCHED
|
|
|
|
|
|
|
|
pytestmark = pytest.mark.airflow
|
|
|
|
|
2022-02-19 21:13:01 -05:00
|
|
|
# Approach suggested by https://stackoverflow.com/a/11887885/5004662.
|
|
|
|
AIRFLOW_VERSION = packaging.version.parse(airflow.version.version)
|
|
|
|
|
2021-04-05 19:11:28 -07:00
|
|
|
lineage_mce = builder.make_lineage_mce(
|
|
|
|
[
|
|
|
|
builder.make_dataset_urn("bigquery", "upstream1"),
|
|
|
|
builder.make_dataset_urn("bigquery", "upstream2"),
|
|
|
|
],
|
|
|
|
builder.make_dataset_urn("bigquery", "downstream1"),
|
|
|
|
)
|
|
|
|
|
|
|
|
datahub_rest_connection_config = Connection(
|
|
|
|
conn_id="datahub_rest_test",
|
|
|
|
conn_type="datahub_rest",
|
|
|
|
host="http://test_host:8080/",
|
|
|
|
extra=None,
|
|
|
|
)
|
2021-08-08 22:30:55 -07:00
|
|
|
datahub_rest_connection_config_with_timeout = Connection(
|
|
|
|
conn_id="datahub_rest_test",
|
|
|
|
conn_type="datahub_rest",
|
|
|
|
host="http://test_host:8080/",
|
|
|
|
extra=json.dumps({"timeout_sec": 5}),
|
|
|
|
)
|
|
|
|
|
2021-04-05 19:11:28 -07:00
|
|
|
datahub_kafka_connection_config = Connection(
|
|
|
|
conn_id="datahub_kafka_test",
|
|
|
|
conn_type="datahub_kafka",
|
|
|
|
host="test_broker:9092",
|
|
|
|
extra=json.dumps(
|
|
|
|
{
|
|
|
|
"connection": {
|
|
|
|
"producer_config": {},
|
|
|
|
"schema_registry_url": "http://localhost:8081",
|
|
|
|
}
|
|
|
|
}
|
|
|
|
),
|
|
|
|
)
|
|
|
|
|
|
|
|
|
2021-04-12 17:40:15 -07:00
|
|
|
def setup_module(module):
|
2021-04-20 20:44:38 -07:00
|
|
|
airflow.configuration.conf.load_test_config()
|
2021-04-12 17:40:15 -07:00
|
|
|
|
|
|
|
|
2021-04-05 19:11:28 -07:00
|
|
|
def test_airflow_provider_info():
|
|
|
|
assert get_provider_info()
|
|
|
|
|
|
|
|
|
2022-11-11 15:04:36 -05:00
|
|
|
@pytest.mark.skipif(
|
|
|
|
AIRFLOW_VERSION < packaging.version.parse("2.0.0"),
|
|
|
|
reason="the examples use list-style lineage, which is only supported on Airflow 2.x",
|
|
|
|
)
|
2021-04-05 19:11:28 -07:00
|
|
|
def test_dags_load_with_no_errors(pytestconfig):
|
2021-05-12 15:01:11 -07:00
|
|
|
airflow_examples_folder = (
|
|
|
|
pytestconfig.rootpath / "src/datahub_provider/example_dags"
|
|
|
|
)
|
2021-04-05 19:11:28 -07:00
|
|
|
|
|
|
|
dag_bag = DagBag(dag_folder=str(airflow_examples_folder), include_examples=False)
|
2021-07-26 13:09:25 -07:00
|
|
|
|
|
|
|
import_errors = dag_bag.import_errors
|
|
|
|
|
|
|
|
assert import_errors == {}
|
2021-04-05 19:11:28 -07:00
|
|
|
assert len(dag_bag.dag_ids) > 0
|
|
|
|
|
|
|
|
|
|
|
|
@contextmanager
|
|
|
|
def patch_airflow_connection(conn: Connection) -> Iterator[Connection]:
|
|
|
|
# The return type should really by ContextManager, but mypy doesn't like that.
|
|
|
|
# See https://stackoverflow.com/questions/49733699/python-type-hints-and-context-managers#comment106444758_58349659.
|
2021-04-12 17:40:15 -07:00
|
|
|
with mock.patch(
|
2021-05-12 15:01:11 -07:00
|
|
|
"datahub_provider.hooks.datahub.BaseHook.get_connection", return_value=conn
|
2021-04-12 17:40:15 -07:00
|
|
|
):
|
2021-04-05 19:11:28 -07:00
|
|
|
yield conn
|
|
|
|
|
|
|
|
|
2021-07-28 14:23:06 -07:00
|
|
|
@mock.patch("datahub.emitter.rest_emitter.DatahubRestEmitter", autospec=True)
|
2021-04-05 19:11:28 -07:00
|
|
|
def test_datahub_rest_hook(mock_emitter):
|
|
|
|
with patch_airflow_connection(datahub_rest_connection_config) as config:
|
2022-11-11 15:04:36 -05:00
|
|
|
assert config.conn_id
|
2021-04-05 19:11:28 -07:00
|
|
|
hook = DatahubRestHook(config.conn_id)
|
|
|
|
hook.emit_mces([lineage_mce])
|
|
|
|
|
2021-08-08 22:30:55 -07:00
|
|
|
mock_emitter.assert_called_once_with(config.host, None, None)
|
|
|
|
instance = mock_emitter.return_value
|
|
|
|
instance.emit_mce.assert_called_with(lineage_mce)
|
|
|
|
|
|
|
|
|
|
|
|
@mock.patch("datahub.emitter.rest_emitter.DatahubRestEmitter", autospec=True)
|
|
|
|
def test_datahub_rest_hook_with_timeout(mock_emitter):
|
|
|
|
with patch_airflow_connection(
|
|
|
|
datahub_rest_connection_config_with_timeout
|
|
|
|
) as config:
|
2022-11-11 15:04:36 -05:00
|
|
|
assert config.conn_id
|
2021-08-08 22:30:55 -07:00
|
|
|
hook = DatahubRestHook(config.conn_id)
|
|
|
|
hook.emit_mces([lineage_mce])
|
|
|
|
|
|
|
|
mock_emitter.assert_called_once_with(config.host, None, 5)
|
2021-04-05 19:11:28 -07:00
|
|
|
instance = mock_emitter.return_value
|
|
|
|
instance.emit_mce.assert_called_with(lineage_mce)
|
|
|
|
|
|
|
|
|
2021-07-28 14:23:06 -07:00
|
|
|
@mock.patch("datahub.emitter.kafka_emitter.DatahubKafkaEmitter", autospec=True)
|
2021-04-05 19:11:28 -07:00
|
|
|
def test_datahub_kafka_hook(mock_emitter):
|
|
|
|
with patch_airflow_connection(datahub_kafka_connection_config) as config:
|
2022-11-11 15:04:36 -05:00
|
|
|
assert config.conn_id
|
2021-04-05 19:11:28 -07:00
|
|
|
hook = DatahubKafkaHook(config.conn_id)
|
|
|
|
hook.emit_mces([lineage_mce])
|
|
|
|
|
|
|
|
mock_emitter.assert_called_once()
|
|
|
|
instance = mock_emitter.return_value
|
|
|
|
instance.emit_mce_async.assert_called()
|
|
|
|
instance.flush.assert_called_once()
|
|
|
|
|
|
|
|
|
2021-05-12 15:01:11 -07:00
|
|
|
@mock.patch("datahub_provider.hooks.datahub.DatahubRestHook.emit_mces")
|
2021-04-12 17:40:15 -07:00
|
|
|
def test_datahub_lineage_operator(mock_emit):
|
|
|
|
with patch_airflow_connection(datahub_rest_connection_config) as config:
|
2022-11-11 15:04:36 -05:00
|
|
|
assert config.conn_id
|
2021-04-12 17:40:15 -07:00
|
|
|
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()
|
|
|
|
|
|
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
|
|
"hook",
|
|
|
|
[
|
|
|
|
DatahubRestHook,
|
|
|
|
DatahubKafkaHook,
|
|
|
|
],
|
|
|
|
)
|
|
|
|
def test_hook_airflow_ui(hook):
|
|
|
|
# Simply ensure that these run without issue. These will also show up
|
|
|
|
# in the Airflow UI, where it will be even more clear if something
|
|
|
|
# is wrong.
|
|
|
|
hook.get_connection_form_widgets()
|
|
|
|
hook.get_ui_field_behaviour()
|
|
|
|
|
|
|
|
|
2021-04-26 23:08:43 -07:00
|
|
|
@pytest.mark.parametrize(
|
2021-04-29 23:18:55 -07:00
|
|
|
["inlets", "outlets"],
|
2021-04-26 23:08:43 -07:00
|
|
|
[
|
2021-10-13 02:25:27 -04:00
|
|
|
pytest.param(
|
2022-11-01 21:12:34 -07:00
|
|
|
# Airflow 2.x uses a flattened list for inlets and outlets.
|
2021-04-26 23:08:43 -07:00
|
|
|
[Dataset("snowflake", "mydb.schema.tableConsumed")],
|
|
|
|
[Dataset("snowflake", "mydb.schema.tableProduced")],
|
2021-10-13 02:25:27 -04:00
|
|
|
id="airflow-2-lineage-syntax",
|
2021-04-26 23:08:43 -07:00
|
|
|
),
|
|
|
|
],
|
|
|
|
)
|
2022-04-13 08:19:39 +02:00
|
|
|
@mock.patch("datahub_provider.hooks.datahub.DatahubRestHook.make_emitter")
|
2021-05-07 15:21:43 -07:00
|
|
|
def test_lineage_backend(mock_emit, inlets, outlets):
|
2021-04-12 17:40:15 -07:00
|
|
|
DEFAULT_DATE = days_ago(2)
|
2022-04-13 08:19:39 +02:00
|
|
|
mock_emitter = Mock()
|
|
|
|
mock_emit.return_value = mock_emitter
|
|
|
|
# Using autospec on xcom_pull and xcom_push methods fails on Python 3.6.
|
|
|
|
with mock.patch.dict(
|
|
|
|
os.environ,
|
|
|
|
{
|
|
|
|
"AIRFLOW__LINEAGE__BACKEND": "datahub_provider.lineage.datahub.DatahubLineageBackend",
|
|
|
|
"AIRFLOW__LINEAGE__DATAHUB_CONN_ID": datahub_rest_connection_config.conn_id,
|
|
|
|
"AIRFLOW__LINEAGE__DATAHUB_KWARGS": json.dumps(
|
|
|
|
{"graceful_exceptions": False, "capture_executions": False}
|
|
|
|
),
|
|
|
|
},
|
|
|
|
), mock.patch("airflow.models.BaseOperator.xcom_pull"), mock.patch(
|
|
|
|
"airflow.models.BaseOperator.xcom_push"
|
|
|
|
), patch_airflow_connection(
|
|
|
|
datahub_rest_connection_config
|
|
|
|
):
|
|
|
|
func = mock.Mock()
|
|
|
|
func.__name__ = "foo"
|
|
|
|
|
|
|
|
dag = DAG(dag_id="test_lineage_is_sent_to_backend", start_date=DEFAULT_DATE)
|
|
|
|
|
|
|
|
with dag:
|
|
|
|
op1 = DummyOperator(
|
|
|
|
task_id="task1_upstream",
|
|
|
|
inlets=inlets,
|
|
|
|
outlets=outlets,
|
|
|
|
)
|
|
|
|
op2 = DummyOperator(
|
|
|
|
task_id="task2",
|
|
|
|
inlets=inlets,
|
|
|
|
outlets=outlets,
|
|
|
|
)
|
|
|
|
op1 >> op2
|
|
|
|
|
|
|
|
# Airflow < 2.2 requires the execution_date parameter. Newer Airflow
|
|
|
|
# versions do not require it, but will attempt to find the associated
|
|
|
|
# run_id in the database if execution_date is provided. As such, we
|
|
|
|
# must fake the run_id parameter for newer Airflow versions.
|
|
|
|
if AIRFLOW_VERSION < packaging.version.parse("2.2.0"):
|
|
|
|
ti = TaskInstance(task=op2, execution_date=DEFAULT_DATE)
|
|
|
|
else:
|
|
|
|
ti = TaskInstance(task=op2, run_id=f"test_airflow-{DEFAULT_DATE}")
|
|
|
|
ctx1 = {
|
|
|
|
"dag": dag,
|
|
|
|
"task": op2,
|
|
|
|
"ti": ti,
|
|
|
|
"task_instance": ti,
|
|
|
|
"execution_date": DEFAULT_DATE,
|
|
|
|
"ts": "2021-04-08T00:54:25.771575+00:00",
|
|
|
|
}
|
|
|
|
|
|
|
|
prep = prepare_lineage(func)
|
|
|
|
prep(op2, ctx1)
|
|
|
|
post = apply_lineage(func)
|
|
|
|
post(op2, ctx1)
|
|
|
|
|
|
|
|
# Verify that the inlets and outlets are registered and recognized by Airflow correctly,
|
|
|
|
# or that our lineage backend forces it to.
|
|
|
|
assert len(op2.inlets) == 1
|
|
|
|
assert len(op2.outlets) == 1
|
|
|
|
assert all(map(lambda let: isinstance(let, Dataset), op2.inlets))
|
|
|
|
assert all(map(lambda let: isinstance(let, Dataset), op2.outlets))
|
|
|
|
|
|
|
|
# Check that the right things were emitted.
|
|
|
|
assert mock_emitter.emit.call_count == 9
|
2022-08-10 22:00:31 +00:00
|
|
|
# Running further checks based on python version because args only exists in python 3.8+
|
|
|
|
if sys.version_info > (3, 8):
|
2022-04-13 08:19:39 +02:00
|
|
|
assert mock_emitter.method_calls[0].args[0].aspectName == "dataFlowInfo"
|
|
|
|
assert (
|
|
|
|
mock_emitter.method_calls[0].args[0].entityUrn
|
|
|
|
== "urn:li:dataFlow:(airflow,test_lineage_is_sent_to_backend,prod)"
|
|
|
|
)
|
|
|
|
|
|
|
|
assert mock_emitter.method_calls[1].args[0].aspectName == "ownership"
|
|
|
|
assert (
|
|
|
|
mock_emitter.method_calls[1].args[0].entityUrn
|
|
|
|
== "urn:li:dataFlow:(airflow,test_lineage_is_sent_to_backend,prod)"
|
|
|
|
)
|
|
|
|
|
|
|
|
assert mock_emitter.method_calls[2].args[0].aspectName == "globalTags"
|
|
|
|
assert (
|
|
|
|
mock_emitter.method_calls[2].args[0].entityUrn
|
|
|
|
== "urn:li:dataFlow:(airflow,test_lineage_is_sent_to_backend,prod)"
|
|
|
|
)
|
|
|
|
|
|
|
|
assert mock_emitter.method_calls[3].args[0].aspectName == "dataJobInfo"
|
|
|
|
assert (
|
|
|
|
mock_emitter.method_calls[3].args[0].entityUrn
|
|
|
|
== "urn:li:dataJob:(urn:li:dataFlow:(airflow,test_lineage_is_sent_to_backend,prod),task2)"
|
|
|
|
)
|
|
|
|
|
|
|
|
assert (
|
|
|
|
mock_emitter.method_calls[4].args[0].aspectName == "dataJobInputOutput"
|
|
|
|
)
|
|
|
|
assert (
|
|
|
|
mock_emitter.method_calls[4].args[0].entityUrn
|
|
|
|
== "urn:li:dataJob:(urn:li:dataFlow:(airflow,test_lineage_is_sent_to_backend,prod),task2)"
|
|
|
|
)
|
|
|
|
assert (
|
|
|
|
mock_emitter.method_calls[4].args[0].aspect.inputDatajobs[0]
|
|
|
|
== "urn:li:dataJob:(urn:li:dataFlow:(airflow,test_lineage_is_sent_to_backend,prod),task1_upstream)"
|
|
|
|
)
|
|
|
|
assert (
|
|
|
|
mock_emitter.method_calls[4].args[0].aspect.inputDatasets[0]
|
|
|
|
== "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableConsumed,PROD)"
|
|
|
|
)
|
|
|
|
assert (
|
|
|
|
mock_emitter.method_calls[4].args[0].aspect.outputDatasets[0]
|
|
|
|
== "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableProduced,PROD)"
|
|
|
|
)
|
|
|
|
|
|
|
|
assert mock_emitter.method_calls[5].args[0].aspectName == "status"
|
|
|
|
assert (
|
|
|
|
mock_emitter.method_calls[5].args[0].entityUrn
|
|
|
|
== "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableConsumed,PROD)"
|
|
|
|
)
|
|
|
|
|
|
|
|
assert mock_emitter.method_calls[6].args[0].aspectName == "status"
|
|
|
|
assert (
|
|
|
|
mock_emitter.method_calls[6].args[0].entityUrn
|
|
|
|
== "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableProduced,PROD)"
|
|
|
|
)
|
|
|
|
|
|
|
|
assert mock_emitter.method_calls[7].args[0].aspectName == "ownership"
|
|
|
|
assert (
|
|
|
|
mock_emitter.method_calls[7].args[0].entityUrn
|
|
|
|
== "urn:li:dataJob:(urn:li:dataFlow:(airflow,test_lineage_is_sent_to_backend,prod),task2)"
|
|
|
|
)
|
|
|
|
|
|
|
|
assert mock_emitter.method_calls[8].args[0].aspectName == "globalTags"
|
|
|
|
assert (
|
|
|
|
mock_emitter.method_calls[8].args[0].entityUrn
|
|
|
|
== "urn:li:dataJob:(urn:li:dataFlow:(airflow,test_lineage_is_sent_to_backend,prod),task2)"
|
|
|
|
)
|
|
|
|
|
2021-04-12 17:40:15 -07:00
|
|
|
|
2022-04-13 08:19:39 +02:00
|
|
|
@pytest.mark.parametrize(
|
|
|
|
["inlets", "outlets"],
|
|
|
|
[
|
|
|
|
pytest.param(
|
|
|
|
[Dataset("snowflake", "mydb.schema.tableConsumed")],
|
|
|
|
[Dataset("snowflake", "mydb.schema.tableProduced")],
|
|
|
|
id="airflow-2-lineage-syntax",
|
|
|
|
),
|
|
|
|
],
|
|
|
|
)
|
|
|
|
@mock.patch("datahub_provider.hooks.datahub.DatahubRestHook.make_emitter")
|
|
|
|
def test_lineage_backend_capture_executions(mock_emit, inlets, outlets):
|
2022-11-11 15:04:36 -05:00
|
|
|
# TODO: Merge this code into the test above to reduce duplication.
|
2022-04-13 08:19:39 +02:00
|
|
|
DEFAULT_DATE = datetime.datetime(2020, 5, 17)
|
|
|
|
mock_emitter = Mock()
|
|
|
|
mock_emit.return_value = mock_emitter
|
2021-07-23 17:07:13 -07:00
|
|
|
# Using autospec on xcom_pull and xcom_push methods fails on Python 3.6.
|
2021-04-12 17:40:15 -07:00
|
|
|
with mock.patch.dict(
|
|
|
|
os.environ,
|
|
|
|
{
|
2021-05-12 15:01:11 -07:00
|
|
|
"AIRFLOW__LINEAGE__BACKEND": "datahub_provider.lineage.datahub.DatahubLineageBackend",
|
2021-04-12 17:40:15 -07:00
|
|
|
"AIRFLOW__LINEAGE__DATAHUB_CONN_ID": datahub_rest_connection_config.conn_id,
|
2021-05-13 20:02:47 -07:00
|
|
|
"AIRFLOW__LINEAGE__DATAHUB_KWARGS": json.dumps(
|
2022-04-13 08:19:39 +02:00
|
|
|
{"graceful_exceptions": False, "capture_executions": True}
|
2021-05-13 20:02:47 -07:00
|
|
|
),
|
2021-04-12 17:40:15 -07:00
|
|
|
},
|
2021-07-23 17:07:13 -07:00
|
|
|
), mock.patch("airflow.models.BaseOperator.xcom_pull"), mock.patch(
|
|
|
|
"airflow.models.BaseOperator.xcom_push"
|
2021-05-07 15:21:43 -07:00
|
|
|
), patch_airflow_connection(
|
|
|
|
datahub_rest_connection_config
|
|
|
|
):
|
2021-04-12 17:40:15 -07:00
|
|
|
func = mock.Mock()
|
|
|
|
func.__name__ = "foo"
|
|
|
|
|
|
|
|
dag = DAG(dag_id="test_lineage_is_sent_to_backend", start_date=DEFAULT_DATE)
|
|
|
|
|
|
|
|
with dag:
|
2021-04-26 23:08:43 -07:00
|
|
|
op1 = DummyOperator(
|
2021-05-25 22:47:00 -07:00
|
|
|
task_id="task1_upstream",
|
2021-04-26 23:08:43 -07:00
|
|
|
inlets=inlets,
|
|
|
|
outlets=outlets,
|
|
|
|
)
|
2021-05-25 22:47:00 -07:00
|
|
|
op2 = DummyOperator(
|
|
|
|
task_id="task2",
|
|
|
|
inlets=inlets,
|
|
|
|
outlets=outlets,
|
|
|
|
)
|
|
|
|
op1 >> op2
|
2021-04-12 17:40:15 -07:00
|
|
|
|
2022-02-19 21:13:01 -05:00
|
|
|
# Airflow < 2.2 requires the execution_date parameter. Newer Airflow
|
2021-10-13 02:25:27 -04:00
|
|
|
# versions do not require it, but will attempt to find the associated
|
|
|
|
# run_id in the database if execution_date is provided. As such, we
|
|
|
|
# must fake the run_id parameter for newer Airflow versions.
|
2022-02-19 21:13:01 -05:00
|
|
|
if AIRFLOW_VERSION < packaging.version.parse("2.2.0"):
|
2021-10-13 02:25:27 -04:00
|
|
|
ti = TaskInstance(task=op2, execution_date=DEFAULT_DATE)
|
2022-04-13 08:19:39 +02:00
|
|
|
# Ignoring type here because DagRun state is just a sring at Airflow 1
|
|
|
|
dag_run = DagRun(state="success", run_id=f"scheduled_{DEFAULT_DATE}") # type: ignore
|
2021-10-11 20:58:37 -07:00
|
|
|
else:
|
2022-04-13 08:19:39 +02:00
|
|
|
from airflow.utils.state import DagRunState
|
|
|
|
|
2021-10-13 02:25:27 -04:00
|
|
|
ti = TaskInstance(task=op2, run_id=f"test_airflow-{DEFAULT_DATE}")
|
2022-04-13 08:19:39 +02:00
|
|
|
dag_run = DagRun(
|
|
|
|
state=DagRunState.SUCCESS, run_id=f"scheduled_{DEFAULT_DATE}"
|
|
|
|
)
|
2022-11-11 15:04:36 -05:00
|
|
|
|
|
|
|
ti.dag_run = dag_run # type: ignore
|
|
|
|
ti.start_date = datetime.datetime.utcnow()
|
|
|
|
ti.execution_date = DEFAULT_DATE
|
2022-04-13 08:19:39 +02:00
|
|
|
|
2021-04-12 17:40:15 -07:00
|
|
|
ctx1 = {
|
|
|
|
"dag": dag,
|
2021-05-25 22:47:00 -07:00
|
|
|
"task": op2,
|
2021-04-12 17:40:15 -07:00
|
|
|
"ti": ti,
|
2022-04-13 08:19:39 +02:00
|
|
|
"dag_run": dag_run,
|
2021-04-12 17:40:15 -07:00
|
|
|
"task_instance": ti,
|
|
|
|
"execution_date": DEFAULT_DATE,
|
|
|
|
"ts": "2021-04-08T00:54:25.771575+00:00",
|
|
|
|
}
|
|
|
|
|
|
|
|
prep = prepare_lineage(func)
|
2021-05-25 22:47:00 -07:00
|
|
|
prep(op2, ctx1)
|
2021-04-12 17:40:15 -07:00
|
|
|
post = apply_lineage(func)
|
2021-05-25 22:47:00 -07:00
|
|
|
post(op2, ctx1)
|
2021-04-12 17:40:15 -07:00
|
|
|
|
2021-04-26 23:08:43 -07:00
|
|
|
# Verify that the inlets and outlets are registered and recognized by Airflow correctly,
|
|
|
|
# or that our lineage backend forces it to.
|
2021-05-25 22:47:00 -07:00
|
|
|
assert len(op2.inlets) == 1
|
|
|
|
assert len(op2.outlets) == 1
|
|
|
|
assert all(map(lambda let: isinstance(let, Dataset), op2.inlets))
|
|
|
|
assert all(map(lambda let: isinstance(let, Dataset), op2.outlets))
|
2021-04-26 23:08:43 -07:00
|
|
|
|
|
|
|
# Check that the right things were emitted.
|
2022-04-13 08:19:39 +02:00
|
|
|
assert mock_emitter.emit.call_count == 17
|
2022-08-10 22:00:31 +00:00
|
|
|
# Running further checks based on python version because args only exists in python 3.8+
|
|
|
|
if sys.version_info > (3, 8):
|
2022-04-13 08:19:39 +02:00
|
|
|
assert mock_emitter.method_calls[0].args[0].aspectName == "dataFlowInfo"
|
|
|
|
assert (
|
|
|
|
mock_emitter.method_calls[0].args[0].entityUrn
|
|
|
|
== "urn:li:dataFlow:(airflow,test_lineage_is_sent_to_backend,prod)"
|
|
|
|
)
|
|
|
|
|
|
|
|
assert mock_emitter.method_calls[1].args[0].aspectName == "ownership"
|
|
|
|
assert (
|
|
|
|
mock_emitter.method_calls[1].args[0].entityUrn
|
|
|
|
== "urn:li:dataFlow:(airflow,test_lineage_is_sent_to_backend,prod)"
|
|
|
|
)
|
|
|
|
|
|
|
|
assert mock_emitter.method_calls[2].args[0].aspectName == "globalTags"
|
|
|
|
assert (
|
|
|
|
mock_emitter.method_calls[2].args[0].entityUrn
|
|
|
|
== "urn:li:dataFlow:(airflow,test_lineage_is_sent_to_backend,prod)"
|
|
|
|
)
|
|
|
|
|
|
|
|
assert mock_emitter.method_calls[3].args[0].aspectName == "dataJobInfo"
|
|
|
|
assert (
|
|
|
|
mock_emitter.method_calls[3].args[0].entityUrn
|
|
|
|
== "urn:li:dataJob:(urn:li:dataFlow:(airflow,test_lineage_is_sent_to_backend,prod),task2)"
|
|
|
|
)
|
|
|
|
|
|
|
|
assert (
|
|
|
|
mock_emitter.method_calls[4].args[0].aspectName == "dataJobInputOutput"
|
|
|
|
)
|
|
|
|
assert (
|
|
|
|
mock_emitter.method_calls[4].args[0].entityUrn
|
|
|
|
== "urn:li:dataJob:(urn:li:dataFlow:(airflow,test_lineage_is_sent_to_backend,prod),task2)"
|
|
|
|
)
|
|
|
|
assert (
|
|
|
|
mock_emitter.method_calls[4].args[0].aspect.inputDatajobs[0]
|
|
|
|
== "urn:li:dataJob:(urn:li:dataFlow:(airflow,test_lineage_is_sent_to_backend,prod),task1_upstream)"
|
|
|
|
)
|
|
|
|
assert (
|
|
|
|
mock_emitter.method_calls[4].args[0].aspect.inputDatasets[0]
|
|
|
|
== "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableConsumed,PROD)"
|
|
|
|
)
|
|
|
|
assert (
|
|
|
|
mock_emitter.method_calls[4].args[0].aspect.outputDatasets[0]
|
|
|
|
== "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableProduced,PROD)"
|
|
|
|
)
|
|
|
|
|
|
|
|
assert mock_emitter.method_calls[5].args[0].aspectName == "status"
|
|
|
|
assert (
|
|
|
|
mock_emitter.method_calls[5].args[0].entityUrn
|
|
|
|
== "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableConsumed,PROD)"
|
|
|
|
)
|
|
|
|
|
|
|
|
assert mock_emitter.method_calls[6].args[0].aspectName == "status"
|
|
|
|
assert (
|
|
|
|
mock_emitter.method_calls[6].args[0].entityUrn
|
|
|
|
== "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableProduced,PROD)"
|
|
|
|
)
|
|
|
|
|
|
|
|
assert mock_emitter.method_calls[7].args[0].aspectName == "ownership"
|
|
|
|
assert (
|
|
|
|
mock_emitter.method_calls[7].args[0].entityUrn
|
|
|
|
== "urn:li:dataJob:(urn:li:dataFlow:(airflow,test_lineage_is_sent_to_backend,prod),task2)"
|
|
|
|
)
|
|
|
|
|
|
|
|
assert mock_emitter.method_calls[8].args[0].aspectName == "globalTags"
|
|
|
|
assert (
|
|
|
|
mock_emitter.method_calls[8].args[0].entityUrn
|
|
|
|
== "urn:li:dataJob:(urn:li:dataFlow:(airflow,test_lineage_is_sent_to_backend,prod),task2)"
|
|
|
|
)
|
|
|
|
|
|
|
|
assert (
|
|
|
|
mock_emitter.method_calls[9].args[0].aspectName
|
|
|
|
== "dataProcessInstanceProperties"
|
|
|
|
)
|
|
|
|
assert (
|
|
|
|
mock_emitter.method_calls[9].args[0].entityUrn
|
|
|
|
== "urn:li:dataProcessInstance:b6375e5f5faeb543cfb5d7d8a47661fb"
|
|
|
|
)
|
|
|
|
|
|
|
|
assert (
|
|
|
|
mock_emitter.method_calls[10].args[0].aspectName
|
|
|
|
== "dataProcessInstanceRelationships"
|
|
|
|
)
|
|
|
|
assert (
|
|
|
|
mock_emitter.method_calls[10].args[0].entityUrn
|
|
|
|
== "urn:li:dataProcessInstance:b6375e5f5faeb543cfb5d7d8a47661fb"
|
|
|
|
)
|
|
|
|
assert (
|
|
|
|
mock_emitter.method_calls[11].args[0].aspectName
|
|
|
|
== "dataProcessInstanceInput"
|
|
|
|
)
|
|
|
|
assert (
|
|
|
|
mock_emitter.method_calls[11].args[0].entityUrn
|
|
|
|
== "urn:li:dataProcessInstance:b6375e5f5faeb543cfb5d7d8a47661fb"
|
|
|
|
)
|
|
|
|
assert (
|
|
|
|
mock_emitter.method_calls[12].args[0].aspectName
|
|
|
|
== "dataProcessInstanceOutput"
|
|
|
|
)
|
|
|
|
assert (
|
|
|
|
mock_emitter.method_calls[12].args[0].entityUrn
|
|
|
|
== "urn:li:dataProcessInstance:b6375e5f5faeb543cfb5d7d8a47661fb"
|
|
|
|
)
|
|
|
|
assert mock_emitter.method_calls[13].args[0].aspectName == "status"
|
|
|
|
assert (
|
|
|
|
mock_emitter.method_calls[13].args[0].entityUrn
|
|
|
|
== "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableConsumed,PROD)"
|
|
|
|
)
|
|
|
|
assert mock_emitter.method_calls[14].args[0].aspectName == "status"
|
|
|
|
assert (
|
|
|
|
mock_emitter.method_calls[14].args[0].entityUrn
|
|
|
|
== "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableProduced,PROD)"
|
|
|
|
)
|
|
|
|
assert (
|
|
|
|
mock_emitter.method_calls[15].args[0].aspectName
|
|
|
|
== "dataProcessInstanceRunEvent"
|
|
|
|
)
|
|
|
|
assert (
|
|
|
|
mock_emitter.method_calls[15].args[0].entityUrn
|
|
|
|
== "urn:li:dataProcessInstance:b6375e5f5faeb543cfb5d7d8a47661fb"
|
|
|
|
)
|
|
|
|
assert (
|
|
|
|
mock_emitter.method_calls[16].args[0].aspectName
|
|
|
|
== "dataProcessInstanceRunEvent"
|
|
|
|
)
|
|
|
|
assert (
|
|
|
|
mock_emitter.method_calls[16].args[0].entityUrn
|
|
|
|
== "urn:li:dataProcessInstance:b6375e5f5faeb543cfb5d7d8a47661fb"
|
|
|
|
)
|