Skip to content

spark_on_k8s.airflow.ui.task_binding

Implementation reference. These symbols describe how the project works; their presence here does not make them supported application imports. Use the documented entry points and support policy.

load_task_binding(reference, dag_bag)

Read only the current, authorized task attempt's small JSON binding.

Source code in spark_on_k8s/airflow/ui/task_binding.py 
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
def load_task_binding(reference: TaskReference, dag_bag) -> TaskBinding:
    """Read only the current, authorized task attempt's small JSON binding."""
    from airflow.api_fastapi.core_api.services.public.task_coordinates import resolve_task_scope
    from airflow.models.task_coordinates import TaskCoordinateResolver
    from airflow.models.taskinstance import TaskInstance
    from airflow.models.xcom import XComModel
    from airflow.utils.session import create_session
    from sqlalchemy import select

    with create_session() as session:
        resolver = TaskCoordinateResolver(dag_bag, session)
        scope = resolve_task_scope(
            dag_id=reference.dag_id,
            run_id=reference.run_id,
            task_id=reference.task_id,
            map_index=reference.map_index,
            region_id=reference.region_id,
            region_index=reference.region_index,
            resolver=resolver,
        )
        task = session.scalar(
            select(TaskInstance).where(
                TaskInstance.dag_id == reference.dag_id,
                TaskInstance.run_id == reference.run_id,
                TaskInstance.task_id == reference.task_id,
                TaskInstance.region_id == scope.region_id,
                TaskInstance.region_index == scope.region_index,
            )
        )
        if task is None or not runs_provider_operator(task, resolver):
            raise failure("task_not_found", "This task has no Spark application binding.")
        if task.try_number != reference.try_number:
            raise failure(
                "stale_attempt", "Select the current task attempt before managing its Spark application.", 409
            )
        producer = select(TaskInstance.id).where(TaskInstance.id == task.id)
        row = session.scalar(
            XComModel.get_many(run_id=reference.run_id, key=CONSOLE_XCOM_KEY, producer_ids=producer)
        )
        if row is None:
            raise failure(
                "binding_not_found", "The Spark application has not been submitted by this task attempt."
            )
        try:
            # Do not invoke custom XCom deserializers or fetch arbitrary remote XCom objects.
            value = json.loads(row.value) if isinstance(row.value, str) else row.value
            binding = TaskBinding.model_validate(value)
        except (ValueError, TypeError, ValidationError) as error:
            raise failure(
                "invalid_binding", "The task's Spark application binding is unavailable."
            ) from error
        if binding.try_number != reference.try_number:
            raise failure("stale_binding", "The Spark binding belongs to an earlier task attempt.", 409)
        # Ownership comes from the authorized database task, never editable XCom data.
        binding.task_context = AirflowTaskContext(
            dag_id=task.dag_id,
            run_id=task.run_id,
            task_id=task.task_id,
            map_index=resolver.public_map_index(task),
            try_number=task.try_number,
            region_id=task.region_id,
            region_index=task.region_index,
        )
        return binding

runs_provider_operator(task, resolver)

Whether the task's operator is a provider operator or inherits from one.

Source code in spark_on_k8s/airflow/ui/task_binding.py 
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
def runs_provider_operator(task, resolver) -> bool:
    """Whether the task's operator is a provider operator or inherits from one."""
    if task.operator in PROVIDER_OPERATORS:
        return True
    from spark_on_k8s.airflow.operator_links import SparkOnK8SOperatorLink

    # Airflow serializes an operator's class name, not its base classes, and does not keep
    # custom class attributes on mapped tasks. Every provider operator declares the Spark job
    # link; subclasses inherit it and Airflow serializes it for plain and mapped tasks alike.
    operator = resolver.find_task(task)
    return any(
        getattr(link, "name", None) == SparkOnK8SOperatorLink.name
        and getattr(link, "xcom_key", None) == SparkOnK8SOperatorLink.xcom_key
        for link in getattr(operator, "operator_extra_links", None) or ()
    )