From 3efebe3eec0bb30e82664f09e3f28b461ef31b3c Mon Sep 17 00:00:00 2001 From: Brian Wylie Date: Mon, 9 Dec 2024 10:40:26 -0700 Subject: [PATCH] changing logic on when we try to assume the sageworks role --- .../core/cloud_platform/aws/aws_session.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/src/sageworks/core/cloud_platform/aws/aws_session.py b/src/sageworks/core/cloud_platform/aws/aws_session.py index e2d86df60..8b6a3370e 100644 --- a/src/sageworks/core/cloud_platform/aws/aws_session.py +++ b/src/sageworks/core/cloud_platform/aws/aws_session.py @@ -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"""