def test_task_states_for_dag_run(self):
dag2 = DagBag().dags['example_python_operator']
task2 = dag2.get_task(task_id='print_the_context')
default_date2 = timezone.make_aware(datetime(2016, 1, 9))
dag2.clear()
ti2 = TaskInstance(task2, default_date2)
ti2.set_state(State.SUCCESS)
ti_start = ti2.start_date
ti_end = ti2.end_date
with redirect_stdout(io.StringIO()) as stdout:
task_command.task_states_for_dag_run(
self.parser.parse_args([
'tasks',
'states-for-dag-run',
'example_python_operator',
default_date2.isoformat(),
'--output',
"json",
]))
actual_out = json.loads(stdout.getvalue())
assert len(actual_out) == 1
assert actual_out[0] == {
'dag_id': 'example_python_operator',
'execution_date': '2016-01-09T00:00:00+00:00',
'task_id': 'print_the_context',
'state': 'success',
'start_date': ti_start.isoformat(),
'end_date': ti_end.isoformat(),
}
def test_log_file_template_with_run_task(self):
"""Verify that the taskinstance has the right context for log_filename_template"""
with mock.patch.object(task_command, "_run_task_by_selected_method"):
with conf_vars({('core', 'dags_folder'): self.dag_path}):
# increment the try_number of the task to be run
dag = DagBag().get_dag(self.dag_id)
task = dag.get_task(self.task_id)
with create_session() as session:
dag.create_dagrun(
execution_date=self.execution_date,
start_date=timezone.utcnow(),
state=State.RUNNING,
run_type=DagRunType.MANUAL,
session=session,
)
ti = TaskInstance(task, self.execution_date)
ti.refresh_from_db(session=session, lock_for_update=True)
ti.try_number = 1 # not running, so starts at 0
session.merge(ti)
log_file_path = os.path.join(
os.path.dirname(self.ti_log_file_path), "2.log")
try:
task_command.task_run(
self.parser.parse_args(self.task_args))
assert os.path.exists(log_file_path)
finally:
try:
os.remove(log_file_path)
except OSError:
pass
def test_task_states_for_dag_run(self):
dag2 = DagBag().dags['example_python_operator']
task2 = dag2.get_task(task_id='print_the_context')
defaut_date2 = timezone.make_aware(datetime(2016, 1, 9))
ti2 = TaskInstance(task2, defaut_date2)
ti2.set_state(State.SUCCESS)
ti_start = ti2.start_date
ti_end = ti2.end_date
with redirect_stdout(io.StringIO()) as stdout:
task_command.task_states_for_dag_run(
self.parser.parse_args([
'tasks', 'states_for_dag_run', 'example_python_operator',
defaut_date2.isoformat()
]))
actual_out = stdout.getvalue()
formatted_rows = [
('example_python_operator', '2016-01-09 00:00:00+00:00',
'print_the_context', 'success', ti_start, ti_end)
]
expected = tabulate(
formatted_rows,
['dag', 'exec_date', 'task', 'state', 'start_date', 'end_date'],
tablefmt="fancy_grid")
# Check that prints, and log messages, are shown
self.assertEqual(expected.replace("\n", ""),
actual_out.replace("\n", ""))
def get_task_log(job_id):
logging.info("Executing custom 'get_task_log' function")
try:
session = settings.Session()
ti = session.query(TaskInstance).filter(
TaskInstance.job_id == job_id).first()
session.close()
if ti is not None:
dag = DagBag("dags").get_dag(ti.dag_id)
ti.task = dag.get_task(ti.task_id)
logger = logging.getLogger('airflow.task')
task_log_reader = configuration.conf.get('core', 'task_log_reader')
handler = next((handler for handler in logger.handlers
if handler.name == task_log_reader), None)
logs = handler.read(ti, None)
if logs:
return ApiResponse.success(logs[0])
else:
return ApiResponse.success("No More Log")
else:
return ApiResponse.success(
"*** Task instance did not exist in the DB\n")
except Exception as e:
error_message = "An error occurred while trying to get_task_log'" + "': " + str(
e)
logging.error(error_message)
return ApiResponse.server_error(error_message)
def test_t1():
dag_id = 'arsalan_outfittery_challenge'
dag = DagBag().get_dag(dag_id)
t1_task = dag.get_task('t1')
execution_date = datetime.now()
t1_ti = TaskInstance(task=t1_task, execution_date=execution_date)
context = t1_ti.get_template_context()
t1_task.execute(context)
def find_task_log():
logging.info("Executing custom 'find_task_log' function")
try:
dag_id = request.args.get('dag_id')
task_id = request.args.get('task_id')
execution_date = request.args.get('execution_date')
dttm = pendulum.parse(execution_date)
try:
from airflow.utils import timezone
execution_date = timezone.parse(execution_date)
except ValueError:
error_message = (
'Given execution date, {}, could not be identified '
'as a date. Example date format: 2015-11-16T14:34:15+00:00'.
format(execution_date))
return ApiResponse.bad_request(error_message)
if request.args.get('try_number') is not None:
try_number = int(request.args.get('try_number'))
else:
try_number = None
session = settings.Session()
ti = session.query(TaskInstance).filter(
TaskInstance.dag_id == dag_id, TaskInstance.task_id == task_id,
TaskInstance.execution_date == dttm).first()
session.close()
if ti is None:
return ApiResponse.success(
"*** Task instance did not exist in the DB\n")
else:
dag = DagBag("dags").get_dag(dag_id)
ti.task = dag.get_task(ti.task_id)
logger = logging.getLogger('airflow.task')
task_log_reader = configuration.conf.get('core', 'task_log_reader')
handler = next((handler for handler in logger.handlers
if handler.name == task_log_reader), None)
logs = handler.read(ti, try_number)
if logs:
return ApiResponse.success(logs[0])
else:
return ApiResponse.success("No More Log")
except Exception as e:
error_message = "An error occurred while trying to find_task_log'" + "': " + str(
e)
logging.error(error_message)
return ApiResponse.server_error(error_message)
def running_dag_run(session):
dag = DagBag().get_dag("example_bash_operator")
task0 = dag.get_task("runme_0")
task1 = dag.get_task("runme_1")
execution_date = timezone.datetime(2016, 1, 9)
tis = [
TaskInstance(task0, execution_date, state="success"),
TaskInstance(task1, execution_date, state="failed"),
]
session.bulk_save_objects(tis)
dr = dag.create_dagrun(
state="running",
execution_date=execution_date,
run_id="test_dag_runs_action",
session=session,
)
return dr
