Repository: apache/airflow
File: airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py (Python)

Code before:
@connections_router.post("/test", dependencies=[Depends(requires_access_connection(method="POST"))])
def test_connection(
    test_body: ConnectionBody,
    session: SessionDep,
) -> ConnectionTestResponse:
    """
    Test an API connection.

    This method first creates an in-memory transient conn_id & exports that to an env var,
    as some hook classes tries to find out the `conn` from their __init__ method & errors out if not found.
    It also deletes the conn id env connection after the test.
    """
    if conf.get("core", "test_connection", fallback="Disabled").lower().strip() != "enabled":
        raise HTTPException(
            status.HTTP_403_FORBIDDEN,
            "Testing connections is disabled in Airflow configuration. "
            "Contact your deployment admin to enable it.",
        )

    transient_conn_id = get_random_string()
    conn_env_var = f"{CONN_ENV_PREFIX}{transient_conn_id.upper()}"
    try:
        data = test_body.model_dump(by_alias=True)
        data["conn_id"] = transient_conn_id

        existing_conn = session.scalar(
            select(Connection).where(Connection.conn_id == test_body.connection_id).limit(1)
        )
        if existing_conn:
            if data.get("password") is not None and existing_conn.password is not None:
                data["password"] = merge(data["password"], existing_conn.password, "password")
            if data.get("extra") is not None and existing_conn.extra is not None:
                try:
                    merged_extra = merge(json.loads(data["extra"]), json.loads(existing_conn.extra))
                    data["extra"] = json.dumps(merged_extra)
                except json.JSONDecodeError:
                    # Can't merge unstructured extra, keep the submitted value
                    pass

        conn = Connection(**data)
        os.environ[conn_env_var] = conn.get_uri()
        test_status, test_message = conn.test_connection()
        return ConnectionTestResponse.model_validate({"status": test_status, "message": test_message})
    finally:
        os.environ.pop(conn_env_var, None)


@connections_router.post(
    "/defaults",
    status_code=status.HTTP_204_NO_CONTENT,
    dependencies=[Depends(requires_access_connection(method="POST")), Depends(action_logging())],

Review comment (security):
Should not we use `Connection.get_connection_from_secrets` ? -> this is the "canonical" (even if badly named) way of retrieving connection. We should not read the connections directly from the database, because the connections might be defined in secrets?

Surely you won't be able to test those "secret backends only" connections via UI, but since this is an API call, it should still handle running "/test" API with the conn_id that is defined in secrets.

Code after:
def test_connection(test_body: ConnectionBody) -> ConnectionTestResponse:
    """
    Test an API connection.

    This method first creates an in-memory transient conn_id & exports that to an env var,
    as some hook classes tries to find out the `conn` from their __init__ method & errors out if not found.
    It also deletes the conn id env connection after the test.
    """
    if conf.get("core", "test_connection", fallback="Disabled").lower().strip() != "enabled":
        raise HTTPException(
            status.HTTP_403_FORBIDDEN,
            "Testing connections is disabled in Airflow configuration. "
            "Contact your deployment admin to enable it.",
        )

    transient_conn_id = get_random_string()
    conn_env_var = f"{CONN_ENV_PREFIX}{transient_conn_id.upper()}"
    try:
        # Try to get existing connection and merge with provided values
        try:
            existing_conn = Connection.get_connection_from_secrets(test_body.connection_id)
            existing_conn.conn_id = transient_conn_id
            update_orm_from_pydantic(existing_conn, test_body)
            conn = existing_conn
        except AirflowNotFoundException:
            data = test_body.model_dump(by_alias=True)
            data["conn_id"] = transient_conn_id
            conn = Connection(**data)

        os.environ[conn_env_var] = conn.get_uri()
        test_status, test_message = conn.test_connection()
        return ConnectionTestResponse.model_validate({"status": test_status, "message": test_message})
    finally:
        os.environ.pop(conn_env_var, None)


@connections_router.post(
    "/defaults",
    status_code=status.HTTP_204_NO_CONTENT,
    dependencies=[Depends(requires_access_connection(method="POST")), Depends(action_logging())],
)
def create_default_connections(
    session: SessionDep,
):
    """Create default connections."""
    db_create_default_connections(session)