From 96b23d4a107a2628520295dff90b00bde8bd9fbb Mon Sep 17 00:00:00 2001 From: hankehly Date: Thu, 16 Jun 2022 09:23:44 +0900 Subject: [PATCH 1/3] Add kerberos click command --- airflow/cli/__main__.py | 1 + airflow/cli/commands/kerberos.py | 54 ++++++++++++++++++++++++++++++++ 2 files changed, 55 insertions(+) create mode 100644 airflow/cli/commands/kerberos.py diff --git a/airflow/cli/__main__.py b/airflow/cli/__main__.py index a559530b7af3d..bc694e80446b1 100644 --- a/airflow/cli/__main__.py +++ b/airflow/cli/__main__.py @@ -23,6 +23,7 @@ from airflow.cli.commands import db # noqa: F401 from airflow.cli.commands import info # noqa: F401 from airflow.cli.commands import jobs # noqa: F401 +from airflow.cli.commands import kerberos # noqa: F401 from airflow.cli.commands import scheduler # noqa: F401 from airflow.cli.commands import standalone # noqa: F401 from airflow.cli.commands import sync_perm # noqa: F401 diff --git a/airflow/cli/commands/kerberos.py b/airflow/cli/commands/kerberos.py new file mode 100644 index 0000000000000..cd9451cb3bc96 --- /dev/null +++ b/airflow/cli/commands/kerberos.py @@ -0,0 +1,54 @@ +# 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. + +"""Kerberos command""" +import daemon +import rich_click as click +from daemon.pidfile import TimeoutPIDLockFile + +from airflow import settings +from airflow.cli import airflow_cmd, click_daemon, click_log_file, click_pid, click_stderr, click_stdout +from airflow.configuration import conf +from airflow.security import kerberos as krb +from airflow.utils.cli import setup_locations + + +@airflow_cmd.command() +@click.argument("principal") +@click_stdout +@click_stderr +@click_pid +@click_daemon +@click_log_file +@click.option("-k", "--keytab", metavar="KEYTAB", help="keytab", default=conf.get("kerberos", "keytab")) +def kerberos(principal, stdout, stderr, pid, daemon, log_file, keytab): + """Start a kerberos ticket renewer""" + print(settings.HEADER) + + if daemon: + pid, stdout, stderr, _ = setup_locations("kerberos", pid, stdout, stderr, log_file) + with open(stdout, "w+") as stdout_handle, open(stderr, "w+") as stderr_handle: + ctx = daemon.DaemonContext( + pidfile=TimeoutPIDLockFile(pid, -1), + stdout=stdout_handle, + stderr=stderr_handle, + ) + + with ctx: + krb.run(principal=principal, keytab=keytab) + else: + krb.run(principal=principal, keytab=keytab) From c5046713a196c3a5f35f8523d68f19fff5bb8ec1 Mon Sep 17 00:00:00 2001 From: hankehly Date: Thu, 16 Jun 2022 10:22:58 +0900 Subject: [PATCH 2/3] Update unit tests --- airflow/cli/commands/kerberos.py | 4 +- tests/cli/commands/test_kerberos.py | 99 +++++++++++++++++++++++++++++ 2 files changed, 101 insertions(+), 2 deletions(-) create mode 100644 tests/cli/commands/test_kerberos.py diff --git a/airflow/cli/commands/kerberos.py b/airflow/cli/commands/kerberos.py index cd9451cb3bc96..84f8dfc2d1325 100644 --- a/airflow/cli/commands/kerberos.py +++ b/airflow/cli/commands/kerberos.py @@ -35,11 +35,11 @@ @click_daemon @click_log_file @click.option("-k", "--keytab", metavar="KEYTAB", help="keytab", default=conf.get("kerberos", "keytab")) -def kerberos(principal, stdout, stderr, pid, daemon, log_file, keytab): +def kerberos(principal, stdout, stderr, pid, daemon_, log_file, keytab): """Start a kerberos ticket renewer""" print(settings.HEADER) - if daemon: + if daemon_: pid, stdout, stderr, _ = setup_locations("kerberos", pid, stdout, stderr, log_file) with open(stdout, "w+") as stdout_handle, open(stderr, "w+") as stderr_handle: ctx = daemon.DaemonContext( diff --git a/tests/cli/commands/test_kerberos.py b/tests/cli/commands/test_kerberos.py new file mode 100644 index 0000000000000..f0dab11568171 --- /dev/null +++ b/tests/cli/commands/test_kerberos.py @@ -0,0 +1,99 @@ +# 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. + +import unittest +from unittest import mock + +from click.testing import CliRunner + +from airflow.cli.commands.kerberos import kerberos +from tests.test_utils.config import conf_vars + + +class TestKerberosCommand(unittest.TestCase): + @mock.patch("airflow.cli.commands.kerberos.krb") + @conf_vars({("core", "executor"): "CeleryExecutor"}) + def test_run_command(self, mock_krb): + runner = CliRunner() + runner.invoke(kerberos, ["PRINCIPAL", "--keytab", "/tmp/airflow.keytab"]) + mock_krb.run.assert_called_once_with(keytab="/tmp/airflow.keytab", principal="PRINCIPAL") + + @mock.patch("airflow.cli.commands.kerberos.TimeoutPIDLockFile") + @mock.patch("airflow.cli.commands.kerberos.setup_locations") + @mock.patch("airflow.cli.commands.kerberos.daemon") + @mock.patch("airflow.cli.commands.kerberos.krb") + @conf_vars({("core", "executor"): "CeleryExecutor"}) + def test_run_command_daemon(self, mock_krb, mock_daemon, mock_setup_locations, mock_pid_file): + mock_setup_locations.return_value = ( + mock.MagicMock(name="pidfile"), + mock.MagicMock(name="stdout"), + mock.MagicMock(name="stderr"), + mock.MagicMock(name="INVALID"), + ) + + mock_open = mock.mock_open() + with mock.patch("airflow.cli.commands.kerberos.open", mock_open): + runner = CliRunner() + runner.invoke( + kerberos, + [ + "PRINCIPAL", + "--keytab", + "/tmp/airflow.keytab", + "--log-file", + "/tmp/kerberos.log", + "--pid", + "/tmp/kerberos.pid", + "--stderr", + "/tmp/kerberos-stderr.log", + "--stdout", + "/tmp/kerberos-stdout.log", + "--daemon", + ], + ) + + mock_krb.run.assert_called_once_with(keytab="/tmp/airflow.keytab", principal="PRINCIPAL") + assert mock_daemon.mock_calls == [ + mock.call.DaemonContext( + pidfile=mock_pid_file.return_value, + stderr=mock_open.return_value, + stdout=mock_open.return_value, + ), + mock.call.DaemonContext().__enter__(), + mock.call.DaemonContext().__exit__(None, None, None), + ] + + mock_setup_locations.assert_has_calls( + [ + mock.call( + "kerberos", + "/tmp/kerberos.pid", + "/tmp/kerberos-stdout.log", + "/tmp/kerberos-stderr.log", + "/tmp/kerberos.log", + ) + ] + ) + mock_pid_file.assert_has_calls([mock.call(mock_setup_locations.return_value[0], -1)]) + assert mock_open.mock_calls == [ + mock.call(mock_setup_locations.return_value[1], "w+"), + mock.call().__enter__(), + mock.call(mock_setup_locations.return_value[2], "w+"), + mock.call().__enter__(), + mock.call().__exit__(None, None, None), + mock.call().__exit__(None, None, None), + ] From 5c0691382eb443db4775bcc10559deeb4e12c3c6 Mon Sep 17 00:00:00 2001 From: hankehly Date: Fri, 17 Jun 2022 08:52:47 +0900 Subject: [PATCH 3/3] Replace print with rich console print --- airflow/cli/commands/kerberos.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/airflow/cli/commands/kerberos.py b/airflow/cli/commands/kerberos.py index 84f8dfc2d1325..69ae2970a5325 100644 --- a/airflow/cli/commands/kerberos.py +++ b/airflow/cli/commands/kerberos.py @@ -19,6 +19,7 @@ import daemon import rich_click as click from daemon.pidfile import TimeoutPIDLockFile +from rich.console import Console from airflow import settings from airflow.cli import airflow_cmd, click_daemon, click_log_file, click_pid, click_stderr, click_stdout @@ -37,7 +38,8 @@ @click.option("-k", "--keytab", metavar="KEYTAB", help="keytab", default=conf.get("kerberos", "keytab")) def kerberos(principal, stdout, stderr, pid, daemon_, log_file, keytab): """Start a kerberos ticket renewer""" - print(settings.HEADER) + console = Console() + console.print(settings.HEADER) if daemon_: pid, stdout, stderr, _ = setup_locations("kerberos", pid, stdout, stderr, log_file)