diff --git a/airflow/api_connexion/endpoints/connection_endpoint.py b/airflow/api_connexion/endpoints/connection_endpoint.py index 7e2a7dd8da42d..641c3288d39e8 100644 --- a/airflow/api_connexion/endpoints/connection_endpoint.py +++ b/airflow/api_connexion/endpoints/connection_endpoint.py @@ -39,6 +39,7 @@ from airflow.models import Connection from airflow.secrets.environment_variables import CONN_ENV_PREFIX from airflow.security import permissions +from airflow.utils import helpers from airflow.utils.log.action_logger import action_event_from_permission from airflow.utils.session import NEW_SESSION, provide_session from airflow.utils.strings import get_random_string @@ -157,6 +158,10 @@ def post_connection(*, session: Session = NEW_SESSION) -> APIResponse: except ValidationError as err: raise BadRequest(detail=str(err.messages)) conn_id = data["conn_id"] + try: + helpers.validate_key(conn_id, max_length=200) + except Exception as e: + raise BadRequest(detail=str(e)) query = session.query(Connection) connection = query.filter_by(conn_id=conn_id).first() if not connection: diff --git a/airflow/cli/commands/connection_command.py b/airflow/cli/commands/connection_command.py index dc8c1792e3519..6d137bc896855 100644 --- a/airflow/cli/commands/connection_command.py +++ b/airflow/cli/commands/connection_command.py @@ -35,7 +35,7 @@ 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 import cli as cli_utils, helpers, yaml from airflow.utils.cli import suppress_logs_and_warning from airflow.utils.session import create_session @@ -203,6 +203,12 @@ def connections_add(args): has_json = bool(args.conn_json) has_type = bool(args.conn_type) + # Validate connection-id + try: + helpers.validate_key(args.conn_id, max_length=200) + except Exception as e: + raise SystemExit(f"Could not create connection. {e}") + if not has_type and not (has_json or has_uri): raise SystemExit("Must supply either conn-uri or conn-json if not supplying conn-type") @@ -313,6 +319,12 @@ def _import_helper(file_path: str, overwrite: bool) -> None: connections_dict = load_connections_dict(file_path) with create_session() as session: for conn_id, conn in connections_dict.items(): + try: + helpers.validate_key(conn_id, max_length=200) + except Exception as e: + print(f"Could not import connection. {e}") + continue + existing_conn_id = session.query(Connection.id).filter(Connection.conn_id == conn_id).scalar() if existing_conn_id is not None: if not overwrite: diff --git a/tests/api_connexion/endpoints/test_connection_endpoint.py b/tests/api_connexion/endpoints/test_connection_endpoint.py index a7df127d90f96..bce7150f8bda1 100644 --- a/tests/api_connexion/endpoints/test_connection_endpoint.py +++ b/tests/api_connexion/endpoints/test_connection_endpoint.py @@ -562,6 +562,20 @@ def test_post_should_respond_400_for_invalid_payload(self): "type": EXCEPTIONS_LINK_MAP[400], } + def test_post_should_respond_400_for_invalid_conn_id(self): + payload = {"connection_id": "****", "conn_type": "test_type"} + response = self.client.post( + "/api/v1/connections", json=payload, environ_overrides={"REMOTE_USER": "test"} + ) + assert response.status_code == 400 + assert response.json == { + "detail": "The key '****' has to be made of " + "alphanumeric characters, dashes, dots and underscores exclusively", + "status": 400, + "title": "Bad Request", + "type": EXCEPTIONS_LINK_MAP[400], + } + def test_post_should_respond_409_already_exist(self): payload = {"connection_id": "test-connection-id", "conn_type": "test_type"} response = self.client.post( diff --git a/tests/cli/commands/test_connection_command.py b/tests/cli/commands/test_connection_command.py index 692039fb9eefa..42f6fc8cc204e 100644 --- a/tests/cli/commands/test_connection_command.py +++ b/tests/cli/commands/test_connection_command.py @@ -635,6 +635,16 @@ def test_cli_connections_add_invalid_type(self): ) ) + def test_cli_connections_add_invalid_conn_id(self): + with pytest.raises(SystemExit) as e: + connection_command.connections_add( + self.parser.parse_args(["connections", "add", "Test$", f"--conn-uri={TEST_URL}"]) + ) + assert ( + e.value.args[0] == "Could not create connection. The key 'Test$' has to be made of " + "alphanumeric characters, dashes, dots and underscores exclusively" + ) + class TestCliDeleteConnections: parser = cli_parser.get_parser()