def get_dag_runs(dag_id: str, state: Optional[str] = None) -> List[Dict[str, Any]]: """ Returns a list of Dag Runs for a specific DAG ID. :param dag_id: String identifier of a DAG :param state: queued|running|success... :return: List of DAG runs of a DAG with requested state, or all runs if the state is not specified """ check_and_get_dag(dag_id=dag_id) dag_runs = [] state = DagRunState(state.lower()) if state else None for run in DagRun.find(dag_id=dag_id, state=state): dag_runs.append({ 'id': run.id, 'run_id': run.run_id, 'state': run.state, 'dag_id': run.dag_id, 'execution_date': run.execution_date.isoformat(), 'start_date': ((run.start_date or '') and run.start_date.isoformat()), 'dag_run_url': url_for('Airflow.graph', dag_id=run.dag_id, execution_date=run.execution_date), }) return dag_runs
def update_dag_run_state(dag_id: str, dag_run_id: str, session) -> dict: """Set a state of a dag run.""" dag_run: Optional[DagRun] = (session.query(DagRun).filter( DagRun.dag_id == dag_id, DagRun.run_id == dag_run_id).one_or_none()) if dag_run is None: error_message = f'Dag Run id {dag_run_id} not found in dag {dag_id}' raise NotFound(error_message) try: post_body = set_dagrun_state_form_schema.load(request.json) except ValidationError as err: raise BadRequest(detail=str(err)) state = post_body['state'] dag_run.set_state(state=DagRunState(state)) session.merge(dag_run) return dagrun_schema.dump(dag_run)
def set_state( *, tasks: Iterable[BaseOperator], dag_run_id: Optional[str] = None, execution_date: Optional[datetime] = None, upstream: bool = False, downstream: bool = False, future: bool = False, past: bool = False, state: TaskInstanceState = TaskInstanceState.SUCCESS, commit: bool = False, session: SASession = NEW_SESSION, ) -> List[TaskInstance]: """ Set the state of a task instance and if needed its relatives. Can set state for future tasks (calculated from run_id) and retroactively for past tasks. Will verify integrity of past dag runs in order to create tasks that did not exist. It will not create dag runs that are missing on the schedule (but it will as for subdag dag runs if needed). :param tasks: the iterable of tasks from which to work. task.task.dag needs to be set :param dag_run_id: the run_id of the dagrun to start looking from :param execution_date: the execution date from which to start looking(deprecated) :param upstream: Mark all parents (upstream tasks) :param downstream: Mark all siblings (downstream tasks) of task_id, including SubDags :param future: Mark all future tasks on the interval of the dag up until last execution date. :param past: Retroactively mark all tasks starting from start_date of the DAG :param state: State to which the tasks need to be set :param commit: Commit tasks to be altered to the database :param session: database session :return: list of tasks that have been created and updated """ if not tasks: return [] if not exactly_one(execution_date, dag_run_id): raise ValueError( "Exactly one of dag_run_id and execution_date must be set") if execution_date and not timezone.is_localized(execution_date): raise ValueError(f"Received non-localized date {execution_date}") task_dags = {task.dag for task in tasks} if len(task_dags) > 1: raise ValueError(f"Received tasks from multiple DAGs: {task_dags}") dag = next(iter(task_dags)) if dag is None: raise ValueError("Received tasks with no DAG") if execution_date: dag_run_id = dag.get_dagrun(execution_date=execution_date).run_id if not dag_run_id: raise ValueError("Received tasks with no dag_run_id") dag_run_ids = get_run_ids(dag, dag_run_id, future, past) task_ids = list(find_task_relatives(tasks, downstream, upstream)) confirmed_infos = list(_iter_existing_dag_run_infos(dag, dag_run_ids)) confirmed_dates = [info.logical_date for info in confirmed_infos] sub_dag_run_ids = list( _iter_subdag_run_ids(dag, session, DagRunState(state), task_ids, commit, confirmed_infos), ) # now look for the task instances that are affected qry_dag = get_all_dag_task_query(dag, session, state, task_ids, confirmed_dates) if commit: tis_altered = qry_dag.with_for_update().all() if sub_dag_run_ids: qry_sub_dag = all_subdag_tasks_query(sub_dag_run_ids, session, state, confirmed_dates) tis_altered += qry_sub_dag.with_for_update().all() for task_instance in tis_altered: task_instance.set_state(state) else: tis_altered = qry_dag.all() if sub_dag_run_ids: qry_sub_dag = all_subdag_tasks_query(sub_dag_run_ids, session, state, confirmed_dates) tis_altered += qry_sub_dag.all() return tis_altered