From 7c01cd4145dadde2368eda6867a8450c08872491 Mon Sep 17 00:00:00 2001 From: Brian Wylie Date: Fri, 29 Nov 2024 15:21:40 -0700 Subject: [PATCH] adding a color print since we can't use logs and levels --- .../core/cloud_platform/aws/aws_session.py | 25 +++++++++++++------ 1 file changed, 17 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 218025bc5..10075889b 100644 --- a/src/sageworks/core/cloud_platform/aws/aws_session.py +++ b/src/sageworks/core/cloud_platform/aws/aws_session.py @@ -106,25 +106,34 @@ def _assume_sageworks_role_session_credentials(self): } # Note: We can't use a log messages, because they go through the CloudWatchHandler # which require AWS credentials so results in a refresh_lock() deadlock. - print(f"AWS Credentials Refreshed: Expires at {response['Expiration'].astimezone()}") + self.c_print(f"AWS Credentials Refreshed: Expires at {response['Expiration'].astimezone()}") return credentials except Exception as e: # Note: We can't use a log message, because they go through the CloudWatchHandler # which require AWS credentials so results in a refresh_lock() deadlock. - print(f"Error during Refresh Credentials: {e}") + self.c_print(f"Error during Refresh Credentials: {e}", critical=True) raise + @staticmethod + def c_print(text: str, critical: bool = False): + if critical: + print(f"\x1b[38;5;198m{text}\x1b[0m") + else: + print(f"\x1b[38;5;69m{text}\x1b[0m") + if __name__ == "__main__": """Exercise the AWS Session Class""" - my_session = AWSSession().boto3_session - my_sts_client = my_session.client("sts") + my_aws_session = AWSSession() + my_boto_session = AWSSession().boto3_session + my_sts_client = my_boto_session.client("sts") print(f"Account ID: {my_sts_client.get_caller_identity()['Account']}") - print(f"Region: {my_session.region_name}") + print(f"Region: {my_boto_session.region_name}") print(f"Identity: {my_sts_client.get_caller_identity()['Arn']}") - print(f"Session Token: {my_session.get_credentials().token}") - print(f"Session Expiry: {my_session.get_credentials()._expiry_time}") - print(f"Assumed Role: {my_session.get_credentials().method}") + print(f"Session Token: {my_boto_session.get_credentials().token}") + print(f"Session Expiry: {my_boto_session.get_credentials()._expiry_time}") + my_aws_session.c_print(f"Assumed Role: {my_boto_session.get_credentials().method}") + my_aws_session.c_print(f"Fake Critical Message", critical=True)