Repository: apache/airflow
File: task-sdk/src/airflow/sdk/definitions/dag.py (Python)

Code before:


def _run_inline_trigger(trigger, task_sdk_ti):
    from airflow.sdk.execution_time.supervisor import InProcessTestSupervisor

    return InProcessTestSupervisor.run_trigger_in_process(trigger=trigger, ti=task_sdk_ti)


# Since we define all the attributes of the class with attrs, we can compute this statically at parse time
DAG._DAG__serialized_fields = frozenset(a.name for a in attrs.fields(DAG)) - {  # type: ignore[attr-defined]
    "schedule_asset_references",
    "schedule_asset_alias_references",
    "task_outlet_asset_references",
    "_old_context_manager_dags",
    "safe_dag_id",
    "last_loaded",
    "user_defined_filters",
    "user_defined_macros",
    "partial",
    "params",
    "_log",
    "task_dict",
    "template_searchpath",
    "sla_miss_callback",
    "on_success_callback",
    "on_failure_callback",
    "template_undefined",
    "jinja_environment_kwargs",
    # has_on_*_callback are only stored if the value is True, as the default is False
    "has_on_success_callback",
    "has_on_failure_callback",
    "auto_register",
    "schedule",
}

if TYPE_CHECKING:
    # NOTE: Please keep the list of arguments in sync with DAG.__init__.
    # Only exception: dag_id here should have a default value, but not in DAG.
    @overload
    def dag(
        dag_id: str = "",
        *,
        description: str | None = None,
        schedule: ScheduleArg = None,
        start_date: datetime | None = None,
        end_date: datetime | None = None,
        template_searchpath: str | Iterable[str] | None = None,
        template_undefined: type[jinja2.StrictUndefined] = jinja2.StrictUndefined,
        user_defined_macros: dict | None = None,
        user_defined_filters: dict | None = None,
        default_args: dict[str, Any] | None = None,

Review comment (suggestion):
The fact that it is used in `_run_task` might be a good indication to place it in _generated.py, so that the api schema has it and it can be used, but that can come later if / when we decide to split `State` appropriately.

Code after:

            if ti.state == TaskInstanceState.DEFERRED and isinstance(msg, DeferTask) and run_triggerer:
                from airflow.utils.session import create_session

                # API Server expects the task instance to be in QUEUED state before
                # resuming from deferral.
                ti.set_state(TaskInstanceState.QUEUED)

                log.info("[DAG TEST] running trigger in line")
                trigger = import_string(msg.classpath)(**msg.trigger_kwargs)
                event = _run_inline_trigger(trigger, task_sdk_ti)
                ti.next_method = msg.next_method
                ti.next_kwargs = {"event": event.payload} if event else msg.next_kwargs
                log.info("[DAG TEST] Trigger completed")

                # Set the state to SCHEDULED so that the task can be resumed.
                with create_session() as session:
                    ti.state = TaskInstanceState.SCHEDULED
                    session.add(ti)
                continue

            break
        except Exception:
            log.exception("[DAG TEST] Error running task %s", ti)
            if ti.state not in FINISHED_STATES:
                ti.set_state(TaskInstanceState.FAILED)
                taskrun_result = None
                break
            raise

    log.info("[DAG TEST] end task task_id=%s map_index=%s", ti.task_id, ti.map_index)
    return taskrun_result


def _run_inline_trigger(trigger, task_sdk_ti):
    from airflow.sdk.execution_time.supervisor import InProcessTestSupervisor

    return InProcessTestSupervisor.run_trigger_in_process(trigger=trigger, ti=task_sdk_ti)


# Since we define all the attributes of the class with attrs, we can compute this statically at parse time
DAG._DAG__serialized_fields = frozenset(a.name for a in attrs.fields(DAG)) - {  # type: ignore[attr-defined]
    "schedule_asset_references",
    "schedule_asset_alias_references",
    "task_outlet_asset_references",
    "_old_context_manager_dags",
    "safe_dag_id",
    "last_loaded",
    "user_defined_filters",
    "user_defined_macros",
    "partial",