diff --git a/UPDATING.md b/UPDATING.md
index e9b36bb66ae52..6af5e09be8b8b 100644
--- a/UPDATING.md
+++ b/UPDATING.md
@@ -52,6 +52,40 @@ assists users migrating to a new version.
## Master
+### Changes to output argument in commands
+
+From Airflow 2.0, We are replacing [tabulate](https://pypi.org/project/tabulate/) with [rich](https://github.com/willmcgugan/rich) to render commands output. Due to this change, the `--output` argument
+will no longer accept formats of tabulate tables. Instead, it now accepts:
+
+- `table` - will render the output in predefined table
+- `json` - will render the output as a json
+- `yaml` - will render the output as yaml
+
+By doing this we increased consistency and gave users possibility to manipulate the
+output programmatically (when using json or yaml).
+
+Affected commands:
+
+- `airflow dags list`
+- `airflow dags report`
+- `airflow dags list-runs`
+- `airflow dags list-jobs`
+- `airflow connections list`
+- `airflow connections get`
+- `airflow pools list`
+- `airflow pools get`
+- `airflow pools set`
+- `airflow pools delete`
+- `airflow pools import`
+- `airflow pools export`
+- `airflow role list`
+- `airflow providers list`
+- `airflow providers get`
+- `airflow providers hooks`
+- `airflow tasks states-for-dag-run`
+- `airflow users list`
+- `airflow variables list`
+
### Azure Wasb Hook does not work together with Snowflake hook
The WasbHook in Apache Airflow use a legacy version of Azure library. While the conflict is not
diff --git a/airflow/cli/cli_parser.py b/airflow/cli/cli_parser.py
index 6a13db0619b05..556af7a1b39cd 100644
--- a/airflow/cli/cli_parser.py
+++ b/airflow/cli/cli_parser.py
@@ -26,8 +26,6 @@
from functools import lru_cache
from typing import Callable, Dict, Iterable, List, NamedTuple, Optional, Set, Union
-from tabulate import tabulate_formats
-
from airflow import settings
from airflow.cli.commands.legacy_commands import check_legacy_command
from airflow.configuration import conf
@@ -176,13 +174,10 @@ def positive_int(value):
)
ARG_OUTPUT = Arg(
("--output",),
- help=(
- "Output table format. The specified value is passed to "
- "the tabulate module (https://pypi.org/project/tabulate/). "
- ),
- metavar="FORMAT",
- choices=tabulate_formats,
- default="plain",
+ help=("Output format. Allowed values: json, yaml, table (default: table)"),
+ metavar="(table, json, yaml)",
+ choices=("table", "json", "yaml"),
+ default="table",
)
ARG_COLOR = Arg(
('--color',),
@@ -1033,7 +1028,7 @@ class GroupCommand(NamedTuple):
name='list',
help='List variables',
func=lazy_load_command('airflow.cli.commands.variable_command.variables_list'),
- args=(),
+ args=(ARG_OUTPUT,),
),
ActionCommand(
name='get',
@@ -1110,7 +1105,7 @@ class GroupCommand(NamedTuple):
name='get',
help='Get a connection',
func=lazy_load_command('airflow.cli.commands.connection_command.connections_get'),
- args=(ARG_CONN_ID, ARG_COLOR),
+ args=(ARG_CONN_ID, ARG_COLOR, ARG_OUTPUT),
),
ActionCommand(
name='list',
diff --git a/airflow/cli/commands/cheat_sheet_command.py b/airflow/cli/commands/cheat_sheet_command.py
index 28ee67bab3af0..263688d615e3a 100644
--- a/airflow/cli/commands/cheat_sheet_command.py
+++ b/airflow/cli/commands/cheat_sheet_command.py
@@ -21,9 +21,11 @@
from airflow.cli.cli_parser import ActionCommand, GroupCommand, airflow_commands
from airflow.cli.simple_table import SimpleTable
+from airflow.utils.cli import suppress_logs_and_warning
from airflow.utils.helpers import partition
+@suppress_logs_and_warning()
def cheat_sheet(args):
"""Display cheat-sheet."""
display_commands_index()
diff --git a/airflow/cli/commands/connection_command.py b/airflow/cli/commands/connection_command.py
index 279a7e01bee9d..dd96f9f43e43a 100644
--- a/airflow/cli/commands/connection_command.py
+++ b/airflow/cli/commands/connection_command.py
@@ -19,78 +19,54 @@
import json
import os
import sys
-from typing import List
+from typing import Any, Dict, List
from urllib.parse import urlparse, urlunparse
-import pygments
import yaml
-from pygments.lexers.data import YamlLexer
from sqlalchemy.orm import exc
-from tabulate import tabulate
+from airflow.cli.simple_table import AirflowConsole
from airflow.exceptions import AirflowNotFoundException
from airflow.hooks.base_hook import BaseHook
from airflow.models import Connection
from airflow.utils import cli as cli_utils
-from airflow.utils.cli import should_use_colors
-from airflow.utils.code_utils import get_terminal_formatter
+from airflow.utils.cli import suppress_logs_and_warning
from airflow.utils.session import create_session
-def _tabulate_connection(conns: List[Connection], tablefmt: str):
- tabulate_data = [
- {
- 'Conn Id': conn.conn_id,
- 'Conn Type': conn.conn_type,
- 'Description': conn.description,
- 'Host': conn.host,
- 'Port': conn.port,
- 'Is Encrypted': conn.is_encrypted,
- 'Is Extra Encrypted': conn.is_encrypted,
- 'Extra': conn.extra,
- }
- for conn in conns
- ]
-
- msg = tabulate(tabulate_data, tablefmt=tablefmt, headers='keys')
- return msg
-
-
-def _yamulate_connection(conn: Connection):
- yaml_data = {
- 'Id': conn.id,
- 'Conn Id': conn.conn_id,
- 'Conn Type': conn.conn_type,
- 'Description': conn.description,
- 'Host': conn.host,
- 'Schema': conn.schema,
- 'Login': conn.login,
- 'Password': conn.password,
- 'Port': conn.port,
- 'Is Encrypted': conn.is_encrypted,
- 'Is Extra Encrypted': conn.is_encrypted,
- 'Extra': conn.extra_dejson,
- 'URI': conn.get_uri(),
+def _connection_mapper(conn: Connection) -> Dict[str, Any]:
+ return {
+ 'id': conn.id,
+ 'conn_id': conn.conn_id,
+ 'conn_type': conn.conn_type,
+ 'description': conn.description,
+ 'host': conn.host,
+ 'schema': conn.schema,
+ 'login': conn.login,
+ 'password': conn.password,
+ 'port': conn.port,
+ 'is_encrypted': conn.is_encrypted,
+ 'is_extra_encrypted': conn.is_encrypted,
+ 'extra_dejson': conn.extra_dejson,
+ 'get_uri': conn.get_uri(),
}
- return yaml.safe_dump(yaml_data, sort_keys=False)
+@suppress_logs_and_warning()
def connections_get(args):
"""Get a connection."""
try:
conn = BaseHook.get_connection(args.conn_id)
except AirflowNotFoundException:
raise SystemExit("Connection not found.")
-
- yaml_content = _yamulate_connection(conn)
- if should_use_colors(args):
- yaml_content = pygments.highlight(
- code=yaml_content, formatter=get_terminal_formatter(), lexer=YamlLexer()
- )
-
- print(yaml_content)
+ AirflowConsole().print_as(
+ data=[conn],
+ output=args.output,
+ mapper=_connection_mapper,
+ )
+@suppress_logs_and_warning()
def connections_list(args):
"""Lists all connections at the command line"""
with create_session() as session:
@@ -99,9 +75,11 @@ def connections_list(args):
query = query.filter(Connection.conn_id == args.conn_id)
conns = query.all()
- tablefmt = args.output
- msg = _tabulate_connection(conns, tablefmt)
- print(msg)
+ AirflowConsole().print_as(
+ data=conns,
+ output=args.output,
+ mapper=_connection_mapper,
+ )
def _format_connections(conns: List[Connection], fmt: str) -> str:
diff --git a/airflow/cli/commands/dag_command.py b/airflow/cli/commands/dag_command.py
index c02dd9e6f1cb9..f7b57bc895878 100644
--- a/airflow/cli/commands/dag_command.py
+++ b/airflow/cli/commands/dag_command.py
@@ -16,19 +16,19 @@
# under the License.
"""Dag sub-commands"""
+import ast
import errno
import json
import logging
import signal
import subprocess
import sys
-from typing import List
from graphviz.dot import Dot
-from tabulate import tabulate
from airflow import settings
from airflow.api.client import get_current_api_client
+from airflow.cli.simple_table import AirflowConsole
from airflow.configuration import conf
from airflow.exceptions import AirflowException, BackfillUnfinished
from airflow.executors.debug_executor import DebugExecutor
@@ -36,40 +36,18 @@
from airflow.models import DagBag, DagModel, DagRun, TaskInstance
from airflow.models.dag import DAG
from airflow.utils import cli as cli_utils
-from airflow.utils.cli import get_dag, get_dag_by_file_location, process_subdir, sigint_handler
+from airflow.utils.cli import (
+ get_dag,
+ get_dag_by_file_location,
+ process_subdir,
+ sigint_handler,
+ suppress_logs_and_warning,
+)
from airflow.utils.dot_renderer import render_dag
from airflow.utils.session import create_session, provide_session
from airflow.utils.state import State
-def _tabulate_dag_runs(dag_runs: List[DagRun], tablefmt: str = "fancy_grid") -> str:
- tabulate_data = (
- {
- 'ID': dag_run.id,
- 'Run ID': dag_run.run_id,
- 'State': dag_run.state,
- 'DAG ID': dag_run.dag_id,
- 'Execution date': dag_run.execution_date.isoformat(),
- 'Start date': dag_run.start_date.isoformat() if dag_run.start_date else '',
- 'End date': dag_run.end_date.isoformat() if dag_run.end_date else '',
- }
- for dag_run in dag_runs
- )
- return tabulate(tabular_data=tabulate_data, tablefmt=tablefmt)
-
-
-def _tabulate_dags(dags: List[DAG], tablefmt: str = "fancy_grid") -> str:
- tabulate_data = (
- {
- 'DAG ID': dag.dag_id,
- 'Filepath': dag.filepath,
- 'Owner': dag.owner,
- }
- for dag in sorted(dags, key=lambda d: d.dag_id)
- )
- return tabulate(tabular_data=tabulate_data, tablefmt=tablefmt, headers='keys')
-
-
@cli_utils.action_logging
def dag_backfill(args, dag=None):
"""Creates backfill job or dry run for a DAG"""
@@ -293,21 +271,41 @@ def dag_next_execution(args):
@cli_utils.action_logging
+@suppress_logs_and_warning()
def dag_list_dags(args):
"""Displays dags with or without stats at the command line"""
dagbag = DagBag(process_subdir(args.subdir))
- dags = dagbag.dags.values()
- print(_tabulate_dags(dags, tablefmt=args.output))
+ AirflowConsole().print_as(
+ data=sorted(dagbag.dags.values(), key=lambda d: d.dag_id),
+ output=args.output,
+ mapper=lambda x: {
+ "dag_id": x.dag_id,
+ "filepath": x.filepath,
+ "owner": x.owner,
+ },
+ )
@cli_utils.action_logging
+@suppress_logs_and_warning()
def dag_report(args):
"""Displays dagbag stats at the command line"""
dagbag = DagBag(process_subdir(args.subdir))
- print(tabulate(dagbag.dagbag_stats, headers="keys", tablefmt=args.output))
+ AirflowConsole().print_as(
+ data=dagbag.dagbag_stats,
+ output=args.output,
+ mapper=lambda x: {
+ "file": x.file,
+ "duration": x.duration,
+ "dag_num": x.dag_num,
+ "task_num": x.task_num,
+ "dags": sorted(ast.literal_eval(x.dags)),
+ },
+ )
@cli_utils.action_logging
+@suppress_logs_and_warning()
def dag_list_jobs(args, dag=None):
"""Lists latest n jobs"""
queries = []
@@ -324,6 +322,7 @@ def dag_list_jobs(args, dag=None):
if args.state:
queries.append(BaseJob.state == args.state)
+ fields = ['dag_id', 'state', 'job_type', 'start_date', 'end_date']
with create_session() as session:
all_jobs = (
session.query(BaseJob)
@@ -332,15 +331,16 @@ def dag_list_jobs(args, dag=None):
.limit(args.limit)
.all()
)
- fields = ['dag_id', 'state', 'job_type', 'start_date', 'end_date']
- all_jobs = [[job.__getattribute__(field) for field in fields] for job in all_jobs]
- msg = tabulate(
- all_jobs, [field.capitalize().replace('_', ' ') for field in fields], tablefmt=args.output
- )
- print(msg)
+ all_jobs = [{f: str(job.__getattribute__(f)) for f in fields} for job in all_jobs]
+
+ AirflowConsole().print_as(
+ data=all_jobs,
+ output=args.output,
+ )
@cli_utils.action_logging
+@suppress_logs_and_warning()
def dag_list_dag_runs(args, dag=None):
"""Lists dag runs for a given DAG"""
if dag:
@@ -361,13 +361,19 @@ def dag_list_dag_runs(args, dag=None):
execution_end_date=args.end_date,
)
- if not dag_runs:
- print(f'No dag runs for {args.dag_id}')
- return
-
dag_runs.sort(key=lambda x: x.execution_date, reverse=True)
- table = _tabulate_dag_runs(dag_runs, tablefmt=args.output)
- print(table)
+ AirflowConsole().print_as(
+ data=dag_runs,
+ output=args.output,
+ mapper=lambda dr: {
+ "dag_id": dr.dag_id,
+ "run_id": dr.run_id,
+ "state": dr.state,
+ "execution_date": dr.execution_date.isoformat(),
+ "start_date": dr.start_date.isoformat() if dr.start_date else '',
+ "end_date": dr.end_date.isoformat() if dr.end_date else '',
+ },
+ )
@provide_session
diff --git a/airflow/cli/commands/info_command.py b/airflow/cli/commands/info_command.py
index bde7bb47ba47d..a8686bd976f31 100644
--- a/airflow/cli/commands/info_command.py
+++ b/airflow/cli/commands/info_command.py
@@ -33,6 +33,7 @@
from airflow.cli.simple_table import SimpleTable
from airflow.providers_manager import ProvidersManager
from airflow.typing_compat import Protocol
+from airflow.utils.cli import suppress_logs_and_warning
from airflow.version import version as airflow_version
log = logging.getLogger(__name__)
@@ -412,6 +413,7 @@ def _send_report_to_fileio(info):
print(str(ex))
+@suppress_logs_and_warning()
def show_info(args):
"""Show information related to Airflow, system and other."""
# Enforce anonymization, when file_io upload is tuned on.
diff --git a/airflow/cli/commands/plugins_command.py b/airflow/cli/commands/plugins_command.py
index 98728eae0da0b..50a321830ee1c 100644
--- a/airflow/cli/commands/plugins_command.py
+++ b/airflow/cli/commands/plugins_command.py
@@ -25,6 +25,8 @@
from airflow.plugins_manager import PluginsDirectorySource
# list to maintain the order of items.
+from airflow.utils.cli import suppress_logs_and_warning
+
PLUGINS_MANAGER_ATTRIBUTES_TO_DUMP = [
"plugins",
"import_errors",
@@ -64,6 +66,7 @@ def _join_plugins_names(value: Union[List[Any], Any]) -> str:
return ",".join(_get_name(v) for v in value)
+@suppress_logs_and_warning()
def dump_plugins(args):
"""Dump plugins information"""
plugins_manager.ensure_plugins_loaded()
diff --git a/airflow/cli/commands/pool_command.py b/airflow/cli/commands/pool_command.py
index 641b97723163e..20b7959d8f74d 100644
--- a/airflow/cli/commands/pool_command.py
+++ b/airflow/cli/commands/pool_command.py
@@ -21,61 +21,77 @@
import sys
from json import JSONDecodeError
-from tabulate import tabulate
-
from airflow.api.client import get_current_api_client
+from airflow.cli.simple_table import AirflowConsole
+from airflow.exceptions import PoolNotFound
from airflow.utils import cli as cli_utils
+from airflow.utils.cli import suppress_logs_and_warning
-def _tabulate_pools(pools, tablefmt="fancy_grid"):
- return "\n%s" % tabulate(pools, ['Pool', 'Slots', 'Description'], tablefmt=tablefmt)
+def _show_pools(pools, output):
+ AirflowConsole().print_as(
+ data=pools,
+ output=output,
+ mapper=lambda x: {
+ "pool": x[0],
+ "slots": x[1],
+ "description": x[2],
+ },
+ )
+@suppress_logs_and_warning()
def pool_list(args):
"""Displays info of all the pools"""
api_client = get_current_api_client()
pools = api_client.get_pools()
- print(_tabulate_pools(pools=pools, tablefmt=args.output))
+ _show_pools(pools=pools, output=args.output)
+@suppress_logs_and_warning()
def pool_get(args):
"""Displays pool info by a given name"""
api_client = get_current_api_client()
pools = [api_client.get_pool(name=args.pool)]
- print(_tabulate_pools(pools=pools, tablefmt=args.output))
+ _show_pools(pools=pools, output=args.output)
@cli_utils.action_logging
+@suppress_logs_and_warning()
def pool_set(args):
"""Creates new pool with a given name and slots"""
api_client = get_current_api_client()
- pools = [api_client.create_pool(name=args.pool, slots=args.slots, description=args.description)]
- print(_tabulate_pools(pools=pools, tablefmt=args.output))
+ api_client.create_pool(name=args.pool, slots=args.slots, description=args.description)
+ print("Pool created")
@cli_utils.action_logging
+@suppress_logs_and_warning()
def pool_delete(args):
"""Deletes pool by a given name"""
api_client = get_current_api_client()
- pools = [api_client.delete_pool(name=args.pool)]
- print(_tabulate_pools(pools=pools, tablefmt=args.output))
+ try:
+ api_client.delete_pool(name=args.pool)
+ print("Pool deleted")
+ except PoolNotFound:
+ sys.exit(f"Pool {args.pool} does not exist")
@cli_utils.action_logging
+@suppress_logs_and_warning()
def pool_import(args):
"""Imports pools from the file"""
if not os.path.exists(args.file):
sys.exit("Missing pools file.")
- pools, failed = pool_import_helper(args.file)
- print(_tabulate_pools(pools=pools, tablefmt=args.output))
+ _, failed = pool_import_helper(args.file)
if len(failed) > 0:
- sys.exit("Failed to update pool(s): {}".format(", ".join(failed)))
+ sys.exit(f"Failed to update pool(s): {', '.join(failed)}")
def pool_export(args):
"""Exports all of the pools to the file"""
pools = pool_export_helper(args.file)
- print(_tabulate_pools(pools=pools, tablefmt=args.output))
+ print(f"Exported {len(pools)} pools to {args.file}")
def pool_import_helper(filepath):
@@ -108,5 +124,4 @@ def pool_export_helper(filepath):
pool_dict[pool[0]] = {"slots": pool[1], "description": pool[2]}
with open(filepath, 'w') as poolfile:
poolfile.write(json.dumps(pool_dict, sort_keys=True, indent=4))
- print("{} pools successfully exported to {}".format(len(pool_dict), filepath))
return pools
diff --git a/airflow/cli/commands/provider_command.py b/airflow/cli/commands/provider_command.py
index a2df92fa0c6e7..c7830789c60d3 100644
--- a/airflow/cli/commands/provider_command.py
+++ b/airflow/cli/commands/provider_command.py
@@ -15,71 +15,60 @@
# specific language governing permissions and limitations
# under the License.
"""Providers sub-commands"""
-from typing import Dict, List, Tuple
-
-import pygments
-import yaml
-from pygments.lexers.data import YamlLexer
-from tabulate import tabulate
+import re
+from airflow.cli.simple_table import AirflowConsole
from airflow.providers_manager import ProvidersManager
-from airflow.utils.cli import should_use_colors
-from airflow.utils.code_utils import get_terminal_formatter
-
+from airflow.utils.cli import suppress_logs_and_warning
-def _tabulate_providers(providers: List[Dict], tablefmt: str):
- tabulate_data = [
- {
- 'Provider name': provider['package-name'],
- 'Description': provider['description'],
- 'Version': provider['versions'][0],
- }
- for version, provider in providers
- ]
- return tabulate(tabulate_data, tablefmt=tablefmt, headers='keys')
+def _remove_rst_syntax(value: str) -> str:
+ return re.sub("[`_<>]", "", value.strip(" \n."))
+@suppress_logs_and_warning()
def provider_get(args):
"""Get a provider info."""
providers = ProvidersManager().providers
if args.provider_name in providers:
- provider_version, provider_info = providers[args.provider_name]
- print("#")
- print(f"# Provider: {args.provider_name}")
- print(f"# Version: {provider_version}")
- print("#")
+ provider_version = providers[args.provider_name][0]
+ provider_info = providers[args.provider_name][1]
if args.full:
- yaml_content = yaml.dump(provider_info)
- if should_use_colors(args):
- yaml_content = pygments.highlight(
- code=yaml_content, formatter=get_terminal_formatter(), lexer=YamlLexer()
- )
- print(yaml_content)
+ provider_info["description"] = _remove_rst_syntax(provider_info["description"])
+ AirflowConsole().print_as(
+ data=[provider_info],
+ output=args.output,
+ )
+ else:
+ print(f"Provider: {args.provider_name}")
+ print(f"Version: {provider_version}")
else:
raise SystemExit(f"No such provider installed: {args.provider_name}")
+@suppress_logs_and_warning()
def providers_list(args):
"""Lists all providers at the command line"""
- print(_tabulate_providers(ProvidersManager().providers.values(), args.output))
-
-
-def _tabulate_hooks(hook_items: Tuple[str, Tuple[str, str]], tablefmt: str):
- tabulate_data = [
- {
- 'Connection type': hook_item[0],
- 'Hook class': hook_item[1][0],
- 'Hook connection attribute name': hook_item[1][1],
- }
- for hook_item in hook_items
- ]
-
- msg = tabulate(tabulate_data, tablefmt=tablefmt, headers='keys')
- return msg
+ AirflowConsole().print_as(
+ data=ProvidersManager().providers.values(),
+ output=args.output,
+ mapper=lambda x: {
+ "package_name": x[1]["package-name"],
+ "description": _remove_rst_syntax(x[1]["description"]),
+ "version": x[0],
+ },
+ )
+@suppress_logs_and_warning()
def hooks_list(args):
"""Lists all hooks at the command line"""
- msg = _tabulate_hooks(ProvidersManager().hooks.items(), args.output)
- print(msg)
+ AirflowConsole().print_as(
+ data=ProvidersManager().hooks.items(),
+ output=args.output,
+ mapper=lambda x: {
+ "connection_type": x[0],
+ "class": x[1][0],
+ "conn_attribute_name": x[1][1],
+ },
+ )
diff --git a/airflow/cli/commands/role_command.py b/airflow/cli/commands/role_command.py
index fb34d8c9a6003..a76061098ed81 100644
--- a/airflow/cli/commands/role_command.py
+++ b/airflow/cli/commands/role_command.py
@@ -17,20 +17,21 @@
# under the License.
#
"""Roles sub-commands"""
-from tabulate import tabulate
+from airflow.cli.simple_table import AirflowConsole
from airflow.utils import cli as cli_utils
+from airflow.utils.cli import suppress_logs_and_warning
from airflow.www.app import cached_app
+@suppress_logs_and_warning()
def roles_list(args):
"""Lists all existing roles"""
appbuilder = cached_app().appbuilder # pylint: disable=no-member
roles = appbuilder.sm.get_all_roles()
- print("Existing roles:\n")
- role_names = sorted([[r.name] for r in roles])
- msg = tabulate(role_names, headers=['Role'], tablefmt=args.output)
- print(msg)
+ AirflowConsole().print_as(
+ data=sorted([r.name for r in roles]), output=args.output, mapper=lambda x: {"name": x}
+ )
@cli_utils.action_logging
diff --git a/airflow/cli/commands/task_command.py b/airflow/cli/commands/task_command.py
index a3686fb1f584f..356f60cfe5950 100644
--- a/airflow/cli/commands/task_command.py
+++ b/airflow/cli/commands/task_command.py
@@ -24,9 +24,8 @@
from contextlib import contextmanager, redirect_stderr, redirect_stdout
from typing import List
-from tabulate import tabulate
-
from airflow import settings
+from airflow.cli.simple_table import AirflowConsole
from airflow.configuration import conf
from airflow.exceptions import AirflowException
from airflow.executors.executor_loader import ExecutorLoader
@@ -36,7 +35,13 @@
from airflow.ti_deps.dep_context import DepContext
from airflow.ti_deps.dependencies_deps import SCHEDULER_QUEUED_DEPS
from airflow.utils import cli as cli_utils
-from airflow.utils.cli import get_dag, get_dag_by_file_location, get_dag_by_pickle, get_dags
+from airflow.utils.cli import (
+ get_dag,
+ get_dag_by_file_location,
+ get_dag_by_pickle,
+ get_dags,
+ suppress_logs_and_warning,
+)
from airflow.utils.log.logging_mixin import StreamLogWriter
from airflow.utils.net import get_hostname
from airflow.utils.session import create_session
@@ -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() if ti.start_date else "",
+ "end_date": ti.end_date.isoformat() if ti.end_date else "",
+ },
)
- )
-
- session.close()
@cli_utils.action_logging
diff --git a/airflow/cli/commands/user_command.py b/airflow/cli/commands/user_command.py
index 5481c7589ce5c..a5f8ec435c482 100644
--- a/airflow/cli/commands/user_command.py
+++ b/airflow/cli/commands/user_command.py
@@ -24,20 +24,22 @@
import string
import sys
-from tabulate import tabulate
-
+from airflow.cli.simple_table import AirflowConsole
from airflow.utils import cli as cli_utils
+from airflow.utils.cli import suppress_logs_and_warning
from airflow.www.app import cached_app
+@suppress_logs_and_warning()
def users_list(args):
"""Lists users at the command line"""
appbuilder = cached_app().appbuilder # pylint: disable=no-member
users = appbuilder.sm.get_all_users()
fields = ['id', 'username', 'email', 'first_name', 'last_name', 'roles']
- users = [[user.__getattribute__(field) for field in fields] for user in users]
- msg = tabulate(users, [field.capitalize().replace('_', ' ') for field in fields], tablefmt=args.output)
- print(msg)
+
+ AirflowConsole().print_as(
+ data=users, output=args.output, mapper=lambda x: {f: x.__getattribute__(f) for f in fields}
+ )
@cli_utils.action_logging
diff --git a/airflow/cli/commands/variable_command.py b/airflow/cli/commands/variable_command.py
index 8478bf83e0192..b3f34ed9b47ec 100644
--- a/airflow/cli/commands/variable_command.py
+++ b/airflow/cli/commands/variable_command.py
@@ -21,16 +21,19 @@
import sys
from json import JSONDecodeError
+from airflow.cli.simple_table import AirflowConsole
from airflow.models import Variable
from airflow.utils import cli as cli_utils
+from airflow.utils.cli import suppress_logs_and_warning
from airflow.utils.session import create_session
+@suppress_logs_and_warning()
def variables_list(args):
"""Displays all of the variables"""
with create_session() as session:
variables = session.query(Variable)
- print("\n".join(var.key for var in variables))
+ AirflowConsole().print_as(data=variables, output=args.output, mapper=lambda x: {"key": x.key})
def variables_get(args):
diff --git a/airflow/cli/simple_table.py b/airflow/cli/simple_table.py
index abfb705884122..8cdb6de61083c 100644
--- a/airflow/cli/simple_table.py
+++ b/airflow/cli/simple_table.py
@@ -14,11 +14,78 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
+import json
+from typing import Any, Callable, Dict, List, Optional, Union
+import yaml
from rich.box import ASCII_DOUBLE_HEAD
+from rich.console import Console
+from rich.syntax import Syntax
from rich.table import Table
+class AirflowConsole(Console):
+ """Airflow rich console"""
+
+ def print_as_json(self, data: Dict):
+ """Renders dict as json text representation"""
+ json_content = json.dumps(data)
+ self.print(Syntax(json_content, "json", theme="ansi_dark"), soft_wrap=True)
+
+ def print_as_yaml(self, data: Dict):
+ """Renders dict as yaml text representation"""
+ yaml_content = yaml.dump(data)
+ self.print(Syntax(yaml_content, "yaml", theme="ansi_dark"), soft_wrap=True)
+
+ def print_as_table(self, data: List[Dict]):
+ """Renders list of dictionaries as table"""
+ if not data:
+ self.print("No data found")
+ return
+
+ table = SimpleTable(
+ show_header=True,
+ )
+ for col in data[0].keys():
+ table.add_column(col)
+
+ for row in data:
+ table.add_row(*[str(d) for d in row.values()])
+ self.print(table)
+
+ def _normalize_data(self, value: Any, output: str) -> Union[list, str, dict]:
+ if isinstance(value, (tuple, list)):
+ if output == "table":
+ return ",".join(self._normalize_data(x, output) for x in value)
+ return [self._normalize_data(x, output) for x in value]
+ if isinstance(value, dict) and output != "table":
+ return {k: self._normalize_data(v, output) for k, v in value.items()}
+ return str(value)
+
+ def print_as(self, data: List[Union[Dict, Any]], output: str, mapper: Optional[Callable] = None):
+ """Prints provided using format specified by output argument"""
+ output_to_renderer = {
+ "json": self.print_as_json,
+ "yaml": self.print_as_yaml,
+ "table": self.print_as_table,
+ }
+ renderer = output_to_renderer.get(output)
+ if not renderer:
+ raise ValueError(
+ f"Unknown formatter: {output}. Allowed options: {list(output_to_renderer.keys())}"
+ )
+
+ if not all(isinstance(d, dict) for d in data) and not mapper:
+ raise ValueError("To tabulate non-dictionary data you need to provide `mapper` function")
+
+ if mapper:
+ dict_data: List[Dict] = [mapper(d) for d in data]
+ else:
+ dict_data: List[Dict] = data
+ dict_data = [{k: self._normalize_data(v, output) for k, v in d.items()} for d in dict_data]
+ renderer(dict_data)
+
+
class SimpleTable(Table):
"""A rich Table with some default hardcoded for consistency."""
diff --git a/airflow/utils/cli.py b/airflow/utils/cli.py
index 0578acb7597bf..292a896259ec6 100644
--- a/airflow/utils/cli.py
+++ b/airflow/utils/cli.py
@@ -17,7 +17,7 @@
# under the License.
#
"""Utilities module for cli"""
-
+import contextlib
import functools
import getpass
import json
@@ -28,6 +28,7 @@
import sys
import threading
import traceback
+import warnings
from argparse import Namespace
from datetime import datetime
from typing import Callable, Optional, TypeVar, cast
@@ -272,3 +273,15 @@ def should_use_colors(args) -> bool:
if args.color == ColorMode.OFF:
return False
return is_terminal_support_colors()
+
+
+@contextlib.contextmanager
+def suppress_logs_and_warning():
+ """Context manager to suppress logging and warning messages"""
+ with warnings.catch_warnings():
+ warnings.simplefilter("ignore")
+ logging.disable(logging.CRITICAL)
+ yield
+ # logging output again depends on the effective
+ # levels of individual loggers
+ logging.disable(logging.NOTSET)
diff --git a/docs/apache-airflow/usage-cli.rst b/docs/apache-airflow/usage-cli.rst
index a41fcf3fbf1aa..5dadab8466952 100644
--- a/docs/apache-airflow/usage-cli.rst
+++ b/docs/apache-airflow/usage-cli.rst
@@ -174,3 +174,40 @@ You will see a similar result as in the screenshot below.
.. figure:: img/usage_cli_imgcat.png
Preview of DAG in iTerm2
+
+Formatting commands output
+--------------------------
+
+Some Airflow commands like ``airflow dags list`` or ``airflow tasks states-for-dag-run`` support ``--output`` flag
+which allow users to change the formatting of command's output. Possible options:
+
+ - ``table`` - renders the information as a plain text table
+ - ``json`` - renders the information in form of json string
+ - ``yaml`` - render the information in form of valid yaml
+
+Both ``json`` and ``yaml`` formats make it easier to manipulate the data using command line tools like
+`jq `__ or `yq `__. For example:
+
+.. code-block:: bash
+
+ airflow tasks states-for-dag-run example_complex 2020-11-13T00:00:00+00:00 --output json | jq ".[] | {sd: .start_date, ed: .end_date}"
+ {
+ "sd": "2020-11-29T14:53:46.811030+00:00",
+ "ed": "2020-11-29T14:53:46.974545+00:00"
+ }
+ {
+ "sd": "2020-11-29T14:53:56.926441+00:00",
+ "ed": "2020-11-29T14:53:57.118781+00:00"
+ }
+ {
+ "sd": "2020-11-29T14:53:56.915802+00:00",
+ "ed": "2020-11-29T14:53:57.125230+00:00"
+ }
+ {
+ "sd": "2020-11-29T14:53:56.922131+00:00",
+ "ed": "2020-11-29T14:53:57.129091+00:00"
+ }
+ {
+ "sd": "2020-11-29T14:53:56.931243+00:00",
+ "ed": "2020-11-29T14:53:57.126306+00:00"
+ }
diff --git a/scripts/in_container/run_install_and_test_provider_packages.sh b/scripts/in_container/run_install_and_test_provider_packages.sh
index b151505813cfd..4a6f96e23400a 100755
--- a/scripts/in_container/run_install_and_test_provider_packages.sh
+++ b/scripts/in_container/run_install_and_test_provider_packages.sh
@@ -83,7 +83,7 @@ function discover_all_provider_packages() {
local expected_number_of_providers=60
local actual_number_of_providers
- actual_number_of_providers=$(airflow providers list --output simple | grep -c apache-airflow-providers | xargs)
+ actual_number_of_providers=$(airflow providers list --output table | grep -c apache-airflow-providers | xargs)
if [[ ${actual_number_of_providers} != "${expected_number_of_providers}" ]]; then
>&2 echo "ERROR! Number of providers installed is wrong!"
>&2 echo "Expected number was '${expected_number_of_providers}' and got '${actual_number_of_providers}'"
@@ -102,7 +102,7 @@ function discover_all_hooks() {
local expected_number_of_hooks=33
local actual_number_of_hooks
- actual_number_of_hooks=$(airflow providers hooks --output simple | grep -c conn_id | xargs)
+ actual_number_of_hooks=$(airflow providers hooks --output table | grep -c conn_id | xargs)
if [[ ${actual_number_of_hooks} != "${expected_number_of_hooks}" ]]; then
>&2 echo "ERROR! Number of hooks registered is wrong!"
>&2 echo "Expected number was '${expected_number_of_hooks}' and got '${actual_number_of_hooks}'"
diff --git a/tests/cli/commands/test_connection_command.py b/tests/cli/commands/test_connection_command.py
index adc4b423394c6..700d8191e1702 100644
--- a/tests/cli/commands/test_connection_command.py
+++ b/tests/cli/commands/test_connection_command.py
@@ -43,10 +43,10 @@ def tearDown(self):
def test_cli_connection_get(self):
with redirect_stdout(io.StringIO()) as stdout:
connection_command.connections_get(
- self.parser.parse_args(["connections", "get", "google_cloud_default"])
+ self.parser.parse_args(["connections", "get", "google_cloud_default", "--output", "json"])
)
stdout = stdout.getvalue()
- self.assertIn("URI: google-cloud-platform:///default", stdout)
+ self.assertIn("google-cloud-platform:///default", stdout)
def test_cli_connection_get_invalid(self):
with self.assertRaisesRegex(SystemExit, re.escape("Connection not found.")):
@@ -112,36 +112,27 @@ def setUp(self):
def tearDown(self):
clear_db_connections()
- def test_cli_connections_list(self):
- with redirect_stdout(io.StringIO()) as stdout:
- connection_command.connections_list(self.parser.parse_args(["connections", "list"]))
- stdout = stdout.getvalue()
- lines = stdout.split("\n")
-
- for conn_id, conn_type in self.EXPECTED_CONS:
- self.assertTrue(any(conn_id in line and conn_type in line for line in lines))
-
- def test_cli_connections_list_as_tsv(self):
- args = self.parser.parse_args(["connections", "list", "--output", "tsv"])
-
+ def test_cli_connections_list_as_json(self):
+ args = self.parser.parse_args(["connections", "list", "--output", "json"])
with redirect_stdout(io.StringIO()) as stdout:
connection_command.connections_list(args)
+ print(stdout.getvalue())
stdout = stdout.getvalue()
- lines = stdout.split("\n")
for conn_id, conn_type in self.EXPECTED_CONS:
- self.assertTrue(any(conn_id in line and conn_type in line for line in lines))
+ self.assertIn(conn_type, stdout)
+ self.assertIn(conn_id, stdout)
def test_cli_connections_filter_conn_id(self):
- args = self.parser.parse_args(["connections", "list", "--output", "tsv", '--conn-id', 'http_default'])
+ args = self.parser.parse_args(
+ ["connections", "list", "--output", "json", '--conn-id', 'http_default']
+ )
with redirect_stdout(io.StringIO()) as stdout:
connection_command.connections_list(args)
stdout = stdout.getvalue()
- lines = stdout.split("\n")
- conn_ids = [line.split("\t", 2)[0].strip() for line in lines[1:] if line]
- self.assertEqual(conn_ids, ['http_default'])
+ self.assertIn("http_default", stdout)
class TestCliExportConnections(unittest.TestCase):
diff --git a/tests/cli/commands/test_dag_command.py b/tests/cli/commands/test_dag_command.py
index 331bbf811292a..55f4226d51af5 100644
--- a/tests/cli/commands/test_dag_command.py
+++ b/tests/cli/commands/test_dag_command.py
@@ -47,6 +47,7 @@
)
+# TODO: Check if tests needs side effects - locally there's missing DAG
class TestCliDags(unittest.TestCase):
@classmethod
def setUpClass(cls):
@@ -328,22 +329,22 @@ def test_next_execution(self):
@conf_vars({('core', 'load_examples'): 'true'})
def test_cli_report(self):
- args = self.parser.parse_args(['dags', 'report'])
+ args = self.parser.parse_args(['dags', 'report', '--output', 'json'])
with contextlib.redirect_stdout(io.StringIO()) as temp_stdout:
dag_command.dag_report(args)
out = temp_stdout.getvalue()
- self.assertIn("airflow/example_dags/example_complex.py ", out)
- self.assertIn("['example_complex']", out)
+ self.assertIn("airflow/example_dags/example_complex.py", out)
+ self.assertIn("example_complex", out)
@conf_vars({('core', 'load_examples'): 'true'})
def test_cli_list_dags(self):
- args = self.parser.parse_args(['dags', 'list', '--output=fancy_grid'])
+ args = self.parser.parse_args(['dags', 'list', '--output', 'yaml'])
with contextlib.redirect_stdout(io.StringIO()) as temp_stdout:
dag_command.dag_list_dags(args)
out = temp_stdout.getvalue()
- self.assertIn("Owner", out)
- self.assertIn("│ airflow │", out)
+ self.assertIn("owner", out)
+ self.assertIn("airflow", out)
self.assertIn("airflow/example_dags/example_complex.py", out)
def test_cli_list_dag_runs(self):
@@ -383,7 +384,7 @@ def test_cli_list_jobs_with_args(self):
'--limit',
'100',
'--output',
- 'tsv',
+ 'json',
]
)
dag_command.dag_list_jobs(args)
diff --git a/tests/cli/commands/test_pool_command.py b/tests/cli/commands/test_pool_command.py
index 5d2e875461631..d40e18786a1cf 100644
--- a/tests/cli/commands/test_pool_command.py
+++ b/tests/cli/commands/test_pool_command.py
@@ -56,14 +56,13 @@ def _cleanup(session=None):
def test_pool_list(self):
pool_command.pool_set(self.parser.parse_args(['pools', 'set', 'foo', '1', 'test']))
- stdout = io.StringIO()
- with redirect_stdout(stdout):
+ with redirect_stdout(io.StringIO()) as stdout:
pool_command.pool_list(self.parser.parse_args(['pools', 'list']))
self.assertIn('foo', stdout.getvalue())
def test_pool_list_with_args(self):
- pool_command.pool_list(self.parser.parse_args(['pools', 'list', '--output', 'tsv']))
+ pool_command.pool_list(self.parser.parse_args(['pools', 'list', '--output', 'json']))
def test_pool_create(self):
pool_command.pool_set(self.parser.parse_args(['pools', 'set', 'foo', '1', 'test']))
diff --git a/tests/cli/commands/test_role_command.py b/tests/cli/commands/test_role_command.py
index e5a93fcebf8ee..167e3a005435e 100644
--- a/tests/cli/commands/test_role_command.py
+++ b/tests/cli/commands/test_role_command.py
@@ -86,4 +86,4 @@ def test_cli_list_roles(self):
self.assertIn('FakeTeamB', stdout)
def test_cli_list_roles_with_args(self):
- role_command.roles_list(self.parser.parse_args(['roles', 'list', '--output', 'tsv']))
+ role_command.roles_list(self.parser.parse_args(['roles', 'list', '--output', 'yaml']))
diff --git a/tests/cli/commands/test_task_command.py b/tests/cli/commands/test_task_command.py
index 61020ec56ffa9..afa16d149b111 100644
--- a/tests/cli/commands/test_task_command.py
+++ b/tests/cli/commands/test_task_command.py
@@ -17,6 +17,7 @@
# under the License.
#
import io
+import json
import logging
import os
import unittest
@@ -26,7 +27,6 @@
import pytest
from parameterized import parameterized
-from tabulate import tabulate
from airflow.cli import cli_parser
from airflow.cli.commands import task_command
@@ -55,6 +55,7 @@ def reset(dag_id):
runs.delete()
+# TODO: Check if tests needs side effects - locally there's missing DAG
class TestCliTasks(unittest.TestCase):
@classmethod
def setUpClass(cls):
@@ -247,29 +248,31 @@ def test_task_states_for_dag_run(self):
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()]
+ [
+ 'tasks',
+ 'states-for-dag-run',
+ 'example_python_operator',
+ defaut_date2.isoformat(),
+ '--output',
+ "json",
+ ]
)
)
- 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="plain"
+ actual_out = json.loads(stdout.getvalue())
+
+ self.assertEqual(len(actual_out), 1)
+ self.assertDictEqual(
+ 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(),
+ },
)
- # Check that prints, and log messages, are shown
- self.assertIn(expected.replace("\n", ""), actual_out.replace("\n", ""))
-
def test_subdag_clear(self):
args = self.parser.parse_args(['tasks', 'clear', 'example_subdag_operator', '--yes'])
task_command.task_clear(args)
diff --git a/tests/cli/commands/test_user_command.py b/tests/cli/commands/test_user_command.py
index 188a977676b3c..f5bf1671802c0 100644
--- a/tests/cli/commands/test_user_command.py
+++ b/tests/cli/commands/test_user_command.py
@@ -161,7 +161,7 @@ def test_cli_list_users(self):
self.assertIn(f'user{i}', stdout)
def test_cli_list_users_with_args(self):
- user_command.users_list(self.parser.parse_args(['users', 'list', '--output', 'tsv']))
+ user_command.users_list(self.parser.parse_args(['users', 'list', '--output', 'json']))
def test_cli_import_users(self):
def assert_user_in_roles(email, roles):