Add VersionId parameter for S3 reads (#747)

Author: kukushking

awswrangler/s3/_describe.py

@@ -16,7 +16,10 @@
 
 
 def _describe_object(
-    path: str, boto3_session: boto3.Session, s3_additional_kwargs: Optional[Dict[str, Any]]
+    path: str,
+    boto3_session: boto3.Session,
+    s3_additional_kwargs: Optional[Dict[str, Any]],
+    version_id: Optional[str] = None,
 ) -> Tuple[str, Dict[str, Any]]:
     client_s3: boto3.client = _utils.client(service_name="s3", session=boto3_session)
     bucket: str
@@ -28,21 +31,30 @@ def _describe_object(
         )
     else:
         extra_kwargs = {}
-    desc: Dict[str, Any] = _utils.try_it(
+    desc: Dict[str, Any]
+    if version_id:
+        extra_kwargs["VersionId"] = version_id
+    desc = _utils.try_it(
         f=client_s3.head_object, ex=client_s3.exceptions.NoSuchKey, Bucket=bucket, Key=key, **extra_kwargs
     )
     return path, desc
 
 
 def _describe_object_concurrent(
-    path: str, boto3_primitives: _utils.Boto3PrimitivesType, s3_additional_kwargs: Optional[Dict[str, Any]]
+    path: str,
+    boto3_primitives: _utils.Boto3PrimitivesType,
+    s3_additional_kwargs: Optional[Dict[str, Any]],
+    version_id: Optional[str] = None,
 ) -> Tuple[str, Dict[str, Any]]:
     boto3_session = _utils.boto3_from_primitives(primitives=boto3_primitives)
-    return _describe_object(path=path, boto3_session=boto3_session, s3_additional_kwargs=s3_additional_kwargs)
+    return _describe_object(
+        path=path, boto3_session=boto3_session, s3_additional_kwargs=s3_additional_kwargs, version_id=version_id
+    )
 
 
 def describe_objects(
     path: Union[str, List[str]],
+    version_id: Optional[Union[str, Dict[str, str]]] = None,
     use_threads: bool = True,
     last_modified_begin: Optional[datetime.datetime] = None,
     last_modified_end: Optional[datetime.datetime] = None,
@@ -75,6 +87,9 @@ def describe_objects(
     path : Union[str, List[str]]
         S3 prefix (accepts Unix shell-style wildcards)
         (e.g. s3://bucket/prefix) or list of S3 objects paths (e.g. [s3://bucket/key0, s3://bucket/key1]).
+    version_id: Optional[Union[str, Dict[str, str]]]
+        Version id of the object or mapping of object path to version id.
+        (e.g. {'s3://bucket/key0': '121212', 's3://bucket/key1': '343434'})
     use_threads : bool
         True to enable concurrent requests, False to disable multiple threads.
         If enabled os.cpu_count() will be used as the max number of threads.
@@ -116,20 +131,32 @@ def describe_objects(
     resp_list: List[Tuple[str, Dict[str, Any]]]
     if len(paths) == 1:
         resp_list = [
-            _describe_object(path=paths[0], boto3_session=boto3_session, s3_additional_kwargs=s3_additional_kwargs)
+            _describe_object(
+                path=paths[0],
+                version_id=version_id.get(paths[0]) if isinstance(version_id, dict) else version_id,
+                boto3_session=boto3_session,
+                s3_additional_kwargs=s3_additional_kwargs,
+            )
         ]
     elif use_threads is False:
         resp_list = [
-            _describe_object(path=p, boto3_session=boto3_session, s3_additional_kwargs=s3_additional_kwargs)
+            _describe_object(
+                path=p,
+                version_id=version_id.get(p) if isinstance(version_id, dict) else version_id,
+                boto3_session=boto3_session,
+                s3_additional_kwargs=s3_additional_kwargs,
+            )
             for p in paths
         ]
     else:
         cpus: int = _utils.ensure_cpu_count(use_threads=use_threads)
+        versions = [version_id.get(p) if isinstance(version_id, dict) else version_id for p in paths]
         with concurrent.futures.ThreadPoolExecutor(max_workers=cpus) as executor:
             resp_list = list(
                 executor.map(
                     _describe_object_concurrent,
                     paths,
+                    versions,
                     itertools.repeat(_utils.boto3_to_primitives(boto3_session=boto3_session)),
                     itertools.repeat(s3_additional_kwargs),
                 )
@@ -140,6 +167,7 @@ def describe_objects(
 
 def size_objects(
     path: Union[str, List[str]],
+    version_id: Optional[Union[str, Dict[str, str]]] = None,
     use_threads: bool = True,
     s3_additional_kwargs: Optional[Dict[str, Any]] = None,
     boto3_session: Optional[boto3.Session] = None,
@@ -162,6 +190,9 @@ def size_objects(
     path : Union[str, List[str]]
         S3 prefix (accepts Unix shell-style wildcards)
         (e.g. s3://bucket/prefix) or list of S3 objects paths (e.g. [s3://bucket/key0, s3://bucket/key1]).
+    version_id: Optional[Union[str, Dict[str, str]]]
+        Version id of the object or mapping of object path to version id.
+        (e.g. {'s3://bucket/key0': '121212', 's3://bucket/key1': '343434'})
     use_threads : bool
         True to enable concurrent requests, False to disable multiple threads.
         If enabled os.cpu_count() will be used as the max number of threads.
@@ -184,7 +215,11 @@ def size_objects(
 
     """
     desc_list: Dict[str, Dict[str, Any]] = describe_objects(
-        path=path, use_threads=use_threads, boto3_session=boto3_session, s3_additional_kwargs=s3_additional_kwargs
+        path=path,
+        version_id=version_id,
+        use_threads=use_threads,
+        boto3_session=boto3_session,
+        s3_additional_kwargs=s3_additional_kwargs,
     )
     size_dict: Dict[str, Optional[int]] = {k: d.get("ContentLength", None) for k, d in desc_list.items()}
     return size_dict

awswrangler/s3/_download.py

@@ -14,6 +14,7 @@
 def download(
     path: str,
     local_file: Union[str, Any],
+    version_id: Optional[str] = None,
     use_threads: bool = True,
     boto3_session: Optional[boto3.Session] = None,
     s3_additional_kwargs: Optional[Dict[str, Any]] = None,
@@ -31,6 +32,8 @@ def download(
         S3 path (e.g. ``s3://bucket/key0``).
     local_file : Union[str, Any]
         A file-like object in binary mode or a path to local file (e.g. ``./local/path/to/key0``).
+    version_id: Optional[str]
+        Version id of the object.
     use_threads : bool
         True to enable concurrent requests, False to disable multiple threads.
         If enabled os.cpu_count() will be used as the max number of threads.
@@ -64,6 +67,7 @@ def download(
         path=path,
         mode="rb",
         use_threads=use_threads,
+        version_id=version_id,
         s3_block_size=-1,  # One shot download
         s3_additional_kwargs=s3_additional_kwargs,
         boto3_session=session,

awswrangler/s3/_fs.py

@@ -30,6 +30,7 @@
         "SSECustomerKey",
         "RequestPayer",
         "ExpectedBucketOwner",
+        "VersionId",
     },
     "copy_object": {
         "ACL",
@@ -43,6 +44,7 @@
         "Tagging",
         "RequestPayer",
         "ExpectedBucketOwner",
+        "CopySource",
     },
     "create_multipart_upload": {
         "ACL",
@@ -87,10 +89,12 @@
     "delete_objects": {
         "RequestPayer",
         "ExpectedBucketOwner",
+        "Objects",
     },
     "head_object": {
         "RequestPayer",
         "ExpectedBucketOwner",
+        "VersionId",
     },
 }
 
@@ -106,10 +110,14 @@ def _fetch_range(
     key: str,
     s3_client: boto3.client,
     boto3_kwargs: Dict[str, Any],
+    version_id: Optional[str] = None,
 ) -> Tuple[int, bytes]:
     start, end = range_values
-    _logger.debug("Fetching: s3://%s/%s - Range: %s-%s", bucket, key, start, end)
-    resp: Dict[str, Any] = _utils.try_it(
+    _logger.debug("Fetching: s3://%s/%s - VersionId: %s - Range: %s-%s", bucket, key, version_id, start, end)
+    resp: Dict[str, Any]
+    if version_id:
+        boto3_kwargs["VersionId"] = version_id
+    resp = _utils.try_it(
         f=s3_client.get_object,
         ex=_S3_RETRYABLE_ERRORS,
         base=0.5,
@@ -230,12 +238,14 @@ def __init__(
         boto3_session: Optional[boto3.Session],
         newline: Optional[str],
         encoding: Optional[str],
+        version_id: Optional[str] = None,
     ) -> None:
         super().__init__()
         self._use_threads = use_threads
         self._newline: str = "\n" if newline is None else newline
         self._encoding: str = "utf-8" if encoding is None else encoding
         self._bucket, self._key = _utils.parse_path(path=path)
+        self._version_id = version_id
         self._boto3_session: boto3.Session = _utils.ensure_session(session=boto3_session)
         if mode not in {"rb", "wb", "r", "w"}:
             raise NotImplementedError("File mode must be {'rb', 'wb', 'r', 'w'}, not %s" % mode)
@@ -261,6 +271,7 @@ def __init__(
             self._end: int = 0
             size: Optional[int] = size_objects(
                 path=[path],
+                version_id=version_id,
                 use_threads=False,
                 boto3_session=self._boto3_session,
                 s3_additional_kwargs=self._s3_additional_kwargs,
@@ -325,6 +336,7 @@ def _fetch_range_proxy(self, start: int, end: int) -> bytes:
                 key=self._key,
                 s3_client=s3_client,
                 boto3_kwargs=boto3_kwargs,
+                version_id=self._version_id,
             )[1]
         sizes: Tuple[int, ...] = _utils.get_even_chunks_sizes(
             total_size=range_size, chunk_size=_MIN_PARALLEL_READ_BLOCK, upper_bound=False
@@ -344,6 +356,7 @@ def _fetch_range_proxy(self, start: int, end: int) -> bytes:
                         itertools.repeat(self._key),
                         itertools.repeat(s3_client),
                         itertools.repeat(boto3_kwargs),
+                        itertools.repeat(self._version_id),
                     )
                 ),
             )
@@ -596,6 +609,7 @@ def write(self, data: Union[bytes, bytearray, memoryview]) -> int:  # type: igno
 def open_s3_object(
     path: str,
     mode: str,
+    version_id: Optional[str] = None,
     use_threads: Union[bool, int] = False,
     s3_additional_kwargs: Optional[Dict[str, str]] = None,
     s3_block_size: int = -1,  # One shot download
@@ -611,6 +625,7 @@ def open_s3_object(
             path=path,
             s3_block_size=s3_block_size,
             mode=mode,
+            version_id=version_id,
             use_threads=use_threads,
             s3_additional_kwargs=s3_additional_kwargs,
             boto3_session=boto3_session,

awswrangler/s3/_list.py

@@ -137,7 +137,10 @@ def _list_objects(  # pylint: disable=too-many-branches
 
 
 def does_object_exist(
-    path: str, s3_additional_kwargs: Optional[Dict[str, Any]] = None, boto3_session: Optional[boto3.Session] = None
+    path: str,
+    s3_additional_kwargs: Optional[Dict[str, Any]] = None,
+    boto3_session: Optional[boto3.Session] = None,
+    version_id: Optional[str] = None,
 ) -> bool:
     """Check if object exists on S3.
 
@@ -187,6 +190,8 @@ def does_object_exist(
     else:
         extra_kwargs = {}
     try:
+        if version_id:
+            extra_kwargs["VersionId"] = version_id
         client_s3.head_object(Bucket=bucket, Key=key, **extra_kwargs)
         return True
     except botocore.exceptions.ClientError as ex:

awswrangler/s3/_read.py

@@ -128,13 +128,18 @@ def _union(dfs: List[pd.DataFrame], ignore_index: Optional[bool]) -> pd.DataFram
 
 
 def _read_dfs_from_multiple_paths(
-    read_func: Callable[..., pd.DataFrame], paths: List[str], use_threads: Union[bool, int], kwargs: Dict[str, Any]
+    read_func: Callable[..., pd.DataFrame],
+    paths: List[str],
+    version_ids: Optional[Dict[str, str]],
+    use_threads: Union[bool, int],
+    kwargs: Dict[str, Any],
 ) -> List[pd.DataFrame]:
     cpus = ensure_cpu_count(use_threads)
     if cpus < 2:
-        return [read_func(path, **kwargs) for path in paths]
+        return [read_func(path, version_id=version_ids.get(path) if version_ids else None, **kwargs) for path in paths]
 
     with concurrent.futures.ThreadPoolExecutor(max_workers=ensure_cpu_count(use_threads)) as executor:
         kwargs["boto3_session"] = boto3_to_primitives(kwargs["boto3_session"])
         partial_read_func = partial(read_func, **kwargs)
-        return list(df for df in executor.map(partial_read_func, paths))
+        versions = [version_ids.get(p) if isinstance(version_ids, dict) else None for p in paths]
+        return list(df for df in executor.map(partial_read_func, paths, versions))

awswrangler/s3/_read_excel.py

@@ -14,6 +14,7 @@
 
 def read_excel(
     path: str,
+    version_id: Optional[str] = None,
     use_threads: Union[bool, int] = True,
     boto3_session: Optional[boto3.Session] = None,
     s3_additional_kwargs: Optional[Dict[str, Any]] = None,
@@ -38,8 +39,10 @@ def read_excel(
 
     Parameters
     ----------
-    path : Union[str, List[str]]
+    path : str
         S3 path (e.g. ``s3://bucket/key.xlsx``).
+    version_id : Optional[str]
+        Version id of the object.
     use_threads : Union[bool, int]
         True to enable concurrent requests, False to disable multiple threads.
         If enabled os.cpu_count() will be used as the max number of threads.
@@ -77,6 +80,7 @@ def read_excel(
     with open_s3_object(
         path=path,
         mode="rb",
+        version_id=version_id,
         use_threads=use_threads,
         s3_block_size=-1,  # One shot download
         s3_additional_kwargs=s3_additional_kwargs,