def _get_boto3_session()

in sagemaker_studio_sparkmagic_lib/emr.py [0:0]

• 25 lines of code
• 3 McCabe index (conditional complexity)

    def _get_boto3_session(self, role_arn):
        """
        Allows cross account cluster access
        """
        if not role_arn:
            return boto3.session.Session()
        else:
            logger.info(f"Assuming role: {role_arn}")
            sts_client = boto3.client("sts")
            try:
                assume_role_object = sts_client.assume_role(
                    RoleArn=role_arn, RoleSessionName="SageMakerStudioUser"
                )
            except botocore.exceptions.ClientError as ce:
                logger.debug(
                    f"Failed to assume role: ({role_arn}) details. {ce.response}"
                )
                raise ValueError(
                    f"Unable to assume role(arn: {role_arn}). Ensure permissions are setup correctly"
                    f' Error: {ce.response["Error"]}'
                ) from None
            return boto3.Session(
                aws_access_key_id=assume_role_object["Credentials"]["AccessKeyId"],
                aws_secret_access_key=assume_role_object["Credentials"][
                    "SecretAccessKey"
                ],
                aws_session_token=assume_role_object["Credentials"]["SessionToken"],
            )