def test_serialize(self, session): dagrun_model = DagRun( run_id="my-dag-run", run_type=DagRunType.MANUAL.value, execution_date=timezone.parse(self.default_time), start_date=timezone.parse(self.default_time), conf='{"start": "stop"}', ) session.add(dagrun_model) session.commit() dagrun_model = session.query(DagRun).first() deserialized_dagrun = dagrun_schema.dump(dagrun_model) self.assertEqual( deserialized_dagrun, { "dag_id": None, "dag_run_id": "my-dag-run", "end_date": None, "state": "running", "execution_date": self.default_time, "external_trigger": True, "start_date": self.default_time, "conf": {"start": "stop"}, }, )
def update_dag_run_state(*, dag_id: str, dag_run_id: str, session: Session = NEW_SESSION) -> APIResponse: """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(get_json_request_dict()) except ValidationError as err: raise BadRequest(detail=str(err)) state = post_body['state'] dag = get_airflow_app().dag_bag.get_dag(dag_id) if state == DagRunState.SUCCESS: set_dag_run_state_to_success(dag=dag, run_id=dag_run.run_id, commit=True) elif state == DagRunState.QUEUED: set_dag_run_state_to_queued(dag=dag, run_id=dag_run.run_id, commit=True) else: set_dag_run_state_to_failed(dag=dag, run_id=dag_run.run_id, commit=True) dag_run = session.query(DagRun).get(dag_run.id) return dagrun_schema.dump(dag_run)
def test_serialze(self, session): dagrun_model = DagRun(run_id='my-dag-run', run_type=DagRunType.MANUAL.value, execution_date=timezone.parse(self.default_time), start_date=timezone.parse(self.default_time), conf='{"start": "stop"}' ) session.add(dagrun_model) session.commit() dagrun_model = session.query(DagRun).first() deserialized_dagrun = dagrun_schema.dump(dagrun_model) self.assertEqual( deserialized_dagrun, { 'dag_id': None, 'dag_run_id': 'my-dag-run', 'end_date': None, 'state': 'running', 'execution_date': self.default_time, 'external_trigger': True, 'start_date': self.default_time, 'conf': {"start": "stop"} } )
def post_dag_run(dag_id, session): """Trigger a DAG.""" if not session.query(DagModel).filter(DagModel.dag_id == dag_id).first(): raise NotFound(title="DAG not found", detail=f"DAG with dag_id: '{dag_id}' not found") try: post_body = dagrun_schema.load(request.json, session=session) except ValidationError as err: raise BadRequest(detail=str(err)) dagrun_instance = (session.query(DagRun).filter( DagRun.dag_id == dag_id, or_(DagRun.run_id == post_body["run_id"], DagRun.execution_date == post_body["execution_date"]), ).first()) if not dagrun_instance: dag_run = DagRun(dag_id=dag_id, run_type=DagRunType.MANUAL, **post_body) session.add(dag_run) session.commit() return dagrun_schema.dump(dag_run) if dagrun_instance.execution_date == post_body["execution_date"]: raise AlreadyExists( detail=f"DAGRun with DAG ID: '{dag_id}' and " f"DAGRun ExecutionDate: '{post_body['execution_date']}' already exists" ) raise AlreadyExists( detail= f"DAGRun with DAG ID: '{dag_id}' and DAGRun ID: '{post_body['run_id']}' already exists" )
def get_dag_run(dag_id, dag_run_id, session): """ Get a DAG Run. """ dag_run = session.query(DagRun).filter(DagRun.dag_id == dag_id, DagRun.run_id == dag_run_id).one_or_none() if dag_run is None: raise NotFound("DAGRun not found") return dagrun_schema.dump(dag_run)
def get_dag_run(*, dag_id: str, dag_run_id: str, session: Session = NEW_SESSION) -> APIResponse: """Get a DAG Run.""" dag_run = session.query(DagRun).filter(DagRun.dag_id == dag_id, DagRun.run_id == dag_run_id).one_or_none() if dag_run is None: raise NotFound( "DAGRun not found", detail=f"DAGRun with DAG ID: '{dag_id}' and DagRun ID: '{dag_run_id}' not found", ) return dagrun_schema.dump(dag_run)
def post_dag_run(*, dag_id: str, session: Session = NEW_SESSION) -> APIResponse: """Trigger a DAG.""" dm = session.query(DagModel).filter(DagModel.dag_id == dag_id).first() if not dm: raise NotFound(title="DAG not found", detail=f"DAG with dag_id: '{dag_id}' not found") if dm.has_import_errors: raise BadRequest( title="DAG cannot be triggered", detail=f"DAG with dag_id: '{dag_id}' has import errors", ) try: post_body = dagrun_schema.load(get_json_request_dict(), session=session) except ValidationError as err: raise BadRequest(detail=str(err)) logical_date = pendulum.instance(post_body["execution_date"]) run_id = post_body["run_id"] dagrun_instance = (session.query(DagRun).filter( DagRun.dag_id == dag_id, or_(DagRun.run_id == run_id, DagRun.execution_date == logical_date), ).first()) if not dagrun_instance: try: dag = get_airflow_app().dag_bag.get_dag(dag_id) dag_run = dag.create_dagrun( run_type=DagRunType.MANUAL, run_id=run_id, execution_date=logical_date, data_interval=dag.timetable.infer_manual_data_interval( run_after=logical_date), state=DagRunState.QUEUED, conf=post_body.get("conf"), external_trigger=True, dag_hash=get_airflow_app().dag_bag.dags_hash.get(dag_id), ) return dagrun_schema.dump(dag_run) except ValueError as ve: raise BadRequest(detail=str(ve)) if dagrun_instance.execution_date == logical_date: raise AlreadyExists(detail=( f"DAGRun with DAG ID: '{dag_id}' and " f"DAGRun logical date: '{logical_date.isoformat(sep=' ')}' already exists" ), ) raise AlreadyExists( detail= f"DAGRun with DAG ID: '{dag_id}' and DAGRun ID: '{run_id}' already exists" )
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 post_dag_run(dag_id, session): """Trigger a DAG.""" if not session.query(DagModel).filter(DagModel.dag_id == dag_id).first(): raise NotFound(title="DAG not found", detail=f"DAG with dag_id: '{dag_id}' not found") post_body = dagrun_schema.load(request.json, session=session) dagrun_instance = ( session.query(DagRun).filter(DagRun.dag_id == dag_id, DagRun.run_id == post_body["run_id"]).first() ) if not dagrun_instance: dag_run = DagRun(dag_id=dag_id, run_type=DagRunType.MANUAL, **post_body) session.add(dag_run) session.commit() return dagrun_schema.dump(dag_run) raise AlreadyExists( detail=f"DAGRun with DAG ID: '{dag_id}' and DAGRun ID: '{post_body['run_id']}' already exists" )
def post_dag_run(dag_id, session): """Trigger a DAG.""" if not session.query(DagModel).filter(DagModel.dag_id == dag_id).first(): raise NotFound(title="DAG not found", detail=f"DAG with dag_id: '{dag_id}' not found") try: post_body = dagrun_schema.load(request.json, session=session) except ValidationError as err: raise BadRequest(detail=str(err)) logical_date = pendulum.instance(post_body["execution_date"]) run_id = post_body["run_id"] dagrun_instance = ( session.query(DagRun) .filter( DagRun.dag_id == dag_id, or_(DagRun.run_id == run_id, DagRun.execution_date == logical_date), ) .first() ) if not dagrun_instance: try: dag = current_app.dag_bag.get_dag(dag_id) dag_run = dag.create_dagrun( run_type=DagRunType.MANUAL, run_id=run_id, execution_date=logical_date, data_interval=dag.timetable.infer_manual_data_interval(run_after=logical_date), state=State.QUEUED, conf=post_body.get("conf"), external_trigger=True, dag_hash=current_app.dag_bag.dags_hash.get(dag_id), ) return dagrun_schema.dump(dag_run) except ValueError as ve: raise BadRequest(detail=str(ve)) if dagrun_instance.execution_date == logical_date: raise AlreadyExists( detail=( f"DAGRun with DAG ID: '{dag_id}' and " f"DAGRun logical date: '{logical_date.isoformat(sep=' ')}' already exists" ), ) raise AlreadyExists(detail=f"DAGRun with DAG ID: '{dag_id}' and DAGRun ID: '{run_id}' already exists")
def clear_dag_run(*, dag_id: str, dag_run_id: str, session: Session = NEW_SESSION) -> APIResponse: """Clear 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 = clear_dagrun_form_schema.load(get_json_request_dict()) except ValidationError as err: raise BadRequest(detail=str(err)) dry_run = post_body.get('dry_run', False) dag = get_airflow_app().dag_bag.get_dag(dag_id) start_date = dag_run.logical_date end_date = dag_run.logical_date if dry_run: task_instances = dag.clear( start_date=start_date, end_date=end_date, task_ids=None, include_subdags=True, include_parentdag=True, only_failed=False, dry_run=True, ) return task_instance_reference_collection_schema.dump( TaskInstanceReferenceCollection(task_instances=task_instances)) else: dag.clear( start_date=start_date, end_date=end_date, task_ids=None, include_subdags=True, include_parentdag=True, only_failed=False, ) dag_run.refresh_from_db() return dagrun_schema.dump(dag_run)
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 = current_app.dag_bag.get_dag(dag_id) if state == State.SUCCESS: set_dag_run_state_to_success(dag, dag_run.execution_date, commit=True) else: set_dag_run_state_to_failed(dag, dag_run.execution_date, commit=True) dag_run = session.query(DagRun).get(dag_run.id) return dagrun_schema.dump(dag_run)
def post_dag_run(dag_id, session): """Trigger a DAG.""" if not session.query(DagModel).filter(DagModel.dag_id == dag_id).first(): raise NotFound(title="DAG not found", detail=f"DAG with dag_id: '{dag_id}' not found") try: post_body = dagrun_schema.load(request.json, session=session) except ValidationError as err: raise BadRequest(detail=str(err)) execution_date = post_body["execution_date"] run_id = post_body["run_id"] dagrun_instance = (session.query(DagRun).filter( DagRun.dag_id == dag_id, or_(DagRun.run_id == run_id, DagRun.execution_date == execution_date), ).first()) if not dagrun_instance: dag_run = current_app.dag_bag.get_dag(dag_id).create_dagrun( run_type=DagRunType.MANUAL, run_id=run_id, execution_date=execution_date, state=State.QUEUED, conf=post_body.get("conf"), external_trigger=True, dag_hash=current_app.dag_bag.dags_hash.get(dag_id), ) return dagrun_schema.dump(dag_run) if dagrun_instance.execution_date == execution_date: raise AlreadyExists( detail=f"DAGRun with DAG ID: '{dag_id}' and " f"DAGRun ExecutionDate: '{post_body['execution_date']}' already exists" ) raise AlreadyExists( detail= f"DAGRun with DAG ID: '{dag_id}' and DAGRun ID: '{post_body['run_id']}' already exists" )