Skip to content

Commit

Permalink
refactored get_s3_client into a function in a aws_s3_utils.py file
Browse files Browse the repository at this point in the history
* removed some bulk from the `cript.API` class
* wrote docstrings for `get_s3_client`
* tested the code and it is working fine
* formatted the whole thing with black
* changed log to have `file` be capitalized for better UI
  • Loading branch information
nh916 committed Sep 20, 2023
1 parent 4b4c3ff commit d371710
Show file tree
Hide file tree
Showing 2 changed files with 48 additions and 17 deletions.
25 changes: 8 additions & 17 deletions src/cript/api/api.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
from cript.api.paginator import Paginator
from cript.api.utils.get_host_token import resolve_host_and_token
from cript.api.utils.helper_functions import _get_node_type_from_json
from cript.api.utils.aws_s3_utils import get_s3_client
from cript.api.utils.save_helper import (
_fix_node_save,
_get_uuid_from_error_message,
Expand Down Expand Up @@ -347,29 +348,19 @@ def _prepare_host(self, host: str) -> str:
@property
def _s3_client(self) -> boto3.client: # type: ignore
"""
creates or returns a fully authenticated and ready s3 client
Property to use when wanting to interact with AWS S3.
Gets a fully authenticated AWS S3 client if it was never created and stash it,
if the AWS S3 client has been created before, then returns the client that it has
Returns
-------
s3_client: boto3.client
fully prepared and authenticated s3 client ready to be used throughout the script
"""

if self._internal_s3_client is None:
auth = boto3.client("cognito-identity", region_name=self._REGION_NAME)
identity_id = auth.get_id(IdentityPoolId=self._IDENTITY_POOL_ID, Logins={self._COGNITO_LOGIN_PROVIDER: self._storage_token})
# TODO remove this temporary fix to the token, by getting is from back end.
aws_token = self._storage_token

aws_credentials = auth.get_credentials_for_identity(IdentityId=identity_id["IdentityId"], Logins={self._COGNITO_LOGIN_PROVIDER: aws_token})
aws_credentials = aws_credentials["Credentials"]
s3_client = boto3.client(
"s3",
aws_access_key_id=aws_credentials["AccessKeyId"],
aws_secret_access_key=aws_credentials["SecretKey"],
aws_session_token=aws_credentials["SessionToken"],
)
self._internal_s3_client = s3_client
self._internal_s3_client = get_s3_client(region_name=self._REGION_NAME, identity_pool_id=self._IDENTITY_POOL_ID, cognito_login_provider=self._COGNITO_LOGIN_PROVIDER, storage_token=self._storage_token)

return self._internal_s3_client

def __enter__(self):
Expand Down Expand Up @@ -869,7 +860,7 @@ def upload_file(self, file_path: Union[Path, str]) -> str:
# upload file to AWS S3
self._s3_client.upload_file(Filename=file_path, Bucket=self._BUCKET_NAME, Key=object_name) # type: ignore

self.logger.info(f"Uploaded file: '{file_path}' to CRIPT storage")
self.logger.info(f"Uploaded File: '{file_path}' to CRIPT storage")

# return the object_name within AWS S3 for easy retrieval
return object_name
Expand Down
40 changes: 40 additions & 0 deletions src/cript/api/utils/aws_s3_utils.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,40 @@
from typing import Any

import boto3
from beartype import beartype


@beartype
def get_s3_client(region_name: str, identity_pool_id: str, cognito_login_provider: str, storage_token: str) -> Any:
"""
Creates an AWS S3 client and returns it to be used in the `cript.API` class.
Parameters
----------
region_name: str
AWS S3 region name
identity_pool_id: str
AWS S3 identity pool id
cognito_login_provider: str
AWS S3 cognito login provider
storage_token: str
AWS S3 storage token gotten from the CRIPT frontend
Returns
-------
boto3.client
fully working AWS S3 client
"""
auth = boto3.client("cognito-identity", region_name=region_name)
identity_id = auth.get_id(IdentityPoolId=identity_pool_id, Logins={cognito_login_provider: storage_token})
aws_token = storage_token

aws_credentials = auth.get_credentials_for_identity(IdentityId=identity_id["IdentityId"], Logins={cognito_login_provider: aws_token})
aws_credentials = aws_credentials["Credentials"]
s3_client = boto3.client(
"s3",
aws_access_key_id=aws_credentials["AccessKeyId"],
aws_secret_access_key=aws_credentials["SecretKey"],
aws_session_token=aws_credentials["SessionToken"],
)
return s3_client

0 comments on commit d371710

Please sign in to comment.