Skip to content

Aws connection

AWSConnector

AWSConnector - Connects to AWS using Credentials File or IAM Role

Source code in src/super_gradients/common/aws_connection/aws_connector.py 
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
class AWSConnector:
    """
    AWSConnector - Connects to AWS using Credentials File or IAM Role
    """

    @staticmethod
    def __create_boto_3_session(profile_name: str):
        """
        __create_boto_3_session
            :param profile_name:
            :return:
        """
        current_class_name = __class__.__name__
        logger = logging.getLogger(current_class_name)

        try:
            try:
                if profile_name and boto3.session.Session(profile_name=profile_name).get_credentials():
                    # TRY USING A SPECIFIC PROFILE_NAME (USING A CREDENTIALS FILE)
                    logger.info("Trying to connect to AWS using Credentials File with profile_name: " + profile_name)

                    session = boto3.Session(profile_name=profile_name)
                    return session

            except ProfileNotFound as profileNotFoundException:
                logger.debug(
                    "[" + current_class_name + "] - Could not find profile name - Trying using Default Profile/IAM Role" + str(profileNotFoundException)
                )

            # TRY USING AN IAM ROLE (OR *DEFAULT* CREDENTIALS - USING A CREDENTIALS FILE)
            logger.info("Trying to connect to AWS using IAM role or Default Credentials")
            session = boto3.Session()
            return session

        except Exception as ex:
            logger.critical("[" + current_class_name + "] - Caught Exception while trying to connect to AWS Credentials Manager " + str(ex))
            return None

    @staticmethod
    def get_aws_session(profile_name: str) -> boto3.Session:
        """
        get_aws_session - Connects to AWS to retrieve an AWS Session
            :param      profile_name: The Config Profile (Environment Name in Credentials file)
            :return:    boto3 Session
        """
        current_class_name = __class__.__name__
        logger = logging.getLogger(current_class_name)

        aws_session = AWSConnector.__create_boto_3_session(profile_name=profile_name)
        if aws_session is None:
            logger.error("Failed to initiate an AWS Session")

        return aws_session

    @staticmethod
    def get_aws_client_for_service_name(profile_name: str, service_name: str) -> boto3.Session.client:
        """
        get_aws_client_for_service_name - Connects to AWS to retrieve the relevant Client
            :param      profile_name: The Config Profile (Environment Name in Credentials file)
            :param      service_name: The AWS Service name to get the Client for
            :return:    Service client instance
        """
        current_class_name = __class__.__name__
        logger = logging.getLogger(current_class_name)

        aws_session = AWSConnector.__create_boto_3_session(profile_name=profile_name)
        if aws_session is None:
            logger.error("Failed to connect to AWS client: " + str(service_name))

        return aws_session.client(service_name=service_name)

    @staticmethod
    def get_aws_resource_for_service_name(profile_name: str, service_name: str) -> boto3.Session.resource:
        """
        Connects to AWS to retrieve the relevant Resource (More functionality then Client)
            :param      profile_name: The Config Profile (Environment Name in Credentials file)
            :param      service_name: The AWS Service name to get the Client for
            :return:    Service client instance
        """
        current_class_name = __class__.__name__
        logger = logging.getLogger(current_class_name)

        aws_session = AWSConnector.__create_boto_3_session(profile_name=profile_name)
        if aws_session is None:
            logger.error("Failed to connect to AWS client: " + str(service_name))

        return aws_session.resource(service_name=service_name)

    @staticmethod
    def is_client_error(code):
        e = sys.exc_info()[1]
        if isinstance(e, ClientError) and e.response["Error"]["Code"] == code:
            return ClientError
        return type("NeverEverRaisedException", (Exception,), {})

__create_boto_3_session(profile_name) staticmethod

__create_boto_3_session :param profile_name: :return:

Source code in src/super_gradients/common/aws_connection/aws_connector.py 
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
@staticmethod
def __create_boto_3_session(profile_name: str):
    """
    __create_boto_3_session
        :param profile_name:
        :return:
    """
    current_class_name = __class__.__name__
    logger = logging.getLogger(current_class_name)

    try:
        try:
            if profile_name and boto3.session.Session(profile_name=profile_name).get_credentials():
                # TRY USING A SPECIFIC PROFILE_NAME (USING A CREDENTIALS FILE)
                logger.info("Trying to connect to AWS using Credentials File with profile_name: " + profile_name)

                session = boto3.Session(profile_name=profile_name)
                return session

        except ProfileNotFound as profileNotFoundException:
            logger.debug(
                "[" + current_class_name + "] - Could not find profile name - Trying using Default Profile/IAM Role" + str(profileNotFoundException)
            )

        # TRY USING AN IAM ROLE (OR *DEFAULT* CREDENTIALS - USING A CREDENTIALS FILE)
        logger.info("Trying to connect to AWS using IAM role or Default Credentials")
        session = boto3.Session()
        return session

    except Exception as ex:
        logger.critical("[" + current_class_name + "] - Caught Exception while trying to connect to AWS Credentials Manager " + str(ex))
        return None

get_aws_client_for_service_name(profile_name, service_name) staticmethod

get_aws_client_for_service_name - Connects to AWS to retrieve the relevant Client :param profile_name: The Config Profile (Environment Name in Credentials file) :param service_name: The AWS Service name to get the Client for :return: Service client instance

Source code in src/super_gradients/common/aws_connection/aws_connector.py 
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
@staticmethod
def get_aws_client_for_service_name(profile_name: str, service_name: str) -> boto3.Session.client:
    """
    get_aws_client_for_service_name - Connects to AWS to retrieve the relevant Client
        :param      profile_name: The Config Profile (Environment Name in Credentials file)
        :param      service_name: The AWS Service name to get the Client for
        :return:    Service client instance
    """
    current_class_name = __class__.__name__
    logger = logging.getLogger(current_class_name)

    aws_session = AWSConnector.__create_boto_3_session(profile_name=profile_name)
    if aws_session is None:
        logger.error("Failed to connect to AWS client: " + str(service_name))

    return aws_session.client(service_name=service_name)

get_aws_resource_for_service_name(profile_name, service_name) staticmethod

Connects to AWS to retrieve the relevant Resource (More functionality then Client) :param profile_name: The Config Profile (Environment Name in Credentials file) :param service_name: The AWS Service name to get the Client for :return: Service client instance

Source code in src/super_gradients/common/aws_connection/aws_connector.py 
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
@staticmethod
def get_aws_resource_for_service_name(profile_name: str, service_name: str) -> boto3.Session.resource:
    """
    Connects to AWS to retrieve the relevant Resource (More functionality then Client)
        :param      profile_name: The Config Profile (Environment Name in Credentials file)
        :param      service_name: The AWS Service name to get the Client for
        :return:    Service client instance
    """
    current_class_name = __class__.__name__
    logger = logging.getLogger(current_class_name)

    aws_session = AWSConnector.__create_boto_3_session(profile_name=profile_name)
    if aws_session is None:
        logger.error("Failed to connect to AWS client: " + str(service_name))

    return aws_session.resource(service_name=service_name)

get_aws_session(profile_name) staticmethod

get_aws_session - Connects to AWS to retrieve an AWS Session :param profile_name: The Config Profile (Environment Name in Credentials file) :return: boto3 Session

Source code in src/super_gradients/common/aws_connection/aws_connector.py 
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
@staticmethod
def get_aws_session(profile_name: str) -> boto3.Session:
    """
    get_aws_session - Connects to AWS to retrieve an AWS Session
        :param      profile_name: The Config Profile (Environment Name in Credentials file)
        :return:    boto3 Session
    """
    current_class_name = __class__.__name__
    logger = logging.getLogger(current_class_name)

    aws_session = AWSConnector.__create_boto_3_session(profile_name=profile_name)
    if aws_session is None:
        logger.error("Failed to initiate an AWS Session")

    return aws_session

AWSSecretsManagerConnector

AWSSecretsManagerConnector - This class handles the AWS Secrets Manager connection

Source code in src/super_gradients/common/aws_connection/aws_secrets_manager_connector.py 
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
class AWSSecretsManagerConnector:
    """
    AWSSecretsManagerConnector - This class handles the AWS Secrets Manager connection
    """

    __slots__ = []  # Making the class immutable for runtime safety
    current_environment_client = None
    DECI_ENVIRONMENTS = ["research", "development", "staging", "production"]

    @staticmethod
    @explicit_params_validation(validation_type="NoneOrEmpty")
    def get_secret_value_for_secret_key(aws_env: str, secret_name: str, secret_key: str) -> str:
        """
        get_secret_value_for_secret_key - Gets a Secret Value from AWS Secrets Manager for the Provided Key
            :param aws_env:         The environment to get the secret for
            :param secret_name: The Secret Name stored in Secrets Manager
            :param secret_key:  The Secret Key To retrieve it's value from AWS
            :return:  str: The Secret Value
        """
        current_class_name = __class__.__name__
        logger = logging.getLogger(current_class_name)
        secret_key = secret_key.upper()
        aws_secrets_dict = AWSSecretsManagerConnector.__get_secrets_manager_dict_for_secret_name(aws_env=aws_env, secret_name=secret_name)

        secret_key = ".".join([aws_env.upper(), secret_key])
        if secret_key not in aws_secrets_dict.keys():
            error = f"[{current_class_name}] - Secret Key ({secret_key}) not Found in AWS Secret: " + secret_name
            logger.error(error)
            raise EnvironmentError(error)
        else:
            return aws_secrets_dict[secret_key]

    @staticmethod
    @explicit_params_validation(validation_type="NoneOrEmpty")
    def get_secret_values_dict_for_secret_key_properties(env: str, secret_key: str, secret_name: str, db_properties_set: set = None) -> dict:
        """
        get_config_dict - Returns the config dict of the properties from the properties dict
            :param  env:                The environment to open the dict for
            :param  secret_key:         The Secret Key
            :param  secret_name:        The Secret to Retrieve to from AWS secrets manager (usually project name)
            :param  db_properties_set:  The set of the properties to get secrets values for
            :return:  dict The secrets dict for the requested property
        """
        current_class_name = __class__.__name__
        logger = logging.getLogger(current_class_name)
        aws_secrets_dict = AWSSecretsManagerConnector.__get_secrets_manager_dict_for_secret_name(aws_env=env, secret_name=secret_name)

        aws_env_safe_secrets = {}
        # FILL THE DICT VALUES FROM THE AWS SECRETS RESPONSE
        if db_properties_set:
            for secret_key_property in db_properties_set:
                secret_key_to_retrieve = ".".join([env.upper(), secret_key, secret_key_property])
                if secret_key_to_retrieve not in aws_secrets_dict:
                    error = (
                        f'[{current_class_name}] - Error retrieving data from AWS Secrets Manager for Secret Key "{secret_name}": '
                        f'The secret property "{secret_key_property}" Does Not Exist'
                    )
                    logger.error(error)
                    raise EnvironmentError(error)
                else:
                    env_stripped_key_name = secret_key_to_retrieve.lstrip(env.upper()).lstrip(".")
                    aws_env_safe_secrets[env_stripped_key_name] = aws_secrets_dict[secret_key_to_retrieve]
        else:
            # "db_properties_set" is not specified - validating and returning all the secret keys and values for
            # the secret name.
            for secret_key_name, secret_value in aws_secrets_dict.items():
                secret_key_to_retrieve = ".".join([env.upper(), secret_key])
                assert secret_key_name.startswith(env.upper()), (
                    f'The secret key property "{secret_key_name}", found in secret named {secret_name}, is not following the convention of '
                    f'environment prefix. please add the environment prefix "{env.upper()}" to property "{secret_key_name}"'
                )
                if secret_key_name.startswith(secret_key_to_retrieve):
                    env_stripped_key_name = secret_key_name.lstrip(env.upper()).lstrip(".")
                    aws_env_safe_secrets[env_stripped_key_name] = secret_value
        return aws_env_safe_secrets

    @staticmethod
    def __get_secrets_manager_dict_for_secret_name(aws_env: str, secret_name: str) -> dict:
        """
        __get_secrets_manager_dict_for_secret_name
            :param  aws_env:                The environment to open the dict for
            :param  secret_name:        The Secret to Retrieve to from AWS secrets manager (usually project name)
            :return: python Dictionary with the key/value pairs stored in AWS Secrets Manager
        """
        current_class_name = __class__.__name__
        logger = logging.getLogger(current_class_name)

        secrets_path = AWSSecretsManagerConnector.__get_secrets_path_from_secret_name(aws_env, secret_name)

        try:
            if not AWSSecretsManagerConnector.current_environment_client:
                logger.debug("Initializing a new secrets manager client...")
                AWSSecretsManagerConnector.current_environment_client = AWSConnector.get_aws_client_for_service_name(
                    profile_name=aws_env, service_name="secretsmanager"
                )
            logger.debug(f'Fetching the secret "{secret_name}" in env "{aws_env}"')
            aws_secrets = AWSSecretsManagerConnector.current_environment_client.get_secret_value(SecretId=secrets_path)
            aws_secrets_dict = json.loads(aws_secrets["SecretString"])
            return aws_secrets_dict

        except Exception as ex:
            error = (
                f'[{current_class_name}] - Caught Exception while trying to connect to aws to get credentials from secrets manager: "{ex}" for {secrets_path}'
            )
            logger.error(error)
            raise EnvironmentError(error)

    @staticmethod
    def __get_secrets_path_from_secret_name(aws_env: str, secret_name: str) -> str:
        """
        __get_secrets_path_from_secret_name - Extracts the full secret path based on the Environment
        :param aws_env:         Env
        :param secret_name: Secret Name
        :return:     str:   The full secret path
        """
        current_class_name = __class__.__name__
        logger = logging.getLogger(current_class_name)

        # Checking for lowercase exact match, in order to prevent any implicit usage of the environments.
        if aws_env not in AWSSecretsManagerConnector.DECI_ENVIRONMENTS:
            logger.critical("[" + current_class_name + " ] -  wrong environment param... Exiting")
            raise Exception("[" + current_class_name + "] - wrong environment param")

        secrets_path = "/".join([aws_env, secret_name])
        return secrets_path

__get_secrets_manager_dict_for_secret_name(aws_env, secret_name) staticmethod

__get_secrets_manager_dict_for_secret_name :param aws_env: The environment to open the dict for :param secret_name: The Secret to Retrieve to from AWS secrets manager (usually project name) :return: python Dictionary with the key/value pairs stored in AWS Secrets Manager

Source code in src/super_gradients/common/aws_connection/aws_secrets_manager_connector.py 
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
@staticmethod
def __get_secrets_manager_dict_for_secret_name(aws_env: str, secret_name: str) -> dict:
    """
    __get_secrets_manager_dict_for_secret_name
        :param  aws_env:                The environment to open the dict for
        :param  secret_name:        The Secret to Retrieve to from AWS secrets manager (usually project name)
        :return: python Dictionary with the key/value pairs stored in AWS Secrets Manager
    """
    current_class_name = __class__.__name__
    logger = logging.getLogger(current_class_name)

    secrets_path = AWSSecretsManagerConnector.__get_secrets_path_from_secret_name(aws_env, secret_name)

    try:
        if not AWSSecretsManagerConnector.current_environment_client:
            logger.debug("Initializing a new secrets manager client...")
            AWSSecretsManagerConnector.current_environment_client = AWSConnector.get_aws_client_for_service_name(
                profile_name=aws_env, service_name="secretsmanager"
            )
        logger.debug(f'Fetching the secret "{secret_name}" in env "{aws_env}"')
        aws_secrets = AWSSecretsManagerConnector.current_environment_client.get_secret_value(SecretId=secrets_path)
        aws_secrets_dict = json.loads(aws_secrets["SecretString"])
        return aws_secrets_dict

    except Exception as ex:
        error = (
            f'[{current_class_name}] - Caught Exception while trying to connect to aws to get credentials from secrets manager: "{ex}" for {secrets_path}'
        )
        logger.error(error)
        raise EnvironmentError(error)

__get_secrets_path_from_secret_name(aws_env, secret_name) staticmethod

__get_secrets_path_from_secret_name - Extracts the full secret path based on the Environment

Parameters:

Name Type Description Default
aws_env str

Env

 required
secret_name str

Secret Name

 required

Returns:

Type Description
str

str: The full secret path
Source code in src/super_gradients/common/aws_connection/aws_secrets_manager_connector.py 
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
@staticmethod
def __get_secrets_path_from_secret_name(aws_env: str, secret_name: str) -> str:
    """
    __get_secrets_path_from_secret_name - Extracts the full secret path based on the Environment
    :param aws_env:         Env
    :param secret_name: Secret Name
    :return:     str:   The full secret path
    """
    current_class_name = __class__.__name__
    logger = logging.getLogger(current_class_name)

    # Checking for lowercase exact match, in order to prevent any implicit usage of the environments.
    if aws_env not in AWSSecretsManagerConnector.DECI_ENVIRONMENTS:
        logger.critical("[" + current_class_name + " ] -  wrong environment param... Exiting")
        raise Exception("[" + current_class_name + "] - wrong environment param")

    secrets_path = "/".join([aws_env, secret_name])
    return secrets_path

get_secret_value_for_secret_key(aws_env, secret_name, secret_key) staticmethod

get_secret_value_for_secret_key - Gets a Secret Value from AWS Secrets Manager for the Provided Key :param aws_env: The environment to get the secret for :param secret_name: The Secret Name stored in Secrets Manager :param secret_key: The Secret Key To retrieve it's value from AWS :return: str: The Secret Value

Source code in src/super_gradients/common/aws_connection/aws_secrets_manager_connector.py 
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
@staticmethod
@explicit_params_validation(validation_type="NoneOrEmpty")
def get_secret_value_for_secret_key(aws_env: str, secret_name: str, secret_key: str) -> str:
    """
    get_secret_value_for_secret_key - Gets a Secret Value from AWS Secrets Manager for the Provided Key
        :param aws_env:         The environment to get the secret for
        :param secret_name: The Secret Name stored in Secrets Manager
        :param secret_key:  The Secret Key To retrieve it's value from AWS
        :return:  str: The Secret Value
    """
    current_class_name = __class__.__name__
    logger = logging.getLogger(current_class_name)
    secret_key = secret_key.upper()
    aws_secrets_dict = AWSSecretsManagerConnector.__get_secrets_manager_dict_for_secret_name(aws_env=aws_env, secret_name=secret_name)

    secret_key = ".".join([aws_env.upper(), secret_key])
    if secret_key not in aws_secrets_dict.keys():
        error = f"[{current_class_name}] - Secret Key ({secret_key}) not Found in AWS Secret: " + secret_name
        logger.error(error)
        raise EnvironmentError(error)
    else:
        return aws_secrets_dict[secret_key]

get_secret_values_dict_for_secret_key_properties(env, secret_key, secret_name, db_properties_set=None) staticmethod

get_config_dict - Returns the config dict of the properties from the properties dict :param env: The environment to open the dict for :param secret_key: The Secret Key :param secret_name: The Secret to Retrieve to from AWS secrets manager (usually project name) :param db_properties_set: The set of the properties to get secrets values for :return: dict The secrets dict for the requested property

Source code in src/super_gradients/common/aws_connection/aws_secrets_manager_connector.py 
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
@staticmethod
@explicit_params_validation(validation_type="NoneOrEmpty")
def get_secret_values_dict_for_secret_key_properties(env: str, secret_key: str, secret_name: str, db_properties_set: set = None) -> dict:
    """
    get_config_dict - Returns the config dict of the properties from the properties dict
        :param  env:                The environment to open the dict for
        :param  secret_key:         The Secret Key
        :param  secret_name:        The Secret to Retrieve to from AWS secrets manager (usually project name)
        :param  db_properties_set:  The set of the properties to get secrets values for
        :return:  dict The secrets dict for the requested property
    """
    current_class_name = __class__.__name__
    logger = logging.getLogger(current_class_name)
    aws_secrets_dict = AWSSecretsManagerConnector.__get_secrets_manager_dict_for_secret_name(aws_env=env, secret_name=secret_name)

    aws_env_safe_secrets = {}
    # FILL THE DICT VALUES FROM THE AWS SECRETS RESPONSE
    if db_properties_set:
        for secret_key_property in db_properties_set:
            secret_key_to_retrieve = ".".join([env.upper(), secret_key, secret_key_property])
            if secret_key_to_retrieve not in aws_secrets_dict:
                error = (
                    f'[{current_class_name}] - Error retrieving data from AWS Secrets Manager for Secret Key "{secret_name}": '
                    f'The secret property "{secret_key_property}" Does Not Exist'
                )
                logger.error(error)
                raise EnvironmentError(error)
            else:
                env_stripped_key_name = secret_key_to_retrieve.lstrip(env.upper()).lstrip(".")
                aws_env_safe_secrets[env_stripped_key_name] = aws_secrets_dict[secret_key_to_retrieve]
    else:
        # "db_properties_set" is not specified - validating and returning all the secret keys and values for
        # the secret name.
        for secret_key_name, secret_value in aws_secrets_dict.items():
            secret_key_to_retrieve = ".".join([env.upper(), secret_key])
            assert secret_key_name.startswith(env.upper()), (
                f'The secret key property "{secret_key_name}", found in secret named {secret_name}, is not following the convention of '
                f'environment prefix. please add the environment prefix "{env.upper()}" to property "{secret_key_name}"'
            )
            if secret_key_name.startswith(secret_key_to_retrieve):
                env_stripped_key_name = secret_key_name.lstrip(env.upper()).lstrip(".")
                aws_env_safe_secrets[env_stripped_key_name] = secret_value
    return aws_env_safe_secrets