From 4f44376256e04ee8b4cb6f364f446b9f71e00255 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Andr=C3=A9=20Costa?= <ancosta@gmail.com>
Date: Thu, 4 Jul 2024 10:20:29 +0200
Subject: [PATCH] S3 Storage Client Authentication Bug fix (#445)

---
 adapta/security/clients/aws/_aws_client.py |  4 ++--
 adapta/storage/blob/s3_storage_client.py   | 19 +++++++++++++------
 poetry.lock                                | 20 ++++++++++----------
 3 files changed, 25 insertions(+), 18 deletions(-)

diff --git a/adapta/security/clients/aws/_aws_client.py b/adapta/security/clients/aws/_aws_client.py
index 126fb839..acab751a 100644
--- a/adapta/security/clients/aws/_aws_client.py
+++ b/adapta/security/clients/aws/_aws_client.py
@@ -54,7 +54,7 @@ def from_base_client(cls, client: AuthenticationClient) -> Optional["AwsClient"]
         """
         return client if isinstance(client, AwsClient) else None
 
-    def get_credentials(self):
+    def get_credentials(self) -> Optional[AccessKeyCredentials]:
         """
         Returns configured credentials (if any)
         """
@@ -89,7 +89,7 @@ def get_pyarrow_filesystem(self, path: DataPath, connection_options: Optional[Di
         :return:
         """
 
-    def initialize_session(self, session_callable: Optional[Callable[[], None]] = None) -> "AwsClient":
+    def initialize_session(self, session_callable: Optional[Callable[[], Session]] = None) -> "AwsClient":
         """
         Initializes the session by custom session function or a default one if no function is provided."
         :return: AwsClient with established session.
diff --git a/adapta/storage/blob/s3_storage_client.py b/adapta/storage/blob/s3_storage_client.py
index 0b944abb..0671bc9a 100644
--- a/adapta/storage/blob/s3_storage_client.py
+++ b/adapta/storage/blob/s3_storage_client.py
@@ -46,15 +46,22 @@ def __init__(self, *, base_client: AwsClient, s3_resource: Optional[Session] = N
 
     @classmethod
     def create(
-        cls, auth: AwsClient, endpoint_url: Optional[str] = None, session_callable: Optional[Callable[[], None]] = None
+        cls,
+        auth: AwsClient,
+        endpoint_url: Optional[str] = None,
+        session_callable: Optional[Callable[[], Session]] = None,
     ):
-        auth.initialize_session(session_callable)
+        def _get_endpoint_url() -> Optional[str]:
+            if endpoint_url:
+                return endpoint_url
+            if auth.get_credentials():
+                return auth.get_credentials().endpoint
+
+            return None
 
-        s3_resource = auth.session.resource(
-            "s3", endpoint_url=endpoint_url if endpoint_url is not None else auth.get_credentials().endpoint
-        )
+        auth.initialize_session(session_callable)
 
-        return cls(base_client=auth, s3_resource=s3_resource)
+        return cls(base_client=auth, s3_resource=auth.session.resource("s3", endpoint_url=_get_endpoint_url()))
 
     def get_blob_uri(self, blob_path: DataPath, **kwargs) -> str:
         """Returns a signed URL for a blob in S3 storage.
diff --git a/poetry.lock b/poetry.lock
index b4a95bea..76e9edb3 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -397,17 +397,17 @@ uvloop = ["uvloop (>=0.15.2)"]
 
 [[package]]
 name = "boto3"
-version = "1.34.138"
+version = "1.34.139"
 description = "The AWS SDK for Python"
 optional = true
 python-versions = ">=3.8"
 files = [
-    {file = "boto3-1.34.138-py3-none-any.whl", hash = "sha256:81518aa95fad71279411fb5c94da4b4a554a5d53fc876faca62b7b5c8737f1cb"},
-    {file = "boto3-1.34.138.tar.gz", hash = "sha256:f79c15e33eb7706f197d98d828b193cf0891966682ad3ec5e900f6f9e7362e35"},
+    {file = "boto3-1.34.139-py3-none-any.whl", hash = "sha256:98b2a12bcb30e679fa9f60fc74145a39db5ec2ca7b7c763f42896e3bd9b3a38d"},
+    {file = "boto3-1.34.139.tar.gz", hash = "sha256:32b99f0d76ec81fdca287ace2c9744a2eb8b92cb62bf4d26d52a4f516b63a6bf"},
 ]
 
 [package.dependencies]
-botocore = ">=1.34.138,<1.35.0"
+botocore = ">=1.34.139,<1.35.0"
 jmespath = ">=0.7.1,<2.0.0"
 s3transfer = ">=0.10.0,<0.11.0"
 
@@ -416,13 +416,13 @@ crt = ["botocore[crt] (>=1.21.0,<2.0a0)"]
 
 [[package]]
 name = "botocore"
-version = "1.34.138"
+version = "1.34.139"
 description = "Low-level, data-driven core of boto 3."
 optional = true
 python-versions = ">=3.8"
 files = [
-    {file = "botocore-1.34.138-py3-none-any.whl", hash = "sha256:84e96a954c39a6f09cae4ea95b2ae582b5ae01b5040c92507b60509c9be5377a"},
-    {file = "botocore-1.34.138.tar.gz", hash = "sha256:f558bbea96c4a4abbaeeedc477dabb00902311ba1ca6327974a6819b9f384920"},
+    {file = "botocore-1.34.139-py3-none-any.whl", hash = "sha256:dd1e085d4caa2a4c1b7d83e3bc51416111c8238a35d498e9d3b04f3b63b086ba"},
+    {file = "botocore-1.34.139.tar.gz", hash = "sha256:df023d8cf8999d574214dad4645cb90f9d2ccd1494f6ee2b57b1ab7522f6be77"},
 ]
 
 [package.dependencies]
@@ -542,13 +542,13 @@ graph = ["gremlinpython (==3.4.6)"]
 
 [[package]]
 name = "certifi"
-version = "2024.6.2"
+version = "2024.7.4"
 description = "Python package for providing Mozilla's CA Bundle."
 optional = false
 python-versions = ">=3.6"
 files = [
-    {file = "certifi-2024.6.2-py3-none-any.whl", hash = "sha256:ddc6c8ce995e6987e7faf5e3f1b02b302836a0e5d98ece18392cb1a36c72ad56"},
-    {file = "certifi-2024.6.2.tar.gz", hash = "sha256:3cd43f1c6fa7dedc5899d69d3ad0398fd018ad1a17fba83ddaf78aa46c747516"},
+    {file = "certifi-2024.7.4-py3-none-any.whl", hash = "sha256:c198e21b1289c2ab85ee4e67bb4b4ef3ead0892059901a8d5b622f24a1101e90"},
+    {file = "certifi-2024.7.4.tar.gz", hash = "sha256:5a1e7645bc0ec61a09e26c36f6106dd4cf40c6db3a1fb6352b0244e7fb057c7b"},
 ]
 
 [[package]]