Add click-compatible kerberos command

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

airflow/cli/commands/kerberos.py

@@ -0,0 +1,56 @@
+# 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 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
+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"""
+    console = Console()
+    console.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)

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),
+        ]