Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

feat(airflow): retry callback, support ExternalTaskSensor subclasses #8514

Merged
merged 11 commits into from
Oct 27, 2023
Merged

feat(airflow): retry callback, support ExternalTaskSensor subclasses #8514

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
92fb5c5
feat(airflow): retry callback, support ExternalTaskSensor subclasses,…
Oct 17, 2023
4d2de66
Merge branch 'master' into richenc-airflow-plugin-improvements
richenc Oct 17, 2023
3f3fb60
Merge branch 'master' into richenc-airflow-plugin-improvements
richenc Oct 19, 2023
0cdb288
Merge branch 'master' into richenc-airflow-plugin-improvements
richenc Oct 20, 2023
66fa26c
Merge branch 'master' into richenc-airflow-plugin-improvements
richenc Oct 24, 2023
bc14134
fix formatting
Oct 24, 2023
432d510
Merge branch 'master' into richenc-airflow-plugin-improvements
richenc Oct 24, 2023
02de933
Merge branch 'master' into richenc-airflow-plugin-improvements
richenc Oct 24, 2023
6cd58f3
update integration tests
Oct 25, 2023
ac58eef
Merge branch 'master' into richenc-airflow-plugin-improvements
richenc Oct 25, 2023
bb25ce7
Merge branch 'master' into richenc-airflow-plugin-improvements
hsheth2 Oct 27, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 6 additions & 1 deletion ...a-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/client/airflow_generator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -98,7 +98,7 @@ def _get_dependencies(
# It is possible to tie an external sensor to DAG if external_task_id is omitted but currently we can't tie
# jobflow to anothet jobflow.
external_task_upstreams = []
if task.task_type == "ExternalTaskSensor":
if isinstance(task, ExternalTaskSensor):
task = cast(ExternalTaskSensor, task)
if hasattr(task, "external_task_id") and task.external_task_id is not None:
external_task_upstreams = [
Expand Down Expand Up @@ -155,6 +155,8 @@ def generate_dataflow(
"_concurrency",
# "_default_view",
"catchup",
"description",
"doc_md",
"fileloc",
"is_paused_upon_creation",
"start_date",
Expand Down Expand Up @@ -431,6 +433,9 @@ def run_datajob(
job_property_bag["operator"] = str(ti.operator)
job_property_bag["priority_weight"] = str(ti.priority_weight)
job_property_bag["log_url"] = ti.log_url
job_property_bag["orchestrator"] = "airflow"
job_property_bag["dag_id"] = str(dag.dag_id)
job_property_bag["task_id"] = str(ti.task_id)
dpi.properties.update(job_property_bag)
dpi.url = ti.log_url

Expand Down
36 changes: 35 additions & 1 deletion metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/datahub_plugin_v22.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@

TASK_ON_FAILURE_CALLBACK = "on_failure_callback"
TASK_ON_SUCCESS_CALLBACK = "on_success_callback"
TASK_ON_RETRY_CALLBACK = "on_retry_callback"


def get_task_inlets_advanced(task: BaseOperator, context: Any) -> Iterable[Any]:
Expand Down Expand Up @@ -259,6 +260,28 @@ def custom_on_success_callback(context):
return custom_on_success_callback


def _wrap_on_retry_callback(on_retry_callback):
def custom_on_retry_callback(context):
config = get_lineage_config()
if config.enabled:
context["_datahub_config"] = config
try:
datahub_task_status_callback(
context, status=InstanceRunResult.UP_FOR_RETRY
)
except Exception as e:
if not config.graceful_exceptions:
raise e
else:
print(f"Exception: {traceback.format_exc()}")

# Call original policy
if on_retry_callback:
on_retry_callback(context)

return custom_on_retry_callback


def task_policy(task: Union[BaseOperator, MappedOperator]) -> None:
task.log.debug(f"Setting task policy for Dag: {task.dag_id} Task: {task.task_id}")
# task.add_inlets(["auto"])
Expand All @@ -274,7 +297,14 @@ def task_policy(task: Union[BaseOperator, MappedOperator]) -> None:
on_success_callback_prop: property = getattr(
MappedOperator, TASK_ON_SUCCESS_CALLBACK
)
if not on_failure_callback_prop.fset or not on_success_callback_prop.fset:
on_retry_callback_prop: property = getattr(
MappedOperator, TASK_ON_RETRY_CALLBACK
)
if (
not on_failure_callback_prop.fset
or not on_success_callback_prop.fset
or not on_retry_callback_prop.fset
):
task.log.debug(
"Using MappedOperator's partial_kwargs instead of callback properties"
)
Expand All @@ -284,10 +314,14 @@ def task_policy(task: Union[BaseOperator, MappedOperator]) -> None:
task.partial_kwargs[TASK_ON_SUCCESS_CALLBACK] = _wrap_on_success_callback(
task.on_success_callback
)
task.partial_kwargs[TASK_ON_RETRY_CALLBACK] = _wrap_on_retry_callback(
task.on_retry_callback
)
return

task.on_failure_callback = _wrap_on_failure_callback(task.on_failure_callback) # type: ignore
task.on_success_callback = _wrap_on_success_callback(task.on_success_callback) # type: ignore
task.on_retry_callback = _wrap_on_retry_callback(task.on_retry_callback) # type: ignore
# task.pre_execute = _wrap_pre_execution(task.pre_execute)


Expand Down
7 changes: 6 additions & 1 deletion metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v1_basic_iolets.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,8 @@
"customProperties": {
"_access_control": "None",
"catchup": "False",
"description": "None",
"doc_md": "None",
"fileloc": "'/Users/hsheth/projects/datahub/metadata-ingestion-modules/airflow-plugin/tests/integration/dags/basic_iolets.py'",
"is_paused_upon_creation": "None",
"start_date": "DateTime(2023, 1, 1, 0, 0, 0, tzinfo=Timezone('UTC'))",
Expand Down Expand Up @@ -373,7 +375,10 @@
"state": "success",
"operator": "BashOperator",
"priority_weight": "1",
"log_url": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=run_data_task&dag_id=basic_iolets"
"log_url": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=run_data_task&dag_id=basic_iolets",
"orchestrator": "airflow",
"dag_id": "basic_iolets",
"task_id": "run_data_task"
},
"externalUrl": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=run_data_task&dag_id=basic_iolets",
"name": "basic_iolets_run_data_task_manual_run_test",
Expand Down
14 changes: 12 additions & 2 deletions metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v1_simple_dag.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,8 @@
"customProperties": {
"_access_control": "None",
"catchup": "False",
"description": "'A simple DAG that runs a few fake data tasks.'",
"doc_md": "None",
"fileloc": "'/Users/hsheth/projects/datahub/metadata-ingestion-modules/airflow-plugin/tests/integration/dags/simple_dag.py'",
"is_paused_upon_creation": "None",
"start_date": "DateTime(2023, 1, 1, 0, 0, 0, tzinfo=Timezone('UTC'))",
Expand Down Expand Up @@ -302,7 +304,10 @@
"state": "success",
"operator": "BashOperator",
"priority_weight": "2",
"log_url": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=task_1&dag_id=simple_dag"
"log_url": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=task_1&dag_id=simple_dag",
"orchestrator": "airflow",
"dag_id": "simple_dag",
"task_id": "task_1"
},
"externalUrl": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=task_1&dag_id=simple_dag",
"name": "simple_dag_task_1_manual_run_test",
Expand Down Expand Up @@ -433,6 +438,8 @@
"customProperties": {
"_access_control": "None",
"catchup": "False",
"description": "'A simple DAG that runs a few fake data tasks.'",
"doc_md": "None",
"fileloc": "'/Users/hsheth/projects/datahub/metadata-ingestion-modules/airflow-plugin/tests/integration/dags/simple_dag.py'",
"is_paused_upon_creation": "None",
"start_date": "DateTime(2023, 1, 1, 0, 0, 0, tzinfo=Timezone('UTC'))",
Expand Down Expand Up @@ -654,7 +661,10 @@
"state": "success",
"operator": "BashOperator",
"priority_weight": "1",
"log_url": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=run_another_data_task&dag_id=simple_dag"
"log_url": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=run_another_data_task&dag_id=simple_dag",
"orchestrator": "airflow",
"dag_id": "simple_dag",
"task_id": "run_another_data_task"
},
"externalUrl": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=run_another_data_task&dag_id=simple_dag",
"name": "simple_dag_run_another_data_task_manual_run_test",
Expand Down
7 changes: 6 additions & 1 deletion metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v2_basic_iolets.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,8 @@
"customProperties": {
"_access_control": "None",
"catchup": "False",
"description": "None",
"doc_md": "None",
"fileloc": "'/Users/hsheth/projects/datahub/metadata-ingestion-modules/airflow-plugin/tests/integration/dags/basic_iolets.py'",
"is_paused_upon_creation": "None",
"start_date": "DateTime(2023, 1, 1, 0, 0, 0, tzinfo=Timezone('UTC'))",
Expand Down Expand Up @@ -224,7 +226,10 @@
"state": "running",
"operator": "BashOperator",
"priority_weight": "1",
"log_url": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=run_data_task&dag_id=basic_iolets&map_index=-1"
"log_url": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=run_data_task&dag_id=basic_iolets&map_index=-1",
"orchestrator": "airflow",
"dag_id": "basic_iolets",
"task_id": "run_data_task"
},
"externalUrl": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=run_data_task&dag_id=basic_iolets&map_index=-1",
"name": "basic_iolets_run_data_task_manual_run_test",
Expand Down
7 changes: 6 additions & 1 deletion ...ion-modules/airflow-plugin/tests/integration/goldens/v2_basic_iolets_no_dag_listener.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,8 @@
"customProperties": {
"_access_control": "None",
"catchup": "False",
"description": "None",
"doc_md": "None",
"fileloc": "'/Users/hsheth/projects/datahub/metadata-ingestion-modules/airflow-plugin/tests/integration/dags/basic_iolets.py'",
"is_paused_upon_creation": "None",
"start_date": "DateTime(2023, 1, 1, 0, 0, 0, tzinfo=Timezone('UTC'))",
Expand Down Expand Up @@ -224,7 +226,10 @@
"state": "running",
"operator": "BashOperator",
"priority_weight": "1",
"log_url": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=run_data_task&dag_id=basic_iolets&map_index=-1"
"log_url": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=run_data_task&dag_id=basic_iolets&map_index=-1",
"orchestrator": "airflow",
"dag_id": "basic_iolets",
"task_id": "run_data_task"
},
"externalUrl": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=run_data_task&dag_id=basic_iolets&map_index=-1",
"name": "basic_iolets_run_data_task_manual_run_test",
Expand Down
12 changes: 10 additions & 2 deletions metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v2_simple_dag.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,8 @@
"customProperties": {
"_access_control": "None",
"catchup": "False",
"description": "'A simple DAG that runs a few fake data tasks.'",
"doc_md": "None",
"fileloc": "'/Users/hsheth/projects/datahub/metadata-ingestion-modules/airflow-plugin/tests/integration/dags/simple_dag.py'",
"is_paused_upon_creation": "None",
"start_date": "DateTime(2023, 1, 1, 0, 0, 0, tzinfo=Timezone('UTC'))",
Expand Down Expand Up @@ -189,7 +191,10 @@
"state": "running",
"operator": "BashOperator",
"priority_weight": "2",
"log_url": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=task_1&dag_id=simple_dag&map_index=-1"
"log_url": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=task_1&dag_id=simple_dag&map_index=-1",
"orchestrator": "airflow",
"dag_id": "simple_dag",
"task_id": "task_1"
},
"externalUrl": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=task_1&dag_id=simple_dag&map_index=-1",
"name": "simple_dag_task_1_manual_run_test",
Expand Down Expand Up @@ -523,7 +528,10 @@
"state": "running",
"operator": "BashOperator",
"priority_weight": "1",
"log_url": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=run_another_data_task&dag_id=simple_dag&map_index=-1"
"log_url": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=run_another_data_task&dag_id=simple_dag&map_index=-1",
"orchestrator": "airflow",
"dag_id": "simple_dag",
"task_id": "run_another_data_task"
},
"externalUrl": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=run_another_data_task&dag_id=simple_dag&map_index=-1",
"name": "simple_dag_run_another_data_task_manual_run_test",
Expand Down
14 changes: 12 additions & 2 deletions ...stion-modules/airflow-plugin/tests/integration/goldens/v2_simple_dag_no_dag_listener.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,8 @@
"customProperties": {
"_access_control": "None",
"catchup": "False",
"description": "'A simple DAG that runs a few fake data tasks.'",
"doc_md": "None",
"fileloc": "'/Users/hsheth/projects/datahub/metadata-ingestion-modules/airflow-plugin/tests/integration/dags/simple_dag.py'",
"is_paused_upon_creation": "None",
"start_date": "DateTime(2023, 1, 1, 0, 0, 0, tzinfo=Timezone('UTC'))",
Expand Down Expand Up @@ -189,7 +191,10 @@
"state": "running",
"operator": "BashOperator",
"priority_weight": "2",
"log_url": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=task_1&dag_id=simple_dag&map_index=-1"
"log_url": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=task_1&dag_id=simple_dag&map_index=-1",
"orchestrator": "airflow",
"dag_id": "simple_dag",
"task_id": "task_1"
},
"externalUrl": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=task_1&dag_id=simple_dag&map_index=-1",
"name": "simple_dag_task_1_manual_run_test",
Expand Down Expand Up @@ -435,6 +440,8 @@
"customProperties": {
"_access_control": "None",
"catchup": "False",
"description": "'A simple DAG that runs a few fake data tasks.'",
"doc_md": "None",
"fileloc": "'/Users/hsheth/projects/datahub/metadata-ingestion-modules/airflow-plugin/tests/integration/dags/simple_dag.py'",
"is_paused_upon_creation": "None",
"start_date": "DateTime(2023, 1, 1, 0, 0, 0, tzinfo=Timezone('UTC'))",
Expand Down Expand Up @@ -579,7 +586,10 @@
"state": "running",
"operator": "BashOperator",
"priority_weight": "1",
"log_url": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=run_another_data_task&dag_id=simple_dag&map_index=-1"
"log_url": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=run_another_data_task&dag_id=simple_dag&map_index=-1",
"orchestrator": "airflow",
"dag_id": "simple_dag",
"task_id": "run_another_data_task"
},
"externalUrl": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=run_another_data_task&dag_id=simple_dag&map_index=-1",
"name": "simple_dag_run_another_data_task_manual_run_test",
Expand Down
7 changes: 6 additions & 1 deletion ...ata-ingestion-modules/airflow-plugin/tests/integration/goldens/v2_snowflake_operator.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,8 @@
"customProperties": {
"_access_control": "None",
"catchup": "False",
"description": "None",
"doc_md": "None",
"fileloc": "'/Users/hsheth/projects/datahub/metadata-ingestion-modules/airflow-plugin/tests/integration/dags/snowflake_operator.py'",
"is_paused_upon_creation": "None",
"start_date": "DateTime(2023, 1, 1, 0, 0, 0, tzinfo=Timezone('UTC'))",
Expand Down Expand Up @@ -234,7 +236,10 @@
"state": "running",
"operator": "SnowflakeOperator",
"priority_weight": "1",
"log_url": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=transform_cost_table&dag_id=snowflake_operator&map_index=-1"
"log_url": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=transform_cost_table&dag_id=snowflake_operator&map_index=-1",
"orchestrator": "airflow",
"dag_id": "snowflake_operator",
"task_id": "transform_cost_table"
},
"externalUrl": "http://airflow.example.com/log?execution_date=2023-09-27T21%3A34%3A38%2B00%3A00&task_id=transform_cost_table&dag_id=snowflake_operator&map_index=-1",
"name": "snowflake_operator_transform_cost_table_manual_run_test",
Expand Down
Loading
Loading