eodag 3.0.1__py3-none-any.whl → 3.1.0__py3-none-any.whl

eodag/plugins/download/aws.py

@@ -23,21 +23,7 @@ import re
 from datetime import datetime
 from itertools import chain
 from pathlib import Path
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Callable,
-    Dict,
-    Iterator,
-    List,
-    Match,
-    Optional,
-    Set,
-    Tuple,
-    TypedDict,
-    Union,
-    cast,
-)
+from typing import TYPE_CHECKING, Any, Callable, Iterator, Optional, Union, cast
 
 import boto3
 import requests
@@ -74,6 +60,7 @@ from eodag.utils.exceptions import (
     NotAvailableError,
     TimeOutError,
 )
+from eodag.utils.s3 import open_s3_zipped_object
 
 if TYPE_CHECKING:
     from boto3.resources.collection import ResourceCollection
@@ -81,6 +68,7 @@ if TYPE_CHECKING:
     from eodag.api.product import EOProduct
     from eodag.api.search_result import SearchResult
     from eodag.config import PluginConfig
+    from eodag.types import S3SessionKwargs
     from eodag.types.download_args import DownloadConf
     from eodag.utils import DownloadedCallback, Unpack
 
@@ -208,6 +196,7 @@ AWS_AUTH_ERROR_MESSAGES = [
     "AccessDenied",
     "InvalidAccessKeyId",
     "SignatureDoesNotMatch",
+    "InvalidRequest",
 ]
 
 
@@ -230,14 +219,14 @@ class AwsDownload(Download):
         * :attr:`~eodag.config.PluginConfig.bucket_path_level` (``int``): at which level of the
           path part of the url the bucket can be found; If no bucket_path_level is given, the bucket
           is taken from the first element of the netloc part.
-        * :attr:`~eodag.config.PluginConfig.products` (``Dict[str, Dict[str, Any]``): product type
+        * :attr:`~eodag.config.PluginConfig.products` (``dict[str, dict[str, Any]``): product type
           specific config; the keys are the product types, the values are dictionaries which can contain the keys:
 
           * **default_bucket** (``str``): bucket where the product type can be found
           * **complementary_url_key** (``str``): keys to add additional urls
           * **build_safe** (``bool``): if a SAFE (Standard Archive Format for Europe) product should
             be created; used for Sentinel products; default: False
-          * **fetch_metadata** (``Dict[str, Any]``): config for metadata to be fetched for the SAFE product
+          * **fetch_metadata** (``dict[str, Any]``): config for metadata to be fetched for the SAFE product
 
     """
 
@@ -245,14 +234,15 @@ class AwsDownload(Download):
         super(AwsDownload, self).__init__(provider, config)
         self.requester_pays = getattr(self.config, "requester_pays", False)
         self.s3_session: Optional[boto3.session.Session] = None
+        self.s3_resource: Optional[boto3.resources.base.ServiceResource] = None
 
     def download(
         self,
         product: EOProduct,
-        auth: Optional[Union[AuthBase, Dict[str, str]]] = None,
+        auth: Optional[Union[AuthBase, S3SessionKwargs]] = None,
         progress_callback: Optional[ProgressCallback] = None,
-        wait: int = DEFAULT_DOWNLOAD_WAIT,
-        timeout: int = DEFAULT_DOWNLOAD_TIMEOUT,
+        wait: float = DEFAULT_DOWNLOAD_WAIT,
+        timeout: float = DEFAULT_DOWNLOAD_TIMEOUT,
         **kwargs: Unpack[DownloadConf],
     ) -> Optional[str]:
         """Download method for AWS S3 API.
@@ -338,19 +328,32 @@ class AwsDownload(Download):
             bucket_names_and_prefixes, auth
         )
 
+        # files in zip
+        updated_bucket_names_and_prefixes = self._download_file_in_zip(
+            product, bucket_names_and_prefixes, product_local_path, progress_callback
+        )
+        # prevent nothing-to-download errors if download was performed in zip
+        raise_error = (
+            False
+            if len(updated_bucket_names_and_prefixes) != len(bucket_names_and_prefixes)
+            else True
+        )
+
         # downloadable files
         unique_product_chunks = self._get_unique_products(
-            bucket_names_and_prefixes,
+            updated_bucket_names_and_prefixes,
             authenticated_objects,
             asset_filter,
             ignore_assets,
             product,
+            raise_error=raise_error,
         )
 
         total_size = sum([p.size for p in unique_product_chunks]) or None
 
         # download
-        progress_callback.reset(total=total_size)
+        if len(unique_product_chunks) > 0:
+            progress_callback.reset(total=total_size)
         try:
             for product_chunk in unique_product_chunks:
                 try:
@@ -402,17 +405,65 @@ class AwsDownload(Download):
 
         return product_local_path
 
+    def _download_file_in_zip(
+        self, product, bucket_names_and_prefixes, product_local_path, progress_callback
+    ):
+        """
+        Download file in zip from a prefix like `foo/bar.zip!file.txt`
+        """
+        if self.s3_resource is None:
+            logger.debug("Cannot check files in s3 zip without s3 resource")
+            return bucket_names_and_prefixes
+
+        s3_client = self.s3_resource.meta.client
+
+        downloaded = []
+        for i, pack in enumerate(bucket_names_and_prefixes):
+            bucket_name, prefix = pack
+            if ".zip!" in prefix:
+                splitted_path = prefix.split(".zip!")
+                zip_prefix = f"{splitted_path[0]}.zip"
+                rel_path = splitted_path[-1]
+                dest_file = os.path.join(product_local_path, rel_path)
+                dest_abs_path_dir = os.path.dirname(dest_file)
+                if not os.path.isdir(dest_abs_path_dir):
+                    os.makedirs(dest_abs_path_dir)
+
+                with open_s3_zipped_object(
+                    bucket_name, zip_prefix, s3_client, partial=False
+                ) as zip_file:
+                    # file size
+                    file_info = zip_file.getinfo(rel_path)
+                    progress_callback.reset(total=file_info.file_size)
+                    with zip_file.open(rel_path) as extracted, open(
+                        dest_file, "wb"
+                    ) as output_file:
+                        # Read in 1MB chunks
+                        for zchunk in iter(lambda: extracted.read(1024 * 1024), b""):
+                            output_file.write(zchunk)
+                            progress_callback(len(zchunk))
+
+                downloaded.append(i)
+
+        return [
+            pack
+            for i, pack in enumerate(bucket_names_and_prefixes)
+            if i not in downloaded
+        ]
+
     def _download_preparation(
         self,
         product: EOProduct,
         progress_callback: ProgressCallback,
         **kwargs: Unpack[DownloadConf],
-    ) -> Tuple[Optional[str], Optional[str]]:
+    ) -> tuple[Optional[str], Optional[str]]:
         """
-        preparation for the download:
+        Preparation for the download:
+
         - check if file was already downloaded
         - get file path
         - create directories
+
         :param product: product to be downloaded
         :param progress_callback: progress callback to be used
         :param kwargs: additional arguments
@@ -436,7 +487,8 @@ class AwsDownload(Download):
 
     def _configure_safe_build(self, build_safe: bool, product: EOProduct):
         """
-        updates the product properties with fetch metadata if safe build is enabled
+        Updates the product properties with fetch metadata if safe build is enabled
+
         :param build_safe: if safe build is enabled
         :param product: product to be updated
         """
@@ -480,9 +532,10 @@ class AwsDownload(Download):
         product: EOProduct,
         asset_filter: Optional[str] = None,
         ignore_assets: Optional[bool] = False,
-    ) -> List[Tuple[str, Optional[str]]]:
+    ) -> list[tuple[str, Optional[str]]]:
         """
-        retrieves the bucket names and path prefixes for the assets
+        Retrieves the bucket names and path prefixes for the assets
+
         :param product: product for which the assets shall be downloaded
         :param asset_filter: text for which the assets should be filtered
         :param ignore_assets: if product instead of individual assets should be used
@@ -521,14 +574,15 @@ class AwsDownload(Download):
 
     def _do_authentication(
         self,
-        bucket_names_and_prefixes: List[Tuple[str, Optional[str]]],
-        auth: Optional[Union[AuthBase, Dict[str, str]]] = None,
-    ) -> Tuple[Dict[str, Any], ResourceCollection]:
+        bucket_names_and_prefixes: list[tuple[str, Optional[str]]],
+        auth: Optional[Union[AuthBase, S3SessionKwargs]] = None,
+    ) -> tuple[dict[str, Any], ResourceCollection]:
         """
-        authenticates with s3 and retrieves the available objects
-        raises an error when authentication is not possible
+        Authenticates with s3 and retrieves the available objects
+
         :param bucket_names_and_prefixes: list of bucket names and corresponding path prefixes
         :param auth: authentication information
+        :raises AuthenticationError: authentication is not possible
         :return: authenticated objects per bucket, list of available objects
         """
         if not isinstance(auth, (dict, type(None))):
@@ -537,8 +591,8 @@ class AwsDownload(Download):
             )
         if auth is None:
             auth = {}
-        authenticated_objects: Dict[str, Any] = {}
-        auth_error_messages: Set[str] = set()
+        authenticated_objects: dict[str, Any] = {}
+        auth_error_messages: set[str] = set()
         for _, pack in enumerate(bucket_names_and_prefixes):
             try:
                 bucket_name, prefix = pack
@@ -590,22 +644,25 @@ class AwsDownload(Download):
 
     def _get_unique_products(
         self,
-        bucket_names_and_prefixes: List[Tuple[str, Optional[str]]],
-        authenticated_objects: Dict[str, Any],
+        bucket_names_and_prefixes: list[tuple[str, Optional[str]]],
+        authenticated_objects: dict[str, Any],
         asset_filter: Optional[str],
         ignore_assets: bool,
         product: EOProduct,
-    ) -> Set[Any]:
+        raise_error: bool = True,
+    ) -> set[Any]:
         """
-        retrieve unique product chunks based on authenticated objects and asset filters
+        Retrieve unique product chunks based on authenticated objects and asset filters
+
         :param bucket_names_and_prefixes: list of bucket names and corresponding path prefixes
         :param authenticated_objects: available objects per bucket
         :param asset_filter: text for which assets should be filtered
         :param ignore_assets: if product instead of individual assets should be used
         :param product: product that shall be downloaded
+        :param raise_error: raise error if there is nothing to download
         :return: set of product chunks that can be downloaded
         """
-        product_chunks: List[Any] = []
+        product_chunks: list[Any] = []
         for bucket_name, prefix in bucket_names_and_prefixes:
             # unauthenticated items filtered out
             if bucket_name in authenticated_objects.keys():
@@ -624,19 +681,19 @@ class AwsDownload(Download):
                     unique_product_chunks,
                 )
             )
-            if not unique_product_chunks:
+            if not unique_product_chunks and raise_error:
                 raise NotAvailableError(
                     rf"No file basename matching re.fullmatch(r'{asset_filter}') was found in {product.remote_location}"
                 )
 
-        if not unique_product_chunks:
+        if not unique_product_chunks and raise_error:
             raise NoMatchingProductType("No product found to download.")
 
         return unique_product_chunks
 
     def _raise_if_auth_error(self, exception: ClientError) -> None:
         """Raises an error if given exception is an authentication error"""
-        err = cast(Dict[str, str], exception.response["Error"])
+        err = cast(dict[str, str], exception.response["Error"])
         if err["Code"] in AWS_AUTH_ERROR_MESSAGES and "key" in err["Message"].lower():
             raise AuthenticationError(
                 f"Please check your credentials for {self.provider}.",
@@ -647,10 +704,10 @@ class AwsDownload(Download):
     def _stream_download_dict(
         self,
         product: EOProduct,
-        auth: Optional[Union[AuthBase, Dict[str, str]]] = None,
+        auth: Optional[Union[AuthBase, S3SessionKwargs]] = None,
         progress_callback: Optional[ProgressCallback] = None,
-        wait: int = DEFAULT_DOWNLOAD_WAIT,
-        timeout: int = DEFAULT_DOWNLOAD_TIMEOUT,
+        wait: float = DEFAULT_DOWNLOAD_WAIT,
+        timeout: float = DEFAULT_DOWNLOAD_TIMEOUT,
         **kwargs: Unpack[DownloadConf],
     ) -> StreamResponse:
         r"""
@@ -713,6 +770,13 @@ class AwsDownload(Download):
             bucket_names_and_prefixes, auth
         )
 
+        # stream not implemented for prefixes like `foo/bar.zip!file.txt`
+        for _, prefix in bucket_names_and_prefixes:
+            if prefix and ".zip!" in prefix:
+                raise NotImplementedError(
+                    "Download streaming is not implemented for files in zip on S3"
+                )
+
         # downloadable files
         unique_product_chunks = self._get_unique_products(
             bucket_names_and_prefixes,
@@ -731,12 +795,12 @@ class AwsDownload(Download):
             else sanitize(product.properties.get("id", "download"))
         )
 
-        if len(assets_values) == 1:
+        if len(assets_values) <= 1:
             first_chunks_tuple = next(chunks_tuples)
             # update headers
             filename = os.path.basename(list(unique_product_chunks)[0].key)
             headers = {"content-disposition": f"attachment; filename={filename}"}
-            if assets_values[0].get("type", None):
+            if assets_values and assets_values[0].get("type", None):
                 headers["content-type"] = assets_values[0]["type"]
 
             return StreamResponse(
@@ -753,11 +817,11 @@ class AwsDownload(Download):
 
     def _stream_download(
         self,
-        unique_product_chunks: Set[Any],
+        unique_product_chunks: set[Any],
         product: EOProduct,
         build_safe: bool,
         progress_callback: ProgressCallback,
-        assets_values: List[Dict[str, Any]],
+        assets_values: list[dict[str, Any]],
     ) -> Iterator[Any]:
         """Yield product data chunks"""
 
@@ -799,7 +863,6 @@ class AwsDownload(Download):
             common_path = self._get_commonpath(
                 product, unique_product_chunks, build_safe
             )
-
         for product_chunk in unique_product_chunks:
             try:
                 chunk_rel_path = self.get_chunk_dest_path(
@@ -817,8 +880,7 @@ class AwsDownload(Download):
                 # out of SAFE format chunk
                 logger.warning(e)
                 continue
-
-            if len(assets_values) == 1:
+            if len(assets_values) <= 1:
                 yield from get_chunk_parts(product_chunk, progress_callback)
             else:
                 yield (
@@ -830,7 +892,7 @@ class AwsDownload(Download):
                 )
 
     def _get_commonpath(
-        self, product: EOProduct, product_chunks: Set[Any], build_safe: bool
+        self, product: EOProduct, product_chunks: set[Any], build_safe: bool
     ) -> str:
         chunk_paths = []
         for product_chunk in product_chunks:
@@ -840,8 +902,8 @@ class AwsDownload(Download):
         return os.path.commonpath(chunk_paths)
 
     def get_rio_env(
-        self, bucket_name: str, prefix: str, auth_dict: Dict[str, str]
-    ) -> Dict[str, Any]:
+        self, bucket_name: str, prefix: str, auth_dict: S3SessionKwargs
+    ) -> dict[str, Any]:
         """Get rasterio environment variables needed for data access authentication.
 
         :param bucket_name: Bucket containg objects
@@ -849,23 +911,26 @@ class AwsDownload(Download):
         :param auth_dict: Dictionary containing authentication keys
         :returns: The rasterio environement variables
         """
-        if self.s3_session is not None:
-            if self.requester_pays:
-                return {"session": self.s3_session, "requester_pays": True}
-            else:
-                return {"session": self.s3_session}
+        rio_env_kwargs = {}
+        if endpoint_url := getattr(self.config, "s3_endpoint", None):
+            rio_env_kwargs["endpoint_url"] = endpoint_url.split("://")[-1]
+        rio_env_kwargs |= auth_dict
+
+        if self.s3_session is None:
+            _ = self.get_authenticated_objects(bucket_name, prefix, auth_dict)
 
-        _ = self.get_authenticated_objects(bucket_name, prefix, auth_dict)
         if self.s3_session is not None:
             if self.requester_pays:
-                return {"session": self.s3_session, "requester_pays": True}
-            else:
-                return {"session": self.s3_session}
+                rio_env_kwargs["requester_pays"] = True
+            return {
+                "session": self.s3_session,
+                **rio_env_kwargs,
+            }
         else:
-            return {"aws_unsigned": True}
+            return {"aws_unsigned": True, **rio_env_kwargs}
 
     def get_authenticated_objects(
-        self, bucket_name: str, prefix: str, auth_dict: Dict[str, str]
+        self, bucket_name: str, prefix: str, auth_dict: S3SessionKwargs
     ) -> ResourceCollection:
         """Get boto3 authenticated objects for the given bucket using
         the most adapted auth strategy.
@@ -877,8 +942,8 @@ class AwsDownload(Download):
         :param auth_dict: Dictionary containing authentication keys
         :returns: The boto3 authenticated objects
         """
-        auth_methods: List[
-            Callable[[str, str, Dict[str, str]], Optional[ResourceCollection]]
+        auth_methods: list[
+            Callable[[str, str, S3SessionKwargs], Optional[ResourceCollection]]
         ] = [
             self._get_authenticated_objects_unsigned,
             self._get_authenticated_objects_from_auth_profile,
@@ -913,7 +978,7 @@ class AwsDownload(Download):
         )
 
     def _get_authenticated_objects_unsigned(
-        self, bucket_name: str, prefix: str, auth_dict: Dict[str, str]
+        self, bucket_name: str, prefix: str, auth_dict: S3SessionKwargs
     ) -> Optional[ResourceCollection]:
         """Auth strategy using no-sign-request"""
 
@@ -928,7 +993,7 @@ class AwsDownload(Download):
         return objects
 
     def _get_authenticated_objects_from_auth_profile(
-        self, bucket_name: str, prefix: str, auth_dict: Dict[str, str]
+        self, bucket_name: str, prefix: str, auth_dict: S3SessionKwargs
     ) -> Optional[ResourceCollection]:
         """Auth strategy using RequestPayer=requester and ``aws_profile`` from provided credentials"""
 
@@ -946,26 +1011,18 @@ class AwsDownload(Download):
                 objects = s3_resource.Bucket(bucket_name).objects
             list(objects.filter(Prefix=prefix).limit(1))
             self.s3_session = s3_session
+            self.s3_resource = s3_resource
             return objects
         else:
             return None
 
     def _get_authenticated_objects_from_auth_keys(
-        self, bucket_name: str, prefix: str, auth_dict: Dict[str, str]
+        self, bucket_name: str, prefix: str, auth_dict: S3SessionKwargs
     ) -> Optional[ResourceCollection]:
         """Auth strategy using RequestPayer=requester and ``aws_access_key_id``/``aws_secret_access_key``
         from provided credentials"""
 
         if all(k in auth_dict for k in ("aws_access_key_id", "aws_secret_access_key")):
-            S3SessionKwargs = TypedDict(
-                "S3SessionKwargs",
-                {
-                    "aws_access_key_id": str,
-                    "aws_secret_access_key": str,
-                    "aws_session_token": str,
-                },
-                total=False,
-            )
             s3_session_kwargs: S3SessionKwargs = {
                 "aws_access_key_id": auth_dict["aws_access_key_id"],
                 "aws_secret_access_key": auth_dict["aws_secret_access_key"],
@@ -985,12 +1042,13 @@ class AwsDownload(Download):
                 objects = s3_resource.Bucket(bucket_name).objects
             list(objects.filter(Prefix=prefix).limit(1))
             self.s3_session = s3_session
+            self.s3_resource = s3_resource
             return objects
         else:
             return None
 
     def _get_authenticated_objects_from_env(
-        self, bucket_name: str, prefix: str, auth_dict: Dict[str, str]
+        self, bucket_name: str, prefix: str, auth_dict: S3SessionKwargs
     ) -> Optional[ResourceCollection]:
         """Auth strategy using RequestPayer=requester and current environment"""
 
@@ -1006,11 +1064,12 @@ class AwsDownload(Download):
             objects = s3_resource.Bucket(bucket_name).objects
         list(objects.filter(Prefix=prefix).limit(1))
         self.s3_session = s3_session
+        self.s3_resource = s3_resource
         return objects
 
     def get_product_bucket_name_and_prefix(
         self, product: EOProduct, url: Optional[str] = None
-    ) -> Tuple[str, Optional[str]]:
+    ) -> tuple[str, Optional[str]]:
         """Extract bucket name and prefix from product URL
 
         :param product: The EO product to download
@@ -1141,7 +1200,7 @@ class AwsDownload(Download):
         s1_title_suffix: Optional[str] = None
         # S2 common
         if product.product_type and "S2_MSI" in product.product_type:
-            title_search: Optional[Match[str]] = re.search(
+            title_search: Optional[re.Match[str]] = re.search(
                 r"^\w+_\w+_(\w+)_(\w+)_(\w+)_(\w+)_(\w+)$",
                 product.properties["title"],
             )
@@ -1327,13 +1386,13 @@ class AwsDownload(Download):
     def download_all(
         self,
         products: SearchResult,
-        auth: Optional[Union[AuthBase, Dict[str, str]]] = None,
+        auth: Optional[Union[AuthBase, S3SessionKwargs]] = None,
         downloaded_callback: Optional[DownloadedCallback] = None,
         progress_callback: Optional[ProgressCallback] = None,
-        wait: int = DEFAULT_DOWNLOAD_WAIT,
-        timeout: int = DEFAULT_DOWNLOAD_TIMEOUT,
+        wait: float = DEFAULT_DOWNLOAD_WAIT,
+        timeout: float = DEFAULT_DOWNLOAD_TIMEOUT,
         **kwargs: Unpack[DownloadConf],
-    ) -> List[str]:
+    ) -> list[str]:
         """
         download_all using parent (base plugin) method
         """