Skip to content

Commit

Permalink
S3 Storage Client Authentication Bug fix (#445)
Browse files Browse the repository at this point in the history
  • Loading branch information
Andre-Lx-Costa authored Jul 4, 2024
1 parent f145b05 commit 4f44376
Show file tree
Hide file tree
Showing 3 changed files with 25 additions and 18 deletions.
4 changes: 2 additions & 2 deletions adapta/security/clients/aws/_aws_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
"""
Expand Down Expand Up @@ -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.
Expand Down
19 changes: 13 additions & 6 deletions adapta/storage/blob/s3_storage_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
20 changes: 10 additions & 10 deletions poetry.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

0 comments on commit 4f44376

Please sign in to comment.