feat(AwsDownload): zip partial download from s3 (#1561)

Author: sbrunato

eodag/plugins/download/aws.py

@@ -60,6 +60,7 @@
     NotAvailableError,
     TimeOutError,
 )
+from eodag.utils.s3 import open_s3_zipped_object
 
 if TYPE_CHECKING:
     from boto3.resources.collection import ResourceCollection
@@ -233,6 +234,7 @@ def __init__(self, provider: str, config: PluginConfig) -> None:
         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,
@@ -326,19 +328,32 @@ def 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:
@@ -390,17 +405,65 @@ def 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]]:
         """
-        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
@@ -424,7 +487,8 @@ def _download_preparation(
 
     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
         """
@@ -514,10 +578,11 @@ def _do_authentication(
         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))):
@@ -584,14 +649,17 @@ def _get_unique_products(
         asset_filter: Optional[str],
         ignore_assets: bool,
         product: EOProduct,
+        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] = []
@@ -613,12 +681,12 @@ def _get_unique_products(
                     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
@@ -702,6 +770,13 @@ def _stream_download_dict(
             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,
@@ -936,6 +1011,7 @@ def _get_authenticated_objects_from_auth_profile(
                 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
@@ -966,6 +1042,7 @@ def _get_authenticated_objects_from_auth_keys(
                 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
@@ -987,6 +1064,7 @@ def _get_authenticated_objects_from_env(
             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(

eodag/utils/s3.py

@@ -36,7 +36,7 @@
 )
 
 if TYPE_CHECKING:
-    from zipfile import ZipInfo
+    from zipfile import ZipFile, ZipInfo
 
     from mypy_boto3_s3.client import S3Client
 
@@ -78,18 +78,19 @@ def parse_int(bytes: bytes) -> int:
     return val
 
 
-def list_files_in_s3_zipped_object(
-    bucket_name: str, key_name: str, client_s3: S3Client
-) -> List[ZipInfo]:
+def open_s3_zipped_object(
+    bucket_name: str, key_name: str, client_s3: S3Client, partial: bool = True
+) -> ZipFile:
     """
-    List files in s3 zipped object, without downloading it.
+    Open s3 zipped object, without downloading it.
 
     See https://stackoverflow.com/questions/41789176/how-to-count-files-inside-zip-in-aws-s3-without-downloading-it;
     Based on https://stackoverflow.com/questions/51351000/read-zip-files-from-s3-without-downloading-the-entire-file
 
     :param bucket_name: Bucket name of the object to fetch
     :param key_name: Key name of the object to fetch
     :param client_s3: s3 client used to fetch the object
+    :param partial: fetch partial data if only content info is needed
     :returns: List of files in zip
     """
     response = client_s3.head_object(Bucket=bucket_name, Key=key_name)
@@ -104,11 +105,33 @@ def list_files_in_s3_zipped_object(
 
     # fetch central directory, append EOCD, and open as zipfile
     cd = fetch(bucket_name, key_name, cd_start, cd_size, client_s3)
-    zip = zipfile.ZipFile(io.BytesIO(cd + eocd))
 
-    logger.debug("Found %s files in %s" % (len(zip.filelist), key_name))
+    zip_data = (
+        cd + eocd if partial else fetch(bucket_name, key_name, 0, size, client_s3)
+    )
+
+    zip = zipfile.ZipFile(io.BytesIO(zip_data))
+
+    return zip
+
+
+def list_files_in_s3_zipped_object(
+    bucket_name: str, key_name: str, client_s3: S3Client
+) -> List[ZipInfo]:
+    """
+    List files in s3 zipped object, without downloading it.
+
+    See https://stackoverflow.com/questions/41789176/how-to-count-files-inside-zip-in-aws-s3-without-downloading-it;
+    Based on https://stackoverflow.com/questions/51351000/read-zip-files-from-s3-without-downloading-the-entire-file
 
-    return zip.filelist
+    :param bucket_name: Bucket name of the object to fetch
+    :param key_name: Key name of the object to fetch
+    :param client_s3: s3 client used to fetch the object
+    :returns: List of files in zip
+    """
+    with open_s3_zipped_object(bucket_name, key_name, client_s3) as zip_file:
+        logger.debug("Found %s files in %s" % (len(zip_file.filelist), key_name))
+        return zip_file.filelist
 
 
 def update_assets_from_s3(

tests/context.py

@@ -98,7 +98,11 @@
     get_ssl_context,
 )
 from eodag.utils.requests import fetch_json
-from eodag.utils.s3 import list_files_in_s3_zipped_object, update_assets_from_s3
+from eodag.utils.s3 import (
+    list_files_in_s3_zipped_object,
+    update_assets_from_s3,
+    open_s3_zipped_object,
+)
 from eodag.utils.exceptions import (
     AddressNotFound,
     AuthenticationError,

tests/units/test_download_plugins.py

@@ -1729,6 +1729,66 @@ def test_plugins_download_aws_no_safe_build_flatten_top_dirs(
             os.path.join(self.output_dir, self.product.properties["title"]),
         )
 
+    @mock.patch(
+        "eodag.plugins.download.aws.open_s3_zipped_object",
+        autospec=True,
+    )
+    @mock.patch(
+        "eodag.plugins.download.aws.AwsDownload.get_authenticated_objects",
+        autospec=True,
+    )
+    def test_plugins_download_aws_in_zip(
+        self,
+        mock_get_authenticated_objects: mock.Mock,
+        mock_open_s3_zipped_object: mock.Mock,
+    ):
+        """AwsDownload.download() must handle files in zip"""
+
+        def _open_zip(*args, **kwargs):
+            return zipfile.ZipFile(
+                os.path.join(
+                    TEST_RESOURCES_PATH,
+                    "products",
+                    "as_archive",
+                    "S2A_MSIL1C_20180101T105441_N0206_R051_T31TDH_20180101T124911.zip",
+                )
+            )
+
+        mock_open_s3_zipped_object.side_effect = _open_zip
+
+        plugin = self.get_download_plugin(self.product)
+        plugin.s3_resource = mock.Mock()
+        self.product.assets.clear()
+        self.product.assets.update(
+            {
+                "file1": {
+                    "href": (
+                        "http://example.com/path/to/foo.zip!"
+                        "GRANULE/L1C_T31TDH_A013204_20180101T105435/IMG_DATA/T31TDH_20180101T105441_B01.jp2"
+                    )
+                },
+                "file2": {
+                    "href": "http://example.com/path/to/foo.zip!GRANULE/L1C_T31TDH_A013204_20180101T105435/MTD_TL.xml"
+                },
+            }
+        )
+        # no SAFE build and flatten_top_dirs
+        plugin.config.products[self.product.product_type]["build_safe"] = False
+        plugin.config.flatten_top_dirs = True
+
+        path = plugin.download(self.product, output_dir=self.output_dir)
+
+        self.assertEqual(mock_open_s3_zipped_object.call_count, 2)
+        mock_open_s3_zipped_object.assert_called_with(
+            "example", "path/to/foo.zip", plugin.s3_resource.meta.client, partial=False
+        )
+        self.assertTrue(
+            os.path.isfile(
+                os.path.join(path, "IMG_DATA/T31TDH_20180101T105441_B01.jp2")
+            )
+        )
+        self.assertTrue(os.path.isfile(os.path.join(path, "MTD_TL.xml")))
+
     @mock.patch("eodag.plugins.download.aws.flatten_top_directories", autospec=True)
     @mock.patch(
         "eodag.plugins.download.aws.AwsDownload.check_manifest_file_list", autospec=True

tests/units/test_utils_s3.py

@@ -1,4 +1,5 @@
 import os
+import zipfile
 from unittest import TestCase
 
 import boto3
@@ -11,6 +12,7 @@
     MisconfiguredError,
     PluginConfig,
     list_files_in_s3_zipped_object,
+    open_s3_zipped_object,
     update_assets_from_s3,
 )
 
@@ -73,6 +75,14 @@ def test_utils_s3_list_files_in_s3_zipped_object(self):
             ],
         )
 
+    def test_utils_s3_open_s3_zipped_object(self):
+        """list_files_in_s3_zipped_object must list the files in a zipped object stored in S3"""
+        zip_file = open_s3_zipped_object(
+            "mybucket", "path/to/product.zip", self.s3_client
+        )
+        self.assertIsInstance(zip_file, zipfile.ZipFile)
+        self.assertEqual(len(zip_file.filelist), 6)
+
     def test_utils_s3_update_assets_from_s3_zip(self):
         """update_assets_from_s3 must update the assets of a product from a zipped object stored in S3"""
         prod = EOProduct("dummy", dict(geometry="POINT (0 0)", id="foo"))