Skip to content

Fix warning in airflow tasks test command regarding absence of data_interval #27106

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
potiuk merged 3 commits into from
Mar 16, 2023
Merged

Fix warning in airflow tasks test command regarding absence of data_interval #27106

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
4e29197
Fix warning in airflow tasks test command regarding absence of data_i…
tirkarthi Oct 18, 2022
87b6a24
Use manual data interval in task commands
uranusjr Mar 16, 2023
6fef2ea
Update tests/cli/commands/test_task_command.py
potiuk Mar 16, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 4 additions & 3 deletions airflow/cli/commands/task_command.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,6 @@
"""Task sub-commands."""
from __future__ import annotations

import datetime
import importlib
import json
import logging
Expand All @@ -28,6 +27,7 @@
from contextlib import contextmanager, redirect_stderr, redirect_stdout, suppress
from typing import Generator, Union

import pendulum
from pendulum.parsing.exceptions import ParserError
from sqlalchemy.orm.exc import NoResultFound
from sqlalchemy.orm.session import Session
Expand Down Expand Up @@ -102,7 +102,7 @@ def _get_dag_run(
"""
if not exec_date_or_run_id and not create_if_necessary:
raise ValueError("Must provide `exec_date_or_run_id` if not `create_if_necessary`.")
execution_date: datetime.datetime | None = None
execution_date: pendulum.DateTime | None = None
if exec_date_or_run_id:
dag_run = dag.get_dagrun(run_id=exec_date_or_run_id, session=session)
if dag_run:
Expand All @@ -127,7 +127,7 @@ def _get_dag_run(
if execution_date is not None:
dag_run_execution_date = execution_date
else:
dag_run_execution_date = timezone.utcnow()
dag_run_execution_date = pendulum.instance(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)
Expand All @@ -137,6 +137,7 @@ def _get_dag_run(
state=DagRunState.QUEUED,
execution_date=dag_run_execution_date,
run_id=_generate_temporary_run_id(),
data_interval=dag.timetable.infer_manual_data_interval(run_after=dag_run_execution_date),
session=session,
)
return dag_run, True
Expand Down
24 changes: 20 additions & 4 deletions tests/cli/commands/test_task_command.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -88,8 +88,13 @@ def setup_class(cls):

cls.dag = cls.dagbag.get_dag(cls.dag_id)
cls.dagbag.sync_to_db()
data_interval = cls.dag.timetable.infer_manual_data_interval(run_after=DEFAULT_DATE)
cls.dag_run = cls.dag.create_dagrun(
state=State.NONE, run_id=cls.run_id, run_type=DagRunType.MANUAL, execution_date=DEFAULT_DATE
state=State.NONE,
run_id=cls.run_id,
run_type=DagRunType.MANUAL,
execution_date=DEFAULT_DATE,
data_interval=data_interval,
)

@classmethod
Expand Down Expand Up @@ -178,11 +183,14 @@ def test_cli_test_different_path(self, session):
dag = dagbag.get_dag("test_dags_folder")
dagbag.sync_to_db(session=session)

execution_date = pendulum.now("UTC")
data_interval = dag.timetable.infer_manual_data_interval(run_after=execution_date)
dag.create_dagrun(
state=State.NONE,
run_id="abc123",
run_type=DagRunType.MANUAL,
execution_date=pendulum.now("UTC"),
execution_date=execution_date,
data_interval=data_interval,
session=session,
)
session.commit()
Expand Down Expand Up @@ -492,9 +500,11 @@ def test_task_states_for_dag_run(self):
task2 = dag2.get_task(task_id="print_the_context")
default_date2 = timezone.datetime(2016, 1, 9)
dag2.clear()
data_interval = dag2.timetable.infer_manual_data_interval(run_after=default_date2)
dagrun = dag2.create_dagrun(
state=State.RUNNING,
execution_date=default_date2,
data_interval=data_interval,
run_type=DagRunType.MANUAL,
external_trigger=True,
)
Expand Down Expand Up @@ -579,9 +589,12 @@ def setup_method(self) -> None:
self.ti_log_file_path = os.path.join(self.log_dir, self.log_filename)
self.parser = cli_parser.get_parser()

DagBag().get_dag(self.dag_id).create_dagrun(
dag = DagBag().get_dag(self.dag_id)
data_interval = dag.timetable.infer_manual_data_interval(run_after=self.execution_date)
dag.create_dagrun(
run_id=self.run_id,
execution_date=self.execution_date,
data_interval=data_interval,
start_date=timezone.utcnow(),
state=State.RUNNING,
run_type=DagRunType.MANUAL,
Expand Down Expand Up @@ -862,9 +875,12 @@ def test_context_with_run():
task_args = ["tasks", "run", dag_id, task_id, "--local", execution_date_str]
parser = cli_parser.get_parser()

DagBag().get_dag(dag_id).create_dagrun(
dag = DagBag().get_dag(dag_id)
data_interval = dag.timetable.infer_manual_data_interval(run_after=execution_date)
dag.create_dagrun(
run_id=run_id,
execution_date=execution_date,
data_interval=data_interval,
start_date=timezone.utcnow(),
state=State.RUNNING,
run_type=DagRunType.MANUAL,
Expand Down