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

fix(ingest): support airflow mapped operators #6738

Merged
merged 3 commits into from
Dec 14, 2022
Merged

fix(ingest): support airflow mapped operators #6738

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
755c6f3
fix(ingest) - Basic support for dynamic task mapping in Airflow plugin
cccs-seb Dec 12, 2022
ab8cdd2
fix(ingest): Missing return statement
cccs-seb Dec 12, 2022
2da365b
fix(ingest): Move strings to constants.
cccs-seb Dec 13, 2022
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
37 changes: 31 additions & 6 deletions metadata-ingestion/src/datahub_provider/_plugin.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,13 +1,12 @@
from datahub_provider._airflow_compat import Operator
from datahub_provider._airflow_compat import BaseOperator, MappedOperator, Operator

import contextlib
import logging
import traceback
from typing import Any, Callable, Iterable, List, Optional
from typing import Any, Callable, Iterable, List, Optional, Union

from airflow.configuration import conf
from airflow.lineage import PIPELINE_OUTLETS
from airflow.models.baseoperator import BaseOperator
from airflow.plugins_manager import AirflowPlugin
from airflow.utils.module_loading import import_string
from cattr import structure
Expand All @@ -19,6 +18,9 @@

logger = logging.getLogger(__name__)

TASK_ON_FAILURE_CALLBACK = "on_failure_callback"
TASK_ON_SUCCESS_CALLBACK = "on_success_callback"


def get_lineage_config() -> DatahubLineageConfig:
"""Load the lineage config from airflow.cfg."""
Expand Down Expand Up @@ -290,12 +292,35 @@ def custom_on_success_callback(context):
return custom_on_success_callback


def task_policy(task: BaseOperator) -> None:
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"])
# task.pre_execute = _wrap_pre_execution(task.pre_execute)
task.on_failure_callback = _wrap_on_failure_callback(task.on_failure_callback)
task.on_success_callback = _wrap_on_success_callback(task.on_success_callback)

# MappedOperator's callbacks don't have setters until Airflow 2.X.X
# https://github.com/apache/airflow/issues/24547
# We can bypass this by going through partial_kwargs for now
if MappedOperator and isinstance(task, MappedOperator): # type: ignore
on_failure_callback_prop: property = getattr(
MappedOperator, TASK_ON_FAILURE_CALLBACK
)
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:
task.log.debug(
"Using MappedOperator's partial_kwargs instead of callback properties"
)
task.partial_kwargs[TASK_ON_FAILURE_CALLBACK] = _wrap_on_failure_callback(
task.on_failure_callback
)
task.partial_kwargs[TASK_ON_SUCCESS_CALLBACK] = _wrap_on_success_callback(
task.on_success_callback
)
cccs-seb marked this conversation as resolved.
Show resolved Hide resolved
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.pre_execute = _wrap_pre_execution(task.pre_execute)


Expand Down