diff --git a/airflow-core/src/airflow/cli/commands/jobs_command.py b/airflow-core/src/airflow/cli/commands/jobs_command.py index 194d8720db2f1..57ab4a1df7b61 100644 --- a/airflow-core/src/airflow/cli/commands/jobs_command.py +++ b/airflow-core/src/airflow/cli/commands/jobs_command.py @@ -20,6 +20,7 @@ from sqlalchemy import select +from airflow.cli.utils import deprecated_for_airflowctl from airflow.jobs.job import Job, JobState from airflow.utils.net import get_hostname from airflow.utils.providers_configuration_loader import providers_configuration_loaded @@ -29,6 +30,7 @@ from sqlalchemy.orm import Session +@deprecated_for_airflowctl("airflowctl jobs check") @providers_configuration_loaded @provide_session def check(args, *, session: Session = NEW_SESSION) -> None: diff --git a/airflow-core/tests/unit/cli/commands/test_command_deprecations.py b/airflow-core/tests/unit/cli/commands/test_command_deprecations.py index 5e9d0613acd07..f8ba4659a2f24 100644 --- a/airflow-core/tests/unit/cli/commands/test_command_deprecations.py +++ b/airflow-core/tests/unit/cli/commands/test_command_deprecations.py @@ -35,6 +35,7 @@ config_command, connection_command, dag_command, + jobs_command, pool_command, provider_command, task_command, @@ -79,6 +80,7 @@ (config_command.show_config, "airflowctl config list"), (task_command.task_states_for_dag_run, "airflowctl tasks states-for-dag-run"), (task_command.task_clear, "airflowctl tasks clear"), + (jobs_command.check, "airflowctl jobs check"), ] diff --git a/airflow-ctl/docs/images/command_hashes.txt b/airflow-ctl/docs/images/command_hashes.txt index 74b5ede7ce3a2..b7c29753748d2 100644 --- a/airflow-ctl/docs/images/command_hashes.txt +++ b/airflow-ctl/docs/images/command_hashes.txt @@ -6,7 +6,7 @@ config:a3d936cb15fe3b547bf6c82cf93d923f connections:942f9f88cb908c28bf5c19159fc5065b dags:5ad68174a1111563dff870a241914b30 dagrun:07035226eaaff0a557d3ad9aab4b41b5 -jobs:a5b644c5da8889443bb40ee10b599270 +jobs:d4af478f28dae48ee18d43b998edc345 pools:19efe105b9515ab1926ebcaf0e028d71 providers:34502fe09dc0b8b0a13e7e46efdffda6 taskinstances:7e323968c0b585287c2a4ab4339aee1f diff --git a/airflow-ctl/docs/images/output_jobs.svg b/airflow-ctl/docs/images/output_jobs.svg index 13b31d2caedc4..9ff6f8753c162 100644 --- a/airflow-ctl/docs/images/output_jobs.svg +++ b/airflow-ctl/docs/images/output_jobs.svg @@ -1,4 +1,4 @@ - + - - + + - + - + - + - + - + - + - + - + - + + + + - + - + - - Usage:airflowctl jobs [-h] COMMAND... - -Perform Jobs operations - -Positional Arguments: -COMMAND -listList scheduler, triggerer, and other Airflow jobs - -Options: --h, --helpshow this help message and exit + + Usage:airflowctl jobs [-h] COMMAND... + +Perform Jobs operations + +Positional Arguments: +COMMAND +checkCheck if job(s) are still alive. +listList scheduler, triggerer, and other Airflow jobs + +Options: +-h, --helpshow this help message and exit diff --git a/airflow-ctl/src/airflowctl/ctl/cli_config.py b/airflow-ctl/src/airflowctl/ctl/cli_config.py index 899067d341040..936d4356538bb 100755 --- a/airflow-ctl/src/airflowctl/ctl/cli_config.py +++ b/airflow-ctl/src/airflowctl/ctl/cli_config.py @@ -314,6 +314,35 @@ def _load_help_texts_yaml() -> dict[str, dict[str, str]]: choices=("overwrite", "fail", "skip"), ) +# Jobs command args +ARG_JOB_TYPE_FILTER = Arg( + flags=("--job-type",), + choices=("SchedulerJob", "TriggererJob", "DagProcessorJob"), + help="The type of job(s) that will be checked.", +) +ARG_JOB_HOSTNAME_FILTER = Arg( + flags=("--hostname",), + type=str, + default=None, + help="The hostname of job(s) that will be checked.", +) +ARG_JOB_LOCAL_FILTER = Arg( + flags=("--local",), + action="store_true", + help="If passed, this command will only show jobs from the local host.", +) +ARG_JOB_LIMIT = Arg( + flags=("--limit",), + type=positive_int(allow_zero=True), + default=1, + help="The number of recent jobs that will be checked. To disable limit, set 0.", +) +ARG_JOB_ALLOW_MULTIPLE = Arg( + flags=("--allow-multiple",), + action="store_true", + help="If passed, this command will be successful even if multiple matching alive jobs are found.", +) + # Config arguments ARG_CONFIG_SECTION = Arg( flags=("--section",), @@ -1111,6 +1140,21 @@ def merge_commands( ), ) +JOB_COMMANDS = ( + ActionCommand( + name="check", + help="Check if job(s) are still alive.", + func=lazy_load_command("airflowctl.ctl.commands.jobs_command.check"), + args=( + ARG_JOB_TYPE_FILTER, + ARG_JOB_HOSTNAME_FILTER, + ARG_JOB_LOCAL_FILTER, + ARG_JOB_LIMIT, + ARG_JOB_ALLOW_MULTIPLE, + ), + ), +) + core_commands: list[CLICommand] = [ GroupCommand( name="auth", @@ -1133,6 +1177,11 @@ def merge_commands( help="Manage Airflow Dags", subcommands=DAG_COMMANDS, ), + GroupCommand( + name="jobs", + help="Manage Airflow jobs", + subcommands=JOB_COMMANDS, + ), GroupCommand( name="pools", help="Manage Airflow pools", diff --git a/airflow-ctl/src/airflowctl/ctl/commands/jobs_command.py b/airflow-ctl/src/airflowctl/ctl/commands/jobs_command.py new file mode 100644 index 0000000000000..67de7154459b8 --- /dev/null +++ b/airflow-ctl/src/airflowctl/ctl/commands/jobs_command.py @@ -0,0 +1,50 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import socket + +import rich + +from airflowctl.api.client import NEW_API_CLIENT, ClientKind, provide_api_client + + +@provide_api_client(kind=ClientKind.CLI) +def check(args, api_client=NEW_API_CLIENT) -> None: + """Check if job(s) are still alive.""" + if args.allow_multiple and args.limit == 1: + raise SystemExit( + "To use option --allow-multiple, you must set the limit to a value greater than 1 " + "or 0 to disable it." + ) + if args.hostname and args.local: + raise SystemExit("You can't use --hostname and --local at the same time") + + hostname = socket.getfqdn() if args.local else args.hostname + alive_jobs = api_client.jobs.list(job_type=args.job_type, hostname=hostname, is_alive=True).jobs + if args.limit > 0: + alive_jobs = alive_jobs[: args.limit] + + count_alive_jobs = len(alive_jobs) + if count_alive_jobs == 0: + raise SystemExit("No alive jobs found.") + if count_alive_jobs > 1 and not args.allow_multiple: + raise SystemExit(f"Found {count_alive_jobs} alive jobs. Expected only one.") + if count_alive_jobs == 1: + rich.print("Found one alive job.") + else: + rich.print(f"Found {count_alive_jobs} alive jobs.") diff --git a/airflow-ctl/tests/airflow_ctl/ctl/commands/test_jobs_command.py b/airflow-ctl/tests/airflow_ctl/ctl/commands/test_jobs_command.py new file mode 100644 index 0000000000000..697faa682c4a7 --- /dev/null +++ b/airflow-ctl/tests/airflow_ctl/ctl/commands/test_jobs_command.py @@ -0,0 +1,115 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from types import SimpleNamespace +from unittest import mock + +import pytest + +from airflowctl.ctl import cli_parser +from airflowctl.ctl.commands import jobs_command + + +def _make_jobs(count: int) -> SimpleNamespace: + """Build a stand-in for the ``JobCollectionResponse`` of ``count`` already-alive jobs.""" + return SimpleNamespace(jobs=[SimpleNamespace(id=i) for i in range(count)]) + + +class TestCliJobsCheck: + parser = cli_parser.get_parser() + + @pytest.mark.parametrize( + ("argv", "alive_count", "expected_output"), + [ + (["jobs", "check", "--job-type", "SchedulerJob"], 1, "Found one alive job."), + (["jobs", "check", "--limit", "100", "--allow-multiple"], 3, "Found 3 alive jobs."), + (["jobs", "check", "--limit", "0", "--allow-multiple"], 3, "Found 3 alive jobs."), + (["jobs", "check", "--limit", "1"], 5, "Found one alive job."), + ], + ) + def test_reports_alive_jobs(self, capsys, argv, alive_count, expected_output): + api_client = mock.MagicMock() + api_client.jobs.list.return_value = _make_jobs(alive_count) + + jobs_command.check(self.parser.parse_args(argv), api_client=api_client) + + assert expected_output in capsys.readouterr().out + + @pytest.mark.parametrize( + ("argv", "local_hostname", "expected_kwargs"), + [ + ( + ["jobs", "check", "--job-type", "SchedulerJob"], + None, + {"job_type": "SchedulerJob", "hostname": None, "is_alive": True}, + ), + ( + ["jobs", "check", "--hostname", "HOSTNAME"], + None, + {"job_type": None, "hostname": "HOSTNAME", "is_alive": True}, + ), + ( + ["jobs", "check", "--local"], + "local-host", + {"job_type": None, "hostname": "local-host", "is_alive": True}, + ), + ], + ) + def test_forwards_filters_to_jobs_list(self, argv, local_hostname, expected_kwargs): + api_client = mock.MagicMock() + api_client.jobs.list.return_value = _make_jobs(1) + + with mock.patch("airflowctl.ctl.commands.jobs_command.socket.getfqdn", return_value=local_hostname): + jobs_command.check(self.parser.parse_args(argv), api_client=api_client) + + api_client.jobs.list.assert_called_once_with(**expected_kwargs) + + @pytest.mark.parametrize( + ("argv", "alive_count", "expected_error"), + [ + (["jobs", "check"], 0, r"No alive jobs found."), + (["jobs", "check", "--limit", "100"], 3, r"Found 3 alive jobs. Expected only one."), + ], + ) + def test_exits_when_alive_count_unexpected(self, argv, alive_count, expected_error): + api_client = mock.MagicMock() + api_client.jobs.list.return_value = _make_jobs(alive_count) + + with pytest.raises(SystemExit, match=expected_error): + jobs_command.check(self.parser.parse_args(argv), api_client=api_client) + + @pytest.mark.parametrize( + ("argv", "expected_error"), + [ + ( + ["jobs", "check", "--allow-multiple"], + r"To use option --allow-multiple, you must set the limit to a value greater than 1 " + r"or 0 to disable it.", + ), + ( + ["jobs", "check", "--hostname", "h", "--local"], + r"You can't use --hostname and --local at the same time", + ), + ], + ) + def test_rejects_invalid_argument_combinations(self, argv, expected_error): + api_client = mock.MagicMock() + + with pytest.raises(SystemExit, match=expected_error): + jobs_command.check(self.parser.parse_args(argv), api_client=api_client) + api_client.jobs.list.assert_not_called()