kaxil commented on code in PR #50300:
URL: https://github.com/apache/airflow/pull/50300#discussion_r2081448983
##########
task-sdk/src/airflow/sdk/definitions/dag.py:
##########
@@ -1014,6 +1016,281 @@ def _validate_owner_links(self, _, owner_links):
f"Bad formatted links are: {wrong_links}"
)
+ def test(
+ self,
+ run_after: datetime | None = None,
+ logical_date: datetime | None = None,
+ run_conf: dict[str, Any] | None = None,
+ conn_file_path: str | None = None,
+ variable_file_path: str | None = None,
+ use_executor: bool = False,
+ mark_success_pattern: Pattern | str | None = None,
+ ):
+ """
+ Execute one single DagRun for a given DAG and logical date.
+
+ :param run_after: the datetime before which to Dag cannot run.
+ :param logical_date: logical date for the DAG run
+ :param run_conf: configuration to pass to newly created dagrun
+ :param conn_file_path: file path to a connection file in either yaml
or json
+ :param variable_file_path: file path to a variable file in either yaml
or json
+ :param use_executor: if set, uses an executor to test the DAG
+ :param mark_success_pattern: regex of task_ids to mark as success
instead of running
+ """
+ import re
+ import time
+ from contextlib import ExitStack
+
+ from airflow import settings
+ from airflow.configuration import secrets_backend_list
+ from airflow.models.dag import DAG as SchedulerDAG,
_get_or_create_dagrun
+ from airflow.models.dagrun import DagRun
+ from airflow.secrets.local_filesystem import LocalFilesystemBackend
+ from airflow.serialization.serialized_objects import SerializedDAG
+ from airflow.utils import timezone
+ from airflow.utils.state import DagRunState, State, TaskInstanceState
+ from airflow.utils.types import DagRunTriggeredByType, DagRunType
+
+ if TYPE_CHECKING:
+ from airflow.models.taskinstance import TaskInstance
+
+ def add_logger_if_needed(ti: TaskInstance):
+ """
+ Add a formatted logger to the task instance.
+
+ This allows all logs to surface to the command line, instead of
into
+ a task file. Since this is a local test run, it is much better for
+ the user to see logs in the command line, rather than needing to
+ search for a log file.
+
+ :param ti: The task instance that will receive a logger.
+ """
+ format = logging.Formatter("[%(asctime)s]
{%(filename)s:%(lineno)d} %(levelname)s - %(message)s")
+ handler = logging.StreamHandler(sys.stdout)
+ handler.level = logging.INFO
+ handler.setFormatter(format)
+ # only add log handler once
+ if not any(isinstance(h, logging.StreamHandler) for h in
ti.log.handlers):
+ log.debug("Adding Streamhandler to taskinstance %s",
ti.task_id)
+ ti.log.addHandler(handler)
+
+ exit_stack = ExitStack()
+
+ if conn_file_path or variable_file_path:
+ local_secrets = LocalFilesystemBackend(
+ variables_file_path=variable_file_path,
connections_file_path=conn_file_path
+ )
+ secrets_backend_list.insert(0, local_secrets)
+ exit_stack.callback(lambda: secrets_backend_list.pop(0))
+
+ session = settings.Session()
+
+ with exit_stack:
+ self.validate()
+ log.debug("Clearing existing task instances for logical date %s",
logical_date)
+ # TODO: Replace with calling client.dag_run.clear in Execution API
at some point
+ SchedulerDAG.clear_dags(
+ dags=[self],
+ start_date=logical_date,
+ end_date=logical_date,
+ dag_run_state=False, # type: ignore
+ )
+
+ log.debug("Getting dagrun for dag %s", self.dag_id)
+ logical_date = timezone.coerce_datetime(logical_date)
+ run_after = timezone.coerce_datetime(run_after) or
timezone.coerce_datetime(timezone.utcnow())
+ data_interval = (
+
self.timetable.infer_manual_data_interval(run_after=logical_date) if
logical_date else None
+ )
+ scheduler_dag =
SerializedDAG.deserialize_dag(SerializedDAG.serialize_dag(self)) # type:
ignore[arg-type]
+
+ dr: DagRun = _get_or_create_dagrun(
+ dag=scheduler_dag,
+ start_date=logical_date or run_after,
+ logical_date=logical_date,
+ data_interval=data_interval,
+ run_after=run_after,
+ run_id=DagRun.generate_run_id(
+ run_type=DagRunType.MANUAL,
+ logical_date=logical_date,
+ run_after=run_after,
+ ),
+ session=session,
+ conf=run_conf,
+ triggered_by=DagRunTriggeredByType.TEST,
+ )
+ # Start a mock span so that one is present and not started
downstream. We
+ # don't care about otel in dag.test and starting the span during
dagrun update
+ # is not functioning properly in this context anyway.
+ dr.start_dr_spans_if_needed(tis=[])
+ dr.dag = self # type: ignore[assignment]
+
+ tasks = self.task_dict
+ log.debug("starting dagrun")
+ # Instead of starting a scheduler, we run the minimal loop
possible to check
+ # for task readiness and dependency management.
+ # Instead of starting a scheduler, we run the minimal loop
possible to check
+ # for task readiness and dependency management.
+
+ # ``Dag.test()`` works in two different modes depending on
``use_executor``:
+ # - if ``use_executor`` is False, runs the task locally with no
executor using ``_run_task``
+ # - if ``use_executor`` is True, sends the task instances to the
executor with
+ # ``BaseExecutor.queue_task_instance``
+ if use_executor:
+ from pathlib import Path
+
+ from airflow.executors.base_executor import ExecutorLoader
+ from airflow.executors.workloads import BundleInfo
+
+ executor = ExecutorLoader.get_default_executor()
+ executor.start()
+
+ while dr.state == DagRunState.RUNNING:
+ session.expire_all()
+ schedulable_tis, _ = dr.update_state(session=session)
+ for s in schedulable_tis:
+ if s.state != TaskInstanceState.UP_FOR_RESCHEDULE:
+ s.try_number += 1
+ s.state = TaskInstanceState.SCHEDULED
+ s.scheduled_dttm = timezone.utcnow()
+ session.commit()
+ # triggerer may mark tasks scheduled so we read from DB
+ all_tis = set(dr.get_task_instances(session=session))
+ scheduled_tis = {x for x in all_tis if x.state ==
TaskInstanceState.SCHEDULED}
+ ids_unrunnable = {x for x in all_tis if x.state not in
State.finished} - scheduled_tis
+ if not scheduled_tis and ids_unrunnable:
+ log.warning("No tasks to run. unrunnable tasks: %s",
ids_unrunnable)
+ time.sleep(1)
+
+ for ti in scheduled_tis:
+ ti.task = tasks[ti.task_id]
+
+ mark_success = (
+ re.compile(mark_success_pattern).fullmatch(ti.task_id)
is not None
+ if mark_success_pattern is not None
+ else False
+ )
+
+ if use_executor:
+ if executor.has_task(ti):
+ continue
+ # TODO: Task-SDK: This check is transitionary. Remove
once all executors are ported over.
+ from airflow.executors import workloads
+ from airflow.executors.base_executor import
BaseExecutor
+
+ if executor.queue_workload.__func__ is not
BaseExecutor.queue_workload: # type: ignore[attr-defined]
+ workload = workloads.ExecuteTask.make(
+ ti,
+ dag_rel_path=Path(self.fileloc),
+ generator=executor.jwt_generator,
+ # For the system test/debug purpose, we use
the default bundle which uses
+ # local file system. If it turns out to be a
feature people want, we could
+ # plumb the Bundle to use as a parameter to
dag.test
+ bundle_info=BundleInfo(name="dags-folder"),
+ )
+ executor.queue_workload(workload, session=session)
+ ti.state = TaskInstanceState.QUEUED
+ session.commit()
+ else:
+ # Send the task to the executor
+ executor.queue_task_instance(ti,
ignore_ti_state=True)
Review Comment:
Good call, not needed
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
