[airflow] branch main updated: formatting fix (#22688)

[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 d7993dca2f formatting fix (#22688)
d7993dca2f is described below

commit d7993dca2f182c1d0f281f06ac04b47935016cf1
Author: Bowrna <ma...@gmail.com>
AuthorDate: Thu Apr 14 01:15:11 2022 +0530

    formatting fix (#22688)
---
 airflow/cli/commands/connection_command.py    | 25 +++++++++++++++++++++++++
 tests/cli/commands/test_connection_command.py |  9 +++++++++
 2 files changed, 34 insertions(+)

diff --git a/airflow/cli/commands/connection_command.py b/airflow/cli/commands/connection_command.py
index c07cea72f4..9fbd524ff6 100644
--- a/airflow/cli/commands/connection_command.py
+++ b/airflow/cli/commands/connection_command.py
@@ -27,9 +27,11 @@ from urllib.parse import urlparse, urlunparse
 from sqlalchemy.orm import exc
 
 from airflow.cli.simple_table import AirflowConsole
+from airflow.compat.functools import cache
 from airflow.exceptions import AirflowNotFoundException
 from airflow.hooks.base import BaseHook
 from airflow.models import Connection
+from airflow.providers_manager import ProvidersManager
 from airflow.secrets.local_filesystem import load_connections_dict
 from airflow.utils import cli as cli_utils, yaml
 from airflow.utils.cli import suppress_logs_and_warning
@@ -133,6 +135,21 @@ def _valid_uri(uri: str) -> bool:
     return uri_parts.scheme != '' and uri_parts.netloc != ''
 
 
+@cache
+def _get_connection_types():
+    """Returns connection types available."""
+    _connection_types = ['fs', 'mesos_framework-id', 'email', 'generic']
+    providers_manager = ProvidersManager()
+    for connection_type, provider_info in providers_manager.hooks.items():
+        if provider_info:
+            _connection_types.append(connection_type)
+    return _connection_types
+
+
+def _valid_conn_type(conn_type: str) -> bool:
+    return conn_type in _get_connection_types()
+
+
 def connections_export(args):
     """Exports all connections to a file"""
     file_formats = ['.yaml', '.json', '.env']
@@ -195,6 +212,14 @@ def connections_add(args):
     if has_json and has_uri:
         raise SystemExit('Cannot supply both conn-uri and conn-json')
 
+    if has_type and not (args.conn_type in _get_connection_types()):
+        warnings.warn(f'The type provided to --conn-type is invalid: {args.conn_type}')
+        warnings.warn(
+            f'Supported --conn-types are:{_get_connection_types()}.'
+            'Hence overriding the conn-type with generic'
+        )
+        args.conn_type = 'generic'
+
     if has_uri or has_json:
         invalid_args = []
         if has_uri and not _valid_uri(args.conn_uri):
diff --git a/tests/cli/commands/test_connection_command.py b/tests/cli/commands/test_connection_command.py
index a4e1148e83..621be5916a 100644
--- a/tests/cli/commands/test_connection_command.py
+++ b/tests/cli/commands/test_connection_command.py
@@ -18,6 +18,7 @@
 import io
 import json
 import re
+import warnings
 from contextlib import redirect_stdout
 from unittest import mock
 
@@ -577,6 +578,14 @@ class TestCliAddConnections:
                 self.parser.parse_args(["connections", "add", "new1", f"--conn-uri={'nonsense_uri'}"])
             )
 
+    def test_cli_connections_add_invalid_type(self):
+        with warnings.catch_warnings(record=True):
+            connection_command.connections_add(
+                self.parser.parse_args(
+                    ["connections", "add", "fsconn", "--conn-host=/tmp", "--conn-type=File"]
+                )
+            )
+
 
 class TestCliDeleteConnections:
     parser = cli_parser.get_parser()