CLI: secure connection commands by masking sensitive values

airflow/cli/commands/connection_command.py

@@ -0,0 +1,382 @@
+# 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.
+"""Connection sub-commands."""
+
+from __future__ import annotations
+
+import json
+import os
+import warnings
+from functools import cache
+from pathlib import Path
+from typing import Any
+from urllib.parse import urlsplit, urlunsplit
+
+from sqlalchemy import select
+from sqlalchemy.orm import exc
+
+from airflow.cli.simple_table import AirflowConsole
+from airflow.cli.utils import is_stdout, print_export_output
+from airflow.configuration import conf
+from airflow.exceptions import AirflowNotFoundException
+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, helpers, yaml
+from airflow.utils.cli import suppress_logs_and_warning
+from airflow.utils.db import create_default_connections as db_create_default_connections
+from airflow.utils.providers_configuration_loader import providers_configuration_loaded
+from airflow.utils.session import create_session
+
+
+def _connection_mapper(conn: Connection) -> dict[str, Any]:
+    return {
+        "id": conn.id,
+        "conn_id": conn.conn_id,
+        "conn_type": conn.conn_type,
+        "description": conn.description,
+        "host": conn.host,
+        "schema": conn.schema,
+        "login": conn.login,
+        "password": conn.password,
+        "port": conn.port,
+        "is_encrypted": conn.is_encrypted,
+        "is_extra_encrypted": conn.is_encrypted,
+        "extra_dejson": conn.extra_dejson,
+        "get_uri": conn.get_uri(),
+    }
+
+
+@suppress_logs_and_warning
+@providers_configuration_loaded
+def connections_get(args):
+    """Get a connection."""
+    os.environ["_AIRFLOW_PROCESS_CONTEXT"] = "server"
+
+    try:
+        conn = Connection.get_connection_from_secrets(args.conn_id)
+    except AirflowNotFoundException:
+        raise SystemExit("Connection not found.")
+    AirflowConsole().print_as(
+        data=[conn],
+        output=args.output,
+        mapper=_connection_mapper,
+    )
+
+
+@suppress_logs_and_warning
+@providers_configuration_loaded
+def connections_list(args):
+    """List all connections at the command line."""
+    with create_session() as session:
+        query = select(Connection)
+        conns = session.scalars(query).all()
+
+        AirflowConsole().print_as(
+            data=conns,
+            output=args.output,
+            mapper=_connection_mapper,
+        )
+
+
+def _connection_to_dict(conn: Connection) -> dict:
+    return {
+        "conn_type": conn.conn_type,
+        "description": conn.description,
+        "login": conn.login,
+        "password": conn.password,
+        "host": conn.host,
+        "port": conn.port,
+        "schema": conn.schema,
+        "extra": conn.extra,
+    }
+
+
+def create_default_connections(args):
+    db_create_default_connections()
+
+
+def _format_connections(conns: list[Connection], file_format: str, serialization_format: str) -> str:
+    if serialization_format == "json":
+
+        def serializer_func(x):
+            return json.dumps(_connection_to_dict(x))
+
+    elif serialization_format == "uri":
+        serializer_func = Connection.get_uri
+    else:
+        raise SystemExit(f"Received unexpected value for `--serialization-format`: {serialization_format!r}")
+    if file_format == ".env":
+        connections_env = ""
+        for conn in conns:
+            connections_env += f"{conn.conn_id}={serializer_func(conn)}\n"
+        return connections_env
+
+    connections_dict = {}
+    for conn in conns:
+        connections_dict[conn.conn_id] = _connection_to_dict(conn)
+
+    if file_format == ".yaml":
+        return yaml.dump(connections_dict)
+
+    if file_format == ".json":
+        return json.dumps(connections_dict, indent=2)
+
+    return json.dumps(connections_dict)
+
+
+def _valid_uri(uri: str) -> bool:
+    """Check if a URI is valid, by checking if scheme (conn_type) provided."""
+    return urlsplit(uri).scheme != ""
+
+
+@cache
+def _get_connection_types() -> list[str]:
+    """Return connection types available."""
+    _connection_types = []
+    providers_manager = ProvidersManager()
+    for connection_type, provider_info in providers_manager.hooks.items():
+        if provider_info:
+            _connection_types.append(connection_type)
+    return _connection_types
+
+
+@providers_configuration_loaded
+def connections_export(args):
+    """Export all connections to a file."""
+    file_formats = [".yaml", ".json", ".env"]
+    if args.format:
+        warnings.warn(
+            "Option `--format` is deprecated. Use `--file-format` instead.", DeprecationWarning, stacklevel=3
+        )
+    if args.format and args.file_format:
+        raise SystemExit("Option `--format` is deprecated.  Use `--file-format` instead.")
+    default_format = ".json"
+    provided_file_format = None
+    if args.format or args.file_format:
+        provided_file_format = f".{(args.format or args.file_format).lower()}"
+
+    with args.file as f:
+        if is_stdout(f):
+            filetype = provided_file_format or default_format
+        elif provided_file_format:
+            filetype = provided_file_format
+        else:
+            filetype = Path(args.file.name).suffix.lower()
+            if filetype not in file_formats:
+                raise SystemExit(
+                    f"Unsupported file format. The file must have the extension {', '.join(file_formats)}."
+                )
+
+        if args.serialization_format and filetype != ".env":
+            raise SystemExit("Option `--serialization-format` may only be used with file type `env`.")
+
+        with create_session() as session:
+            connections = session.scalars(select(Connection).order_by(Connection.conn_id)).all()
+
+        msg = _format_connections(
+            conns=connections,
+            file_format=filetype,
+            serialization_format=args.serialization_format or "uri",
+        )
+
+        f.write(msg)
+
+    print_export_output("Connections", connections, f)
+
+
+alternative_conn_specs = ["conn_type", "conn_host", "conn_login", "conn_password", "conn_schema", "conn_port"]
+
+
+@cli_utils.action_cli
+@providers_configuration_loaded
+def connections_add(args):
+    """Add new connection."""
+    has_uri = bool(args.conn_uri)
+    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")
+
+    if has_json and has_uri:
+        raise SystemExit("Cannot supply both conn-uri and conn-json")
+
+    if has_type and args.conn_type not in _get_connection_types():
+        warnings.warn(
+            f"The type provided to --conn-type is invalid: {args.conn_type}", UserWarning, stacklevel=4
+        )
+        warnings.warn(
+            f"Supported --conn-types are:{_get_connection_types()}."
+            "Hence overriding the conn-type with generic",
+            UserWarning,
+            stacklevel=4,
+        )
+        args.conn_type = "generic"
+
+    if has_uri or has_json:
+        invalid_args = []
+        if has_uri and not _valid_uri(args.conn_uri):
+            raise SystemExit(f"The URI provided to --conn-uri is invalid: {args.conn_uri}")
+
+        for arg in alternative_conn_specs:
+            if getattr(args, arg) is not None:
+                invalid_args.append(arg)
+
+        if has_json and args.conn_extra:
+            invalid_args.append("--conn-extra")
+
+        if invalid_args:
+            raise SystemExit(
+                "The following args are not compatible with "
+                f"the --conn-{'uri' if has_uri else 'json'} flag: {invalid_args!r}"
+            )
+
+    if args.conn_uri:
+        new_conn = Connection(conn_id=args.conn_id, description=args.conn_description, uri=args.conn_uri)
+        if args.conn_extra is not None:
+            new_conn.set_extra(args.conn_extra)
+    elif args.conn_json:
+        new_conn = Connection.from_json(conn_id=args.conn_id, value=args.conn_json)
+        if not new_conn.conn_type:
+            raise SystemExit("conn-json is invalid; must supply conn-type")
+    else:
+        new_conn = Connection(
+            conn_id=args.conn_id,
+            conn_type=args.conn_type,
+            description=args.conn_description,
+            host=args.conn_host,
+            login=args.conn_login,
+            password=args.conn_password,
+            schema=args.conn_schema,
+            port=args.conn_port,
+        )
+        if args.conn_extra is not None:
+            new_conn.set_extra(args.conn_extra)
+
+    with create_session() as session:
+        if not session.scalar(select(Connection).where(Connection.conn_id == new_conn.conn_id).limit(1)):
+            session.add(new_conn)
+            msg = "Successfully added `conn_id`={conn_id} : {uri}"
+            msg = msg.format(
+                conn_id=new_conn.conn_id,
+                uri=args.conn_uri
+                or urlunsplit(
+                    (
+                        new_conn.conn_type,
+                        f"{new_conn.login or ''}:{'******' if new_conn.password else ''}"
+                        f"@{new_conn.host or ''}:{new_conn.port or ''}",
+                        new_conn.schema or "",
+                        "",
+                        "",
+                    )
+                ),
+            )
+            print(msg)
+        else:
+            msg = f"A connection with `conn_id`={new_conn.conn_id} already exists."
+            raise SystemExit(msg)
+
+
+@cli_utils.action_cli
+@providers_configuration_loaded
+def connections_delete(args):
+    """Delete connection from DB."""
+    with create_session() as session:
+        try:
+            to_delete = session.scalars(select(Connection).where(Connection.conn_id == args.conn_id)).one()
+        except exc.NoResultFound:
+            raise SystemExit(f"Did not find a connection with `conn_id`={args.conn_id}")
+        except exc.MultipleResultsFound:
+            raise SystemExit(f"Found more than one connection with `conn_id`={args.conn_id}")
+        else:
+            session.delete(to_delete)
+            print(f"Successfully deleted connection with `conn_id`={to_delete.conn_id}")
+
+
+@cli_utils.action_cli(check_db=False)
+@providers_configuration_loaded
+def connections_import(args):
+    """Import connections from a file."""
+    if os.path.exists(args.file):
+        _import_helper(args.file, args.overwrite)
+    else:
+        raise SystemExit("Missing connections file.")
+
+
+def _import_helper(file_path: str, overwrite: bool) -> None:
+    """
+    Load connections from a file and save them to the DB.
+
+    :param overwrite: Whether to skip or overwrite on collision.
+    """
+    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.scalar(select(Connection.id).where(Connection.conn_id == conn_id))
+            if existing_conn_id is not None:
+                if not overwrite:
+                    print(f"Could not import connection {conn_id}: connection already exists.")
+                    continue
+
+                # The conn_ids match, but the PK of the new entry must also be the same as the old
+                conn.id = existing_conn_id
+
+            session.merge(conn)
+            session.commit()
+            print(f"Imported connection {conn_id}")
+
+
+@suppress_logs_and_warning
+@providers_configuration_loaded
+def connections_test(args) -> None:
+    """Test an Airflow connection."""
+    os.environ["_AIRFLOW_PROCESS_CONTEXT"] = "server"
+
+    console = AirflowConsole()
+    if conf.get("core", "test_connection", fallback="Disabled").lower().strip() != "enabled":
+        console.print(
+            "[bold yellow]\nTesting connections is disabled in Airflow configuration. "
+            "Contact your deployment admin to enable it.\n"
+        )
+        raise SystemExit(1)
+
+    print(f"Retrieving connection: {args.conn_id!r}")
+    try:
+        conn = Connection.get_connection_from_secrets(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")