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