Add boto3 session serializer/deserializer on _utils.py.

Author: igorborgest

awswrangler/_utils.py

@@ -1,10 +1,11 @@
 """Internal (private) Utilities Module."""
 
+import copy
 import logging
 import math
 import os
 import random
-from typing import Any, Dict, Generator, List, Optional, Tuple
+from typing import Any, Dict, Generator, List, Optional, Tuple, Union
 
 import boto3  # type: ignore
 import botocore.config  # type: ignore
@@ -17,8 +18,10 @@
 _logger: logging.Logger = logging.getLogger(__name__)
 
 
-def ensure_session(session: Optional[boto3.Session] = None) -> boto3.Session:
+def ensure_session(session: Optional[Union[boto3.Session, Dict[str, Optional[str]]]] = None) -> boto3.Session:
     """Ensure that a valid boto3.Session will be returned."""
+    if isinstance(session, dict):  # Primitives received
+        return boto3_from_primitives(primitives=session)
     if session is not None:
         return session
     # Ensure the boto3's default session is used so that its parameters can be
@@ -28,6 +31,30 @@ def ensure_session(session: Optional[boto3.Session] = None) -> boto3.Session:
     return boto3.Session()  # pragma: no cover
 
 
+def boto3_to_primitives(boto3_session: Optional[boto3.Session] = None) -> Dict[str, Optional[str]]:
+    """Convert Boto3 Session to Python primitives."""
+    _boto3_session: boto3.Session = ensure_session(session=boto3_session)
+    credentials = _boto3_session.get_credentials()
+    return {
+        "aws_access_key_id": getattr(credentials, "access_key", None),
+        "aws_secret_access_key": getattr(credentials, "secret_key", None),
+        "aws_session_token": getattr(credentials, "token", None),
+        "region_name": _boto3_session.region_name,
+        "profile_name": _boto3_session.profile_name,
+    }
+
+
+def boto3_from_primitives(primitives: Dict[str, Optional[str]] = None) -> boto3.Session:
+    """Convert Python primitives to Boto3 Session."""
+    if primitives is None:
+        return boto3.DEFAULT_SESSION  # pragma: no cover
+    _primitives: Dict[str, Optional[str]] = copy.deepcopy(primitives)
+    profile_name: Optional[str] = _primitives.get("profile_name", None)
+    _primitives["profile_name"] = None if profile_name in (None, "default") else profile_name
+    args: Dict[str, str] = {k: v for k, v in _primitives.items() if v is not None}
+    return boto3.Session(**args)
+
+
 def client(service_name: str, session: Optional[boto3.Session] = None) -> boto3.client:
     """Create a valid boto3.client."""
     return ensure_session(session=session).client(
@@ -139,7 +166,8 @@ def chunkify(lst: List[Any], num_chunks: int = 1, max_length: Optional[int] = No
 
 
 def get_fs(
-    session: Optional[boto3.Session] = None, s3_additional_kwargs: Optional[Dict[str, str]] = None
+    session: Optional[Union[boto3.Session, Dict[str, Optional[str]]]] = None,
+    s3_additional_kwargs: Optional[Dict[str, str]] = None,
 ) -> s3fs.S3FileSystem:
     """Build a S3FileSystem from a given boto3 session."""
     fs: s3fs.S3FileSystem = s3fs.S3FileSystem(

awswrangler/s3.py

@@ -1614,7 +1614,7 @@ def _read_text(
             path_root=path_root,
         )
         return dfs
-    if (use_threads is False) or (boto3_session is not None):
+    if use_threads is False:
         df: pd.DataFrame = pd.concat(
             objs=[
                 _read_text_full(
@@ -1640,7 +1640,7 @@ def _read_text(
                     repeat(parser_func),
                     repeat(path_root),
                     paths,
-                    repeat(None),  # Boto3.Session
+                    repeat(_utils.boto3_to_primitives(boto3_session=session)),  # Boto3.Session
                     repeat(pandas_kwargs),
                     repeat(s3_additional_kwargs),
                     repeat(dataset),
@@ -1683,7 +1683,7 @@ def _read_text_full(
     parser_func: Callable,
     path_root: str,
     path: str,
-    boto3_session: boto3.Session,
+    boto3_session: Union[boto3.Session, Dict[str, Optional[str]]],
     pandas_args: Dict[str, Any],
     s3_additional_kwargs: Optional[Dict[str, str]] = None,
     dataset: bool = False,
@@ -2354,29 +2354,38 @@ def _wait_objects(
     delay = 5 if delay is None else delay
     max_attempts = 20 if max_attempts is None else max_attempts
     _delay: int = int(delay) if isinstance(delay, float) else delay
-
     if len(paths) < 1:
         return None
     client_s3: boto3.client = _utils.client(service_name="s3", session=boto3_session)
-    waiter = client_s3.get_waiter(waiter_name)
     _paths: List[Tuple[str, str]] = [_utils.parse_path(path=p) for p in paths]
     if use_threads is False:
+        waiter = client_s3.get_waiter(waiter_name)
         for bucket, key in _paths:
             waiter.wait(Bucket=bucket, Key=key, WaiterConfig={"Delay": _delay, "MaxAttempts": max_attempts})
     else:
         cpus: int = _utils.ensure_cpu_count(use_threads=use_threads)
         with concurrent.futures.ThreadPoolExecutor(max_workers=cpus) as executor:
-            futures: List[concurrent.futures.Future] = []
-            for bucket, key in _paths:
-                future: concurrent.futures.Future = executor.submit(
-                    fn=waiter.wait, Bucket=bucket, Key=key, WaiterConfig={"Delay": _delay, "MaxAttempts": max_attempts}
+            list(
+                executor.map(
+                    _wait_objects_concurrent,
+                    _paths,
+                    repeat(waiter_name),
+                    repeat(client_s3),
+                    repeat(_delay),
+                    repeat(max_attempts),
                 )
-                futures.append(future)
-            for future in futures:
-                future.result()
+            )
     return None
 
 
+def _wait_objects_concurrent(
+    path: Tuple[str, str], waiter_name: str, client_s3: boto3.client, delay: int, max_attempts: int
+) -> None:
+    waiter = client_s3.get_waiter(waiter_name)
+    bucket, key = path
+    waiter.wait(Bucket=bucket, Key=key, WaiterConfig={"Delay": delay, "MaxAttempts": max_attempts})
+
+
 def read_parquet_table(
     table: str,
     database: str,