XD-DENG commented on a change in pull request #12704:
URL: https://github.com/apache/airflow/pull/12704#discussion_r532869104
##########
File path: airflow/cli/commands/task_command.py
##########
@@ -304,43 +309,38 @@ def _guess_debugger():
@cli_utils.action_logging
+@suppress_logs_and_warning()
def task_states_for_dag_run(args):
"""Get the status of all task instances in a DagRun"""
- session = settings.Session()
-
- tis = (
- session.query(
- TaskInstance.dag_id,
- TaskInstance.execution_date,
- TaskInstance.task_id,
- TaskInstance.state,
- TaskInstance.start_date,
- TaskInstance.end_date,
- )
- .filter(TaskInstance.dag_id == args.dag_id,
TaskInstance.execution_date == args.execution_date)
- .all()
- )
-
- if len(tis) == 0:
- raise AirflowException("DagRun does not exist.")
-
- formatted_rows = []
-
- for ti in tis:
- formatted_rows.append(
- (ti.dag_id, ti.execution_date, ti.task_id, ti.state,
ti.start_date, ti.end_date)
+ with create_session() as session:
+ tis = (
+ session.query(
+ TaskInstance.dag_id,
+ TaskInstance.execution_date,
+ TaskInstance.task_id,
+ TaskInstance.state,
+ TaskInstance.start_date,
+ TaskInstance.end_date,
+ )
+ .filter(TaskInstance.dag_id == args.dag_id,
TaskInstance.execution_date == args.execution_date)
+ .all()
)
- print(
- "\n%s"
- % tabulate(
- formatted_rows,
- ['dag', 'exec_date', 'task', 'state', 'start_date', 'end_date'],
- tablefmt=args.output,
+ if len(tis) == 0:
+ raise AirflowException("DagRun does not exist.")
+
+ AirflowConsole().print_as(
+ data=tis,
+ output=args.output,
+ mapper=lambda ti: {
+ "dag_id": ti.dag_id,
+ "execution_date": ti.execution_date.isoformat(),
+ "task_id": ti.task_id,
+ "state": ti.state,
+ "start_date": ti.start_date.isoformat(),
+ "end_date": ti.end_date.isoformat(),
Review comment:
Should guards against `None` be added? Similar to what you have in
`dag_list_dag_runs()`, i.e.
```python
... ...
"start_date": ti.start_date.isoformat() if ti.start_date
else '',
"end_date": ti.end_date.isoformat() if ti.end_date else '',
... ....
```
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
