Skip to content

Commit

Permalink
changing logic on when we try to assume the sageworks role
Browse files Browse the repository at this point in the history
  • Loading branch information
brifordwylie committed Dec 9, 2024
1 parent 95ec5d5 commit 3efebe3
Showing 1 changed file with 10 additions and 8 deletions.
18 changes: 10 additions & 8 deletions src/sageworks/core/cloud_platform/aws/aws_session.py
Original file line number Diff line number Diff line change
Expand Up @@ -44,17 +44,19 @@ def __init__(self):
@property
def boto3_session(self):
"""Get the AWS Boto3 Session, defaulting to the SageWorks Role if possible."""

# Check the execution environment and determine if we need to assume the SageWorks Role
if running_on_lambda() or running_on_glue() or self.is_sageworks_role():
self.log.important(f"Using the default Boto3 session...")
return boto3.Session()

# Okay, so we need to assume the SageWorks Role
try:
return self._sageworks_role_boto3_session()
except Exception as e:
self.log.info("Checking Execution Environment...")
if running_on_lambda() or running_on_glue() or self.is_sageworks_role():
self.log.important(f"Using the default Boto3 session: {e}")
return boto3.Session()
else:
msg = "SageWorks Session Failure: Check AWS_PROFILE and/or Renew SSO Token.."
self.log.critical(msg)
raise RuntimeError(msg) from e
msg = "Failed to Assume SageWorks Role: Check AWS_PROFILE and/or Renew SSO Token.."
self.log.critical(msg)
raise RuntimeError(msg) from e

def is_sageworks_role(self) -> bool:
"""Helper: Check if the current AWS Identity is the SageWorks Role"""
Expand Down

0 comments on commit 3efebe3

Please sign in to comment.