# 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.
from __future__ import annotations
import asyncio
import logging
from functools import partial
from typing import TYPE_CHECKING, Any, Callable, Dict
import requests
from botocore import UNSIGNED
from requests import HTTPError
from airflow.providers.amazon.aws.hooks.s3 import S3Hook
if TYPE_CHECKING:
from botocore.awsrequest import AWSRequest
from fsspec import AbstractFileSystem
[docs]Properties = Dict[str, str]
[docs]S3_PROXY_URI = "proxy-uri"
[docs]log = logging.getLogger(__name__)
[docs]schemes = ["s3", "s3a", "s3n"]
[docs]class SignError(Exception):
"""Raises when unable to sign a S3 request."""
[docs]def get_fs(conn_id: str | None, storage_options: dict[str, str] | None = None) -> AbstractFileSystem:
try:
from s3fs import S3FileSystem
except ImportError:
raise ImportError(
"Airflow FS S3 protocol requires the s3fs library, but it is not installed as it requires"
"aiobotocore. Please install the s3 protocol support library by running: "
"pip install apache-airflow-providers-amazon[s3fs]"
)
s3_hook = S3Hook(aws_conn_id=conn_id)
session = s3_hook.get_session(deferrable=True)
endpoint_url = s3_hook.conn_config.get_service_endpoint_url(service_name="s3")
config_kwargs: dict[str, Any] = s3_hook.conn_config.extra_config.get("config_kwargs", {})
config_kwargs.update(storage_options or {})
register_events: dict[str, Callable[[Properties], None]] = {}
s3_service_config = s3_hook.service_config
if signer := s3_service_config.get("signer", None):
log.info("Loading signer %s", signer)
if singer_func := SIGNERS.get(signer):
uri = s3_service_config.get("signer_uri", None)
token = s3_service_config.get("signer_token", None)
if not uri or not token:
raise ValueError(f"Signer {signer} requires uri and token")
properties: Properties = {
"uri": uri,
"token": uri,
}
singer_func_with_properties = partial(singer_func, properties)
register_events["before-sign.s3"] = singer_func_with_properties
# Disable the AWS Signer
config_kwargs["signature_version"] = UNSIGNED
else:
raise ValueError(f"Signer not available: {signer}")
if proxy_uri := s3_service_config.get(S3_PROXY_URI, None):
config_kwargs["proxies"] = {"http": proxy_uri, "https": proxy_uri}
anon = False
if asyncio.run(session.get_credentials()) is None:
log.info("No credentials found, using anonymous access")
anon = True
fs = S3FileSystem(session=session, config_kwargs=config_kwargs, endpoint_url=endpoint_url, anon=anon)
for event_name, event_function in register_events.items():
fs.s3.meta.events.register_last(event_name, event_function, unique_id=1925)
return fs
[docs]def s3v4_rest_signer(properties: Properties, request: AWSRequest, **_: Any) -> AWSRequest:
if "token" not in properties:
raise SignError("Signer set, but token is not available")
signer_url = properties["uri"].rstrip("/")
signer_headers = {"Authorization": f"Bearer {properties['token']}"}
signer_body = {
"method": request.method,
"region": request.context["client_region"],
"uri": request.url,
"headers": {key: [val] for key, val in request.headers.items()},
}
response = requests.post(f"{signer_url}/v1/aws/s3/sign", headers=signer_headers, json=signer_body)
try:
response.raise_for_status()
response_json = response.json()
except HTTPError as e:
raise SignError(f"Failed to sign request {response.status_code}: {signer_body}") from e
for key, value in response_json["headers"].items():
request.headers.add_header(key, ", ".join(value))
request.url = response_json["uri"]
return request
[docs]SIGNERS: dict[str, Callable[[Properties, AWSRequest], AWSRequest]] = {"S3V4RestSigner": s3v4_rest_signer}