diff --git a/sdt_dask/dataplugs/S3Bucket_plug.py b/sdt_dask/dataplugs/S3Bucket_plug.py index e5fde285..6dcf6757 100644 --- a/sdt_dask/dataplugs/S3Bucket_plug.py +++ b/sdt_dask/dataplugs/S3Bucket_plug.py @@ -17,10 +17,15 @@ def __init__(self, bucket_name): """ self.bucket_name = bucket_name + def _create_s3_client(self): + # Creating a new session for each call to ensure thread safety + session = boto3.session.Session() + s3_client = session.client('s3') + return s3_client + def _pull_data(self, key): - - s3_client = boto3.client('s3') - print(f"Loading file from S3 bucket: {key}...") + + s3_client = self._create_s3_client() obj = s3_client.get_object(Bucket=self.bucket_name, Key=key) # Assume file is CSV @@ -58,7 +63,7 @@ def _pull_keys(self) -> list: """ KEYS = [] - s3_client = boto3.client('s3') + s3_client = self._create_s3_client() objects = s3_client.list_objects_v2(Bucket=self.bucket_name) if 'Contents' in objects: