From 904f415f89e832d011c3d6d2effaa40e16caa79d Mon Sep 17 00:00:00 2001 From: Tomasz Urbaszek Date: Sun, 29 Nov 2020 16:14:22 +0100 Subject: [PATCH 1/9] Refactor list rendering in commands This commit unifies the mechanism of rendering output of tabular data. This gives users a possibility to eiter display a tabular representation of data or render it as valid json or yaml payload. Closes: #12699 --- UPDATING.md | 33 +++++++ airflow/cli/cli_parser.py | 17 ++-- airflow/cli/commands/connection_command.py | 79 ++++++----------- airflow/cli/commands/dag_command.py | 86 +++++++++---------- airflow/cli/commands/pool_command.py | 27 +++--- airflow/cli/commands/provider_command.py | 79 +++++++---------- airflow/cli/commands/role_command.py | 9 +- airflow/cli/commands/task_command.py | 61 ++++++------- airflow/cli/commands/user_command.py | 10 +-- airflow/cli/commands/variable_command.py | 3 +- airflow/cli/simple_table.py | 65 ++++++++++++++ airflow/utils/cli.py | 1 + tests/cli/commands/test_connection_command.py | 26 +++--- tests/cli/commands/test_dag_command.py | 10 ++- tests/cli/commands/test_pool_command.py | 5 +- tests/cli/commands/test_role_command.py | 2 +- tests/cli/commands/test_task_command.py | 1 + tests/cli/commands/test_user_command.py | 2 +- 18 files changed, 285 insertions(+), 231 deletions(-) diff --git a/UPDATING.md b/UPDATING.md index e9b36bb66ae52..504884df1973f 100644 --- a/UPDATING.md +++ b/UPDATING.md @@ -52,6 +52,39 @@ assists users migrating to a new version. ## Master +### Changes to output argument in commands + +Instead of using [tabulate](https://pypi.org/project/tabulate/) to render commands output +we use [rich](https://github.com/willmcgugan/rich). Due to this change the `--output` argument +will no longer accept formats of tabulate tables. Instead it 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 export` +- `airflow role list` +- `airflow providers list` +- `airflow providers get` +- `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/connection_command.py b/airflow/cli/commands/connection_command.py index 279a7e01bee9d..05980f24f22d2 100644 --- a/airflow/cli/commands/connection_command.py +++ b/airflow/cli/commands/connection_command.py @@ -19,60 +19,36 @@ 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.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) def connections_get(args): @@ -81,14 +57,11 @@ def connections_get(args): 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, + ) def connections_list(args): @@ -99,9 +72,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..8ff1fd5248757 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 @@ -42,34 +42,6 @@ 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""" @@ -296,15 +268,32 @@ def dag_next_execution(args): 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 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 @@ -324,6 +313,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,12 +322,12 @@ 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 @@ -361,13 +351,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/pool_command.py b/airflow/cli/commands/pool_command.py index 641b97723163e..d04ea5ed4e16e 100644 --- a/airflow/cli/commands/pool_command.py +++ b/airflow/cli/commands/pool_command.py @@ -21,28 +21,35 @@ 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.utils import cli as cli_utils -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], + }, + ) 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) 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 @@ -50,7 +57,7 @@ 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)) + _show_pools(pools=pools, output=args.output) @cli_utils.action_logging @@ -58,7 +65,7 @@ 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)) + _show_pools(pools=pools, output=args.output) @cli_utils.action_logging @@ -67,7 +74,7 @@ def pool_import(args): 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)) + _show_pools(pools=pools, output=args.output) if len(failed) > 0: sys.exit("Failed to update pool(s): {}".format(", ".join(failed))) @@ -75,7 +82,7 @@ def pool_import(args): 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)) + _show_pools(pools=pools, output=args.output) def pool_import_helper(filepath): diff --git a/airflow/cli/commands/provider_command.py b/airflow/cli/commands/provider_command.py index a2df92fa0c6e7..bae022511d0a1 100644 --- a/airflow/cli/commands/provider_command.py +++ b/airflow/cli/commands/provider_command.py @@ -15,71 +15,56 @@ # 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 - -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.")) 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}") 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], + }, + ) 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..1f535c762b45d 100644 --- a/airflow/cli/commands/role_command.py +++ b/airflow/cli/commands/role_command.py @@ -17,8 +17,8 @@ # 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.www.app import cached_app @@ -27,10 +27,9 @@ 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..12078a8f23e58 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 @@ -306,41 +305,35 @@ def _guess_debugger(): @cli_utils.action_logging 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(), + }, ) - ) - - session.close() @cli_utils.action_logging diff --git a/airflow/cli/commands/user_command.py b/airflow/cli/commands/user_command.py index 5481c7589ce5c..4fa10e6c730ff 100644 --- a/airflow/cli/commands/user_command.py +++ b/airflow/cli/commands/user_command.py @@ -24,8 +24,7 @@ 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.www.app import cached_app @@ -35,9 +34,10 @@ def users_list(args): 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..be55fe22b0f40 100644 --- a/airflow/cli/commands/variable_command.py +++ b/airflow/cli/commands/variable_command.py @@ -21,6 +21,7 @@ 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.session import create_session @@ -30,7 +31,7 @@ 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..edc99e3739c45 100644 --- a/airflow/cli/simple_table.py +++ b/airflow/cli/simple_table.py @@ -14,11 +14,76 @@ # 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_rendered = { + "json": self.print_as_json, + "yaml": self.print_as_yaml, + "table": self.print_as_table, + } + renderer = output_to_rendered.get(output) + if not renderer: + raise ValueError("The output") + + if not all(isinstance(d, dict) for d in data) and not mapper: + raise ValueError("To tabulate non-dictionary data you need to provider `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..32cc6ee53f749 100644 --- a/airflow/utils/cli.py +++ b/airflow/utils/cli.py @@ -260,6 +260,7 @@ def sigquit_handler(sig, frame): # pylint: disable=unused-argument class ColorMode: """Coloring modes. If `auto` is then automatically detected.""" + ON = "on" ON = "on" OFF = "off" AUTO = "auto" diff --git a/tests/cli/commands/test_connection_command.py b/tests/cli/commands/test_connection_command.py index adc4b423394c6..ba9ccd504cf60 100644 --- a/tests/cli/commands/test_connection_command.py +++ b/tests/cli/commands/test_connection_command.py @@ -46,7 +46,7 @@ def test_cli_connection_get(self): self.parser.parse_args(["connections", "get", "google_cloud_default"]) ) 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.")): @@ -116,32 +116,34 @@ 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"]) + # Tables sometimes wrap the content so full name may not be present + if len(conn_id) < 20 and len(conn_type) < 20: + self.assertIn(conn_type, stdout) + self.assertIn(conn_id, stdout) + 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..64f166fa50a02 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): @@ -332,18 +333,19 @@ def test_cli_report(self): with contextlib.redirect_stdout(io.StringIO()) as temp_stdout: dag_command.dag_report(args) out = temp_stdout.getvalue() + print(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']) 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 +385,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..5e77e712d7eed 100644 --- a/tests/cli/commands/test_task_command.py +++ b/tests/cli/commands/test_task_command.py @@ -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): 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): From 2d5aebab664c34e0a64388ba9821ff6a9346c8f1 Mon Sep 17 00:00:00 2001 From: Tomasz Urbaszek Date: Sun, 29 Nov 2020 18:51:46 +0100 Subject: [PATCH 2/9] fixup! Refactor list rendering in commands --- airflow/cli/commands/cheat_sheet_command.py | 2 + airflow/cli/commands/connection_command.py | 3 ++ airflow/cli/commands/dag_command.py | 12 +++++- airflow/cli/commands/info_command.py | 2 + airflow/cli/commands/plugins_command.py | 3 ++ airflow/cli/commands/pool_command.py | 6 +++ airflow/cli/commands/provider_command.py | 4 ++ airflow/cli/commands/role_command.py | 2 + airflow/cli/commands/task_command.py | 9 +++- airflow/cli/commands/user_command.py | 2 + airflow/cli/commands/variable_command.py | 2 + airflow/utils/cli.py | 16 ++++++- .../run_install_and_test_provider_packages.sh | 4 +- tests/cli/commands/test_connection_command.py | 11 ----- tests/cli/commands/test_dag_command.py | 5 +-- tests/cli/commands/test_task_command.py | 42 ++++++++++--------- 16 files changed, 85 insertions(+), 40 deletions(-) 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 05980f24f22d2..11020b128594a 100644 --- a/airflow/cli/commands/connection_command.py +++ b/airflow/cli/commands/connection_command.py @@ -30,6 +30,7 @@ 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 suppress_logs_and_warning from airflow.utils.session import create_session @@ -51,6 +52,7 @@ def _connection_mapper(conn: Connection) -> Dict[str, Any]: } +@suppress_logs_and_warning() def connections_get(args): """Get a connection.""" try: @@ -64,6 +66,7 @@ def connections_get(args): ) +@suppress_logs_and_warning() def connections_list(args): """Lists all connections at the command line""" with create_session() as session: diff --git a/airflow/cli/commands/dag_command.py b/airflow/cli/commands/dag_command.py index 8ff1fd5248757..f7b57bc895878 100644 --- a/airflow/cli/commands/dag_command.py +++ b/airflow/cli/commands/dag_command.py @@ -36,7 +36,13 @@ 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 @@ -265,6 +271,7 @@ 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)) @@ -280,6 +287,7 @@ def dag_list_dags(args): @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)) @@ -297,6 +305,7 @@ def dag_report(args): @cli_utils.action_logging +@suppress_logs_and_warning() def dag_list_jobs(args, dag=None): """Lists latest n jobs""" queries = [] @@ -331,6 +340,7 @@ def dag_list_jobs(args, dag=None): @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: 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 d04ea5ed4e16e..bca617a749571 100644 --- a/airflow/cli/commands/pool_command.py +++ b/airflow/cli/commands/pool_command.py @@ -24,6 +24,7 @@ from airflow.api.client import get_current_api_client from airflow.cli.simple_table import AirflowConsole from airflow.utils import cli as cli_utils +from airflow.utils.cli import suppress_logs_and_warning def _show_pools(pools, output): @@ -38,6 +39,7 @@ def _show_pools(pools, output): ) +@suppress_logs_and_warning() def pool_list(args): """Displays info of all the pools""" api_client = get_current_api_client() @@ -45,6 +47,7 @@ def pool_list(args): _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() @@ -53,6 +56,7 @@ def pool_get(args): @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() @@ -61,6 +65,7 @@ def pool_set(args): @cli_utils.action_logging +@suppress_logs_and_warning() def pool_delete(args): """Deletes pool by a given name""" api_client = get_current_api_client() @@ -69,6 +74,7 @@ def pool_delete(args): @cli_utils.action_logging +@suppress_logs_and_warning() def pool_import(args): """Imports pools from the file""" if not os.path.exists(args.file): diff --git a/airflow/cli/commands/provider_command.py b/airflow/cli/commands/provider_command.py index bae022511d0a1..c7830789c60d3 100644 --- a/airflow/cli/commands/provider_command.py +++ b/airflow/cli/commands/provider_command.py @@ -19,12 +19,14 @@ from airflow.cli.simple_table import AirflowConsole from airflow.providers_manager import ProvidersManager +from airflow.utils.cli import suppress_logs_and_warning 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 @@ -44,6 +46,7 @@ def provider_get(args): 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""" AirflowConsole().print_as( @@ -57,6 +60,7 @@ def providers_list(args): ) +@suppress_logs_and_warning() def hooks_list(args): """Lists all hooks at the command line""" AirflowConsole().print_as( diff --git a/airflow/cli/commands/role_command.py b/airflow/cli/commands/role_command.py index 1f535c762b45d..a76061098ed81 100644 --- a/airflow/cli/commands/role_command.py +++ b/airflow/cli/commands/role_command.py @@ -20,9 +20,11 @@ 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 diff --git a/airflow/cli/commands/task_command.py b/airflow/cli/commands/task_command.py index 12078a8f23e58..466788129d7f3 100644 --- a/airflow/cli/commands/task_command.py +++ b/airflow/cli/commands/task_command.py @@ -35,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 @@ -303,6 +309,7 @@ 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""" with create_session() as session: diff --git a/airflow/cli/commands/user_command.py b/airflow/cli/commands/user_command.py index 4fa10e6c730ff..a5f8ec435c482 100644 --- a/airflow/cli/commands/user_command.py +++ b/airflow/cli/commands/user_command.py @@ -26,9 +26,11 @@ 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 diff --git a/airflow/cli/commands/variable_command.py b/airflow/cli/commands/variable_command.py index be55fe22b0f40..b3f34ed9b47ec 100644 --- a/airflow/cli/commands/variable_command.py +++ b/airflow/cli/commands/variable_command.py @@ -24,9 +24,11 @@ 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: diff --git a/airflow/utils/cli.py b/airflow/utils/cli.py index 32cc6ee53f749..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 @@ -260,7 +261,6 @@ def sigquit_handler(sig, frame): # pylint: disable=unused-argument class ColorMode: """Coloring modes. If `auto` is then automatically detected.""" - ON = "on" ON = "on" OFF = "off" AUTO = "auto" @@ -273,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/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 ba9ccd504cf60..afda4f913f189 100644 --- a/tests/cli/commands/test_connection_command.py +++ b/tests/cli/commands/test_connection_command.py @@ -112,17 +112,6 @@ 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() - - for conn_id, conn_type in self.EXPECTED_CONS: - # Tables sometimes wrap the content so full name may not be present - if len(conn_id) < 20 and len(conn_type) < 20: - self.assertIn(conn_type, stdout) - self.assertIn(conn_id, stdout) - def test_cli_connections_list_as_json(self): args = self.parser.parse_args(["connections", "list", "--output", "json"]) with redirect_stdout(io.StringIO()) as stdout: diff --git a/tests/cli/commands/test_dag_command.py b/tests/cli/commands/test_dag_command.py index 64f166fa50a02..5358fdf295b2d 100644 --- a/tests/cli/commands/test_dag_command.py +++ b/tests/cli/commands/test_dag_command.py @@ -333,14 +333,13 @@ def test_cli_report(self): with contextlib.redirect_stdout(io.StringIO()) as temp_stdout: dag_command.dag_report(args) out = temp_stdout.getvalue() - print(out) self.assertIn("airflow/example_dags/example_complex.py ", out) - self.assertIn("['example_complex']", 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']) + 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() diff --git a/tests/cli/commands/test_task_command.py b/tests/cli/commands/test_task_command.py index 5e77e712d7eed..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 @@ -248,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) From 58dc49b07c98c47e69dd0828eb52b6ca60cc22fd Mon Sep 17 00:00:00 2001 From: Tomasz Urbaszek Date: Sun, 29 Nov 2020 19:20:03 +0100 Subject: [PATCH 3/9] Add note in docs --- docs/apache-airflow/usage-cli.rst | 37 +++++++++++++++++++++++++++++++ 1 file changed, 37 insertions(+) diff --git a/docs/apache-airflow/usage-cli.rst b/docs/apache-airflow/usage-cli.rst index a41fcf3fbf1aa..cb96232a2b443 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 ``dags list`` or ``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" + } From e546e5226125897bc3add8bd76bcc2f8d0632c26 Mon Sep 17 00:00:00 2001 From: Tomasz Urbaszek Date: Sun, 29 Nov 2020 21:34:05 +0100 Subject: [PATCH 4/9] fixup! Add note in docs --- tests/cli/commands/test_connection_command.py | 2 +- tests/cli/commands/test_dag_command.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/cli/commands/test_connection_command.py b/tests/cli/commands/test_connection_command.py index afda4f913f189..700d8191e1702 100644 --- a/tests/cli/commands/test_connection_command.py +++ b/tests/cli/commands/test_connection_command.py @@ -43,7 +43,7 @@ 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("google-cloud-platform:///default", stdout) diff --git a/tests/cli/commands/test_dag_command.py b/tests/cli/commands/test_dag_command.py index 5358fdf295b2d..aec2642bb7910 100644 --- a/tests/cli/commands/test_dag_command.py +++ b/tests/cli/commands/test_dag_command.py @@ -329,7 +329,7 @@ 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() From a408961d9c5c3869bd3a1f76636dbdcd93ebcf58 Mon Sep 17 00:00:00 2001 From: Tomasz Urbaszek Date: Sun, 29 Nov 2020 21:39:59 +0100 Subject: [PATCH 5/9] fixup! fixup! Add note in docs --- airflow/cli/simple_table.py | 8 +++++--- docs/apache-airflow/usage-cli.rst | 2 +- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/airflow/cli/simple_table.py b/airflow/cli/simple_table.py index edc99e3739c45..67663d94c8559 100644 --- a/airflow/cli/simple_table.py +++ b/airflow/cli/simple_table.py @@ -64,14 +64,16 @@ def _normalize_data(self, value: Any, output: str) -> Union[list, str, dict]: 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_rendered = { + output_to_renderer = { "json": self.print_as_json, "yaml": self.print_as_yaml, "table": self.print_as_table, } - renderer = output_to_rendered.get(output) + renderer = output_to_renderer.get(output) if not renderer: - raise ValueError("The output") + 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 provider `mapper` function") diff --git a/docs/apache-airflow/usage-cli.rst b/docs/apache-airflow/usage-cli.rst index cb96232a2b443..48b41ace8bf48 100644 --- a/docs/apache-airflow/usage-cli.rst +++ b/docs/apache-airflow/usage-cli.rst @@ -190,7 +190,7 @@ Both ``json`` and ``yaml`` formats make it easier to manipulate the data using c .. 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}" + 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" From 573800228eceba85a81b041a525aabfd638f3963 Mon Sep 17 00:00:00 2001 From: Tomasz Urbaszek Date: Sun, 29 Nov 2020 22:26:32 +0100 Subject: [PATCH 6/9] fixup! fixup! fixup! Add note in docs --- tests/cli/commands/test_dag_command.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/cli/commands/test_dag_command.py b/tests/cli/commands/test_dag_command.py index aec2642bb7910..55f4226d51af5 100644 --- a/tests/cli/commands/test_dag_command.py +++ b/tests/cli/commands/test_dag_command.py @@ -334,7 +334,7 @@ def test_cli_report(self): dag_command.dag_report(args) out = temp_stdout.getvalue() - self.assertIn("airflow/example_dags/example_complex.py ", out) + self.assertIn("airflow/example_dags/example_complex.py", out) self.assertIn("example_complex", out) @conf_vars({('core', 'load_examples'): 'true'}) From 03e4867b2eb4b928eccf96f678febf5c0540280a Mon Sep 17 00:00:00 2001 From: Tomasz Urbaszek Date: Mon, 30 Nov 2020 20:57:33 +0100 Subject: [PATCH 7/9] fixup! fixup! fixup! fixup! Add note in docs --- UPDATING.md | 2 ++ airflow/cli/commands/connection_command.py | 2 +- airflow/cli/simple_table.py | 2 +- docs/apache-airflow/usage-cli.rst | 4 ++-- 4 files changed, 6 insertions(+), 4 deletions(-) diff --git a/UPDATING.md b/UPDATING.md index 504884df1973f..362f0d90ca0ab 100644 --- a/UPDATING.md +++ b/UPDATING.md @@ -77,10 +77,12 @@ Affected commands: - `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` diff --git a/airflow/cli/commands/connection_command.py b/airflow/cli/commands/connection_command.py index 11020b128594a..dd96f9f43e43a 100644 --- a/airflow/cli/commands/connection_command.py +++ b/airflow/cli/commands/connection_command.py @@ -48,7 +48,7 @@ def _connection_mapper(conn: Connection) -> Dict[str, Any]: 'is_encrypted': conn.is_encrypted, 'is_extra_encrypted': conn.is_encrypted, 'extra_dejson': conn.extra_dejson, - 'get_uri()': conn.get_uri(), + 'get_uri': conn.get_uri(), } diff --git a/airflow/cli/simple_table.py b/airflow/cli/simple_table.py index 67663d94c8559..8cdb6de61083c 100644 --- a/airflow/cli/simple_table.py +++ b/airflow/cli/simple_table.py @@ -76,7 +76,7 @@ def print_as(self, data: List[Union[Dict, Any]], output: str, mapper: Optional[C ) if not all(isinstance(d, dict) for d in data) and not mapper: - raise ValueError("To tabulate non-dictionary data you need to provider `mapper` function") + 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] diff --git a/docs/apache-airflow/usage-cli.rst b/docs/apache-airflow/usage-cli.rst index 48b41ace8bf48..5dadab8466952 100644 --- a/docs/apache-airflow/usage-cli.rst +++ b/docs/apache-airflow/usage-cli.rst @@ -178,8 +178,8 @@ You will see a similar result as in the screenshot below. Formatting commands output -------------------------- -Some Airflow commands like ``dags list`` or ``tasks states-for-dag-run`` support ``--output`` flag which allow users -to change the formatting of command's output. Possible options: +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 From d12f31aa04d090530011fa8d24ea51faf74b1705 Mon Sep 17 00:00:00 2001 From: Tomasz Urbaszek Date: Mon, 30 Nov 2020 21:40:01 +0100 Subject: [PATCH 8/9] fixup! fixup! fixup! fixup! fixup! Add note in docs --- airflow/cli/commands/pool_command.py | 20 +++++++++++--------- airflow/cli/commands/task_command.py | 4 ++-- 2 files changed, 13 insertions(+), 11 deletions(-) diff --git a/airflow/cli/commands/pool_command.py b/airflow/cli/commands/pool_command.py index bca617a749571..20b7959d8f74d 100644 --- a/airflow/cli/commands/pool_command.py +++ b/airflow/cli/commands/pool_command.py @@ -23,6 +23,7 @@ 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 @@ -60,8 +61,8 @@ def pool_get(args): 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)] - _show_pools(pools=pools, output=args.output) + api_client.create_pool(name=args.pool, slots=args.slots, description=args.description) + print("Pool created") @cli_utils.action_logging @@ -69,8 +70,11 @@ def pool_set(args): def pool_delete(args): """Deletes pool by a given name""" api_client = get_current_api_client() - pools = [api_client.delete_pool(name=args.pool)] - _show_pools(pools=pools, output=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 @@ -79,16 +83,15 @@ 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) - _show_pools(pools=pools, output=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) - _show_pools(pools=pools, output=args.output) + print(f"Exported {len(pools)} pools to {args.file}") def pool_import_helper(filepath): @@ -121,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/task_command.py b/airflow/cli/commands/task_command.py index 466788129d7f3..356f60cfe5950 100644 --- a/airflow/cli/commands/task_command.py +++ b/airflow/cli/commands/task_command.py @@ -337,8 +337,8 @@ def task_states_for_dag_run(args): "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(), + "start_date": ti.start_date.isoformat() if ti.start_date else "", + "end_date": ti.end_date.isoformat() if ti.end_date else "", }, ) From 261ab74ea85dbc58cfbb1747f04be35945a78802 Mon Sep 17 00:00:00 2001 From: Tomek Urbaszek Date: Tue, 1 Dec 2020 23:25:29 +0000 Subject: [PATCH 9/9] Update UPDATING.md Co-authored-by: Kaxil Naik --- UPDATING.md | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/UPDATING.md b/UPDATING.md index 362f0d90ca0ab..6af5e09be8b8b 100644 --- a/UPDATING.md +++ b/UPDATING.md @@ -54,9 +54,8 @@ assists users migrating to a new version. ### Changes to output argument in commands -Instead of using [tabulate](https://pypi.org/project/tabulate/) to render commands output -we use [rich](https://github.com/willmcgugan/rich). Due to this change the `--output` argument -will no longer accept formats of tabulate tables. Instead it accepts: +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