def _get_dag_run(dag, exec_date_or_run_id, create_if_necessary, session):
dag_run = dag.get_dagrun(run_id=exec_date_or_run_id, session=session)
if dag_run:
return dag_run
execution_date = None
with suppress(ParserError, TypeError):
execution_date = timezone.parse(exec_date_or_run_id)
if create_if_necessary and not execution_date:
return DagRun(dag_id=dag.dag_id, run_id=exec_date_or_run_id)
try:
return (
session.query(DagRun)
.filter(
DagRun.dag_id == dag.dag_id,
DagRun.execution_date == execution_date,
)
.one()
)
except NoResultFound:
if create_if_necessary:
return DagRun(dag.dag_id, execution_date=execution_date)
raise DagRunNotFound(
f"DagRun for {dag.dag_id} with run_id or execution_date of {exec_date_or_run_id!r} not found"
) from None
def test_dagrun_id_is_not_backfill(self):
"""
Task instances whose dagrun ID is not a backfill dagrun ID should pass this dep.
"""
dagrun = DagRun()
dagrun.run_type = 'custom_type'
ti = Mock(get_dagrun=Mock(return_value=dagrun))
self.assertTrue(DagrunIdDep().is_met(ti=ti))
dagrun = DagRun()
dagrun.run_id = None
ti = Mock(get_dagrun=Mock(return_value=dagrun))
self.assertTrue(DagrunIdDep().is_met(ti=ti))
def _get_dag_run(
*,
dag: DAG,
exec_date_or_run_id: str,
create_if_necessary: bool,
session: Session,
) -> DagRun:
"""Try to retrieve a DAG run from a string representing either a run ID or logical date.
This checks DAG runs like this:
1. If the input ``exec_date_or_run_id`` matches a DAG run ID, return the run.
2. Try to parse the input as a date. If that works, and the resulting
date matches a DAG run's logical date, return the run.
3. If ``create_if_necessary`` is *False* and the input works for neither of
the above, raise ``DagRunNotFound``.
4. Try to create a new DAG run. If the input looks like a date, use it as
the logical date; otherwise use it as a run ID and set the logical date
to the current time.
"""
dag_run = dag.get_dagrun(run_id=exec_date_or_run_id, session=session)
if dag_run:
return dag_run
try:
execution_date: Optional[datetime.datetime] = timezone.parse(
exec_date_or_run_id)
except (ParserError, TypeError):
execution_date = None
try:
return (session.query(DagRun).filter(
DagRun.dag_id == dag.dag_id,
DagRun.execution_date == execution_date).one())
except NoResultFound:
if not create_if_necessary:
raise DagRunNotFound(
f"DagRun for {dag.dag_id} with run_id or execution_date of {exec_date_or_run_id!r} not found"
) from None
if execution_date is not None:
return DagRun(dag.dag_id,
run_id=IN_MEMORY_RUN_ID,
execution_date=execution_date)
return DagRun(dag.dag_id,
run_id=exec_date_or_run_id,
execution_date=timezone.utcnow())
def _get_dag_run(
*,
dag: DAG,
exec_date_or_run_id: str,
create_if_necessary: CreateIfNecessary,
session: Session,
) -> Tuple[DagRun, bool]:
"""Try to retrieve a DAG run from a string representing either a run ID or logical date.
This checks DAG runs like this:
1. If the input ``exec_date_or_run_id`` matches a DAG run ID, return the run.
2. Try to parse the input as a date. If that works, and the resulting
date matches a DAG run's logical date, return the run.
3. If ``create_if_necessary`` is *False* and the input works for neither of
the above, raise ``DagRunNotFound``.
4. Try to create a new DAG run. If the input looks like a date, use it as
the logical date; otherwise use it as a run ID and set the logical date
to the current time.
"""
dag_run = dag.get_dagrun(run_id=exec_date_or_run_id, session=session)
if dag_run:
return dag_run, False
try:
execution_date: Optional[datetime.datetime] = timezone.parse(exec_date_or_run_id)
except (ParserError, TypeError):
execution_date = None
try:
dag_run = (
session.query(DagRun)
.filter(DagRun.dag_id == dag.dag_id, DagRun.execution_date == execution_date)
.one()
)
except NoResultFound:
if not create_if_necessary:
raise DagRunNotFound(
f"DagRun for {dag.dag_id} with run_id or execution_date of {exec_date_or_run_id!r} not found"
) from None
else:
return dag_run, False
if execution_date is not None:
dag_run_execution_date = execution_date
else:
dag_run_execution_date = timezone.utcnow()
if create_if_necessary == "memory":
dag_run = DagRun(dag.dag_id, run_id=exec_date_or_run_id, execution_date=dag_run_execution_date)
return dag_run, True
elif create_if_necessary == "db":
dag_run = dag.create_dagrun(
state=DagRunState.QUEUED,
execution_date=dag_run_execution_date,
run_id=_generate_temporary_run_id(),
session=session,
)
return dag_run, True
raise ValueError(f"unknown create_if_necessary value: {create_if_necessary!r}")
def test_dagrun_id_is_backfill(self):
"""
Task instances whose dagrun ID is a backfill dagrun ID should fail this dep.
"""
dagrun = DagRun()
dagrun.run_id = f"{DagRunType.BACKFILL_JOB.value}__something"
ti = Mock(get_dagrun=Mock(return_value=dagrun))
self.assertFalse(DagrunIdDep().is_met(ti=ti))
def test_dagrun_id_is_backfill(self):
"""
Task instances whose dagrun ID is a backfill dagrun ID should fail this dep.
"""
dagrun = DagRun()
dagrun.run_id = "anything"
dagrun.run_type = DagRunType.BACKFILL_JOB
ti = Mock(get_dagrun=Mock(return_value=dagrun))
assert not DagrunIdDep().is_met(ti=ti)
def setUp(self, session) -> None:
self.default_time = datetime(2020, 1, 1)
clear_db_runs()
clear_db_xcom()
self.dag = self._create_dag()
self.app.dag_bag.dags = {self.dag.dag_id: self.dag} # type: ignore # pylint: disable=no-member
self.app.dag_bag.sync_to_db() # type: ignore # pylint: disable=no-member
dr = DagRun(
dag_id=self.dag.dag_id,
run_id="TEST_DAG_RUN_ID",
execution_date=self.default_time,
run_type=DagRunType.MANUAL,
)
session.add(dr)
session.commit()
self.client = self.app.test_client() # type:ignore
def setup_attrs(self, configured_app, session) -> None:
self.default_time = datetime(2020, 1, 1)
clear_db_runs()
clear_db_xcom()
self.app = configured_app
self.dag = self._create_dag()
self.app.dag_bag = DagBag(os.devnull, include_examples=False)
self.app.dag_bag.dags = {self.dag.dag_id: self.dag} # type: ignore # pylint: disable=no-member
self.app.dag_bag.sync_to_db() # type: ignore # pylint: disable=no-member
dr = DagRun(
dag_id=self.dag.dag_id,
run_id="TEST_DAG_RUN_ID",
execution_date=self.default_time,
run_type=DagRunType.MANUAL,
)
session.add(dr)
session.commit()
self.client = self.app.test_client() # type:ignore