[airflow] branch main updated: Add testing a connection via Airflow CLI (#29892)

[po...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git


The following commit(s) were added to refs/heads/main by this push:
     new d2e5b097e6 Add testing a connection via Airflow CLI (#29892)
d2e5b097e6 is described below

commit d2e5b097e6251e31fb4c9bb5bf16dc9c77b56f75
Author: Josh Fell <48...@users.noreply.github.com>
AuthorDate: Thu Mar 9 04:26:10 2023 -0500

    Add testing a connection via Airflow CLI (#29892)
    
    Closes: #29875
    
    This PR introduces a new Airflow CLI command -- `airflow connections test <conn_id>` -- to test a single, predefined Airflow connection.
---
 airflow/cli/cli_parser.py                     |  8 ++++++-
 airflow/cli/commands/connection_command.py    | 20 ++++++++++++++++
 docs/apache-airflow/howto/connection.rst      | 31 ++++++++++++++++--------
 tests/cli/commands/test_connection_command.py | 34 +++++++++++++++++++++++++++
 4 files changed, 82 insertions(+), 11 deletions(-)

diff --git a/airflow/cli/cli_parser.py b/airflow/cli/cli_parser.py
index 993be83c0a..66beadfde1 100644
--- a/airflow/cli/cli_parser.py
+++ b/airflow/cli/cli_parser.py
@@ -836,7 +836,7 @@ ARG_ENV_VARS = Arg(
 )
 
 # connections
-ARG_CONN_ID = Arg(("conn_id",), help="Connection id, required to get/add/delete a connection", type=str)
+ARG_CONN_ID = Arg(("conn_id",), help="Connection ID, required to get/add/delete/test a connection", type=str)
 ARG_CONN_ID_FILTER = Arg(
     ("--conn-id",), help="If passed, only items with the specified connection ID will be displayed", type=str
 )
@@ -1709,6 +1709,12 @@ CONNECTIONS_COMMANDS = (
             ARG_VERBOSE,
         ),
     ),
+    ActionCommand(
+        name="test",
+        help="Test a connection",
+        func=lazy_load_command("airflow.cli.commands.connection_command.connections_test"),
+        args=(ARG_CONN_ID, ARG_VERBOSE),
+    ),
 )
 PROVIDERS_COMMANDS = (
     ActionCommand(
diff --git a/airflow/cli/commands/connection_command.py b/airflow/cli/commands/connection_command.py
index 40f81c1c5d..dc8c1792e3 100644
--- a/airflow/cli/commands/connection_command.py
+++ b/airflow/cli/commands/connection_command.py
@@ -325,3 +325,23 @@ def _import_helper(file_path: str, overwrite: bool) -> None:
             session.merge(conn)
             session.commit()
             print(f"Imported connection {conn_id}")
+
+
+@suppress_logs_and_warning
+def connections_test(args) -> None:
+    """Test an Airflow connection."""
+    console = AirflowConsole()
+
+    print(f"Retrieving connection: {args.conn_id!r}")
+    try:
+        conn = BaseHook.get_connection(args.conn_id)
+    except AirflowNotFoundException:
+        console.print("[bold yellow]\nConnection not found.\n")
+        raise SystemExit(1)
+
+    print("\nTesting...")
+    status, message = conn.test_connection()
+    if status is True:
+        console.print("[bold green]\nConnection success!\n")
+    else:
+        console.print(f"[bold][red]\nConnection failed![/bold]\n{message}\n")
diff --git a/docs/apache-airflow/howto/connection.rst b/docs/apache-airflow/howto/connection.rst
index 00b99de8ae..b43094a88b 100644
--- a/docs/apache-airflow/howto/connection.rst
+++ b/docs/apache-airflow/howto/connection.rst
@@ -191,19 +191,30 @@ Passwords cannot be manipulated or read without the key. For information on conf
 Testing Connections
 ^^^^^^^^^^^^^^^^^^^
 
-Airflow Web UI & API allows to test connections. The test connection feature can be used from
-:ref:`create <creating_connection_ui>` or :ref:`edit <editing_connection_ui>` connection page, or through calling
-:doc:`Connections REST API </stable-rest-api-ref/>`.
+Airflow Web UI, REST API, and CLI allow you to test connections. The test connection feature can be used from
+:ref:`create <creating_connection_ui>` or :ref:`edit <editing_connection_ui>` connection page in the UI, through calling
+:doc:`Connections REST API </stable-rest-api-ref/>`, or running the ``airflow connections test`` :ref:`CLI command <cli>`.
 
-To test a connection Airflow calls out the ``test_connection`` method from the associated hook class and reports the
-results of it. It may happen that the connection type does not have any associated hook or the hook doesn't have the
-``test_connection`` method implementation, in either case the error message will throw the proper error message.
+.. warning::
 
-One important point to note is that the connections will be tested from the webserver only, so this feature is
-subject to network egress rules setup for your webserver. Also, if webserver & worker machines have different libs or
-provider packages installed then the test results might differ.
+    This feature won't be available for the connections residing in external secrets backends when using the
+    Airflow UI or REST API.
+
+To test a connection, Airflow calls the ``test_connection`` method from the associated hook class and reports the
+results. It may happen that the connection type does not have any associated hook or the hook doesn't have the
+``test_connection`` method implementation, in either case an error message will be displayed or functionality
+will be disabled (if you are testing in the UI).
+
+.. note::
+
+    When testing in the Airflow UI, the test executes from the webserver so this feature is subject to network
+    egress rules setup for your webserver.
+
+.. note::
+
+    If webserver & worker machines (if testing via the Airflow UI) or machines/pods (if testing via the
+    Airflow CLI) have different libs or provider packages installed, test results *might* differ.
 
-Last caveat is that this feature won't be available for the connections coming out of the secrets backends.
 
 Custom connection types
 ^^^^^^^^^^^^^^^^^^^^^^^
diff --git a/tests/cli/commands/test_connection_command.py b/tests/cli/commands/test_connection_command.py
index b3a3a7342c..e11a83607d 100644
--- a/tests/cli/commands/test_connection_command.py
+++ b/tests/cli/commands/test_connection_command.py
@@ -919,3 +919,37 @@ class TestCliImportConnections:
 
         # The existing connection should have been overwritten
         assert current_conns_as_dicts["new3"] == expected_connections["new3"]
+
+
+class TestCliTestConnections:
+    parser = cli_parser.get_parser()
+
+    def setup_class(self):
+        clear_db_connections()
+
+    @mock.patch("airflow.providers.http.hooks.http.HttpHook.test_connection")
+    def test_cli_connections_test_success(self, mock_test_conn):
+        """Check that successful connection test result is displayed properly."""
+        conn_id = "http_default"
+        mock_test_conn.return_value = True, None
+        with redirect_stdout(io.StringIO()) as stdout:
+            connection_command.connections_test(self.parser.parse_args(["connections", "test", conn_id]))
+
+            assert "Connection success!" in stdout.getvalue()
+
+    @mock.patch("airflow.providers.http.hooks.http.HttpHook.test_connection")
+    def test_cli_connections_test_fail(self, mock_test_conn):
+        """Check that failed connection test result is displayed properly."""
+        conn_id = "http_default"
+        mock_test_conn.return_value = False, "Failed."
+        with redirect_stdout(io.StringIO()) as stdout:
+            connection_command.connections_test(self.parser.parse_args(["connections", "test", conn_id]))
+
+            assert "Connection failed!\nFailed.\n\n" in stdout.getvalue()
+
+    def test_cli_connections_test_missing_conn(self):
+        """Check a connection test on a non-existent connection raises a "Connection not found" message."""
+        with redirect_stdout(io.StringIO()) as stdout, pytest.raises(SystemExit):
+            connection_command.connections_test(self.parser.parse_args(["connections", "test", "missing"]))
+
+            assert "Connection not found.\n\n" in stdout.getvalue()