def lambda_handler()

in customer-stack/create-resources.py [0:0]


def lambda_handler(event, context):

    # Get variables from os
    api_gateway_url = os.environ['ApiGatewayURL']
    api_gateway_role_arn = os.environ['ApiGatewayRoleARN']
    api_gateway_role_name = os.environ['ApiGatewayRoleName']
    auto_ml_role_arn = os.environ['AutoMLRoleARN']
    auto_ml_role_name = os.environ['AutoMLRoleName']
    region_name = os.environ['Region']
    s3_bucket_name = os.environ['S3BucketName']
    secret_name = os.environ['SecretArn']
    kms_key_arn = os.environ['KmsKeyArn']
    vpc_security_group_ids = os.environ['VpcSecurityGroupIds']
    vpc_subnet_ids = os.environ['VpcSubnetIds']
    snowflake_role_name = os.environ['SnowflakeRole']
    stack_name = os.environ['StackName']
    database_name = os.environ['DatabaseName']
    schema_name = os.environ['SchemaName']

    logger.info("api_gateway_url: " + api_gateway_url)
    logger.info("api_gateway_role_arn: " + api_gateway_role_arn)
    logger.info("api_gateway_role_name: " + api_gateway_role_name)
    logger.info("auto_ml_role_arn: " + auto_ml_role_arn)
    logger.info("auto_ml_role_name: " + auto_ml_role_name)
    logger.info("region_name: " + region_name)
    logger.info("s3_bucket_name: " + s3_bucket_name)
    logger.info("secret_name: " + secret_name)
    logger.info("kms_key_arn: " + kms_key_arn)
    logger.info("vpc_security_group_ids: " + vpc_security_group_ids)
    logger.info("vpc_subnet_ids: " + vpc_subnet_ids)
    logger.info("snowflake_role_name: " + snowflake_role_name)
    logger.info("stack_name: " + stack_name)
    logger.info("database_name: " + database_name)
    logger.info("schema_name: " + schema_name)
    logger.info("Snowflake resource suffix: " + os.environ['SnowflakeResourceSuffix'])

    # Delete
    if event['RequestType'] == 'Delete':
        logger.info("No action for Delete. Exiting.")
        sendResponse(event, context, SUCCESS, EMPTY_RESPONSE_DATA)
        return

    # Get the information connection from Secrets Manager
    try:
        get_secret_value_response = get_secret_information(region_name, secret_name)
    except:
        sendResponse(event, context, FAILED, EMPTY_RESPONSE_DATA)
        return

    # Decrypted secret using the associated KMS CMK
    # Ensure the Secret is in String mode
    if 'SecretString' not in get_secret_value_response:
        logger.error("The Secret is not in String mode")
        sendResponse(event, context, FAILED, EMPTY_RESPONSE_DATA)
        return

    # Create Snowflake resource
    try:
        snowflake_connection = connect_to_snowflake(get_secret_value_response, snowflake_role_name)
        snowflake_cursor = snowflake_connection.cursor()

        snowflake_cursor.execute(("use database %s;") % (database_name))
        
        snowflake_cursor.execute(("use schema %s;") % (schema_name))

        storage_integration_name = "AWS_AUTOPILOT_STORAGE_INTEGRATION" + "_" + stack_name
        api_integration_name = "AWS_AUTOPILOT_API_INTEGRATION" + "_" + stack_name

        # Create Snowflake Integrations
        create_storage_integration(snowflake_cursor, storage_integration_name, auto_ml_role_arn, s3_bucket_name)
        create_api_integration(snowflake_cursor, api_integration_name, api_gateway_role_arn, api_gateway_url)
        create_external_functions(snowflake_cursor, api_integration_name, auto_ml_role_arn, api_gateway_url,
                                  s3_bucket_name, secret_name, storage_integration_name, snowflake_role_name,
                                  kms_key_arn, vpc_security_group_ids, vpc_subnet_ids)

        # Describe Snowflake integrations
        storage_integration_info = get_storage_integration_info_for_policy(snowflake_cursor, storage_integration_name)
        api_integration_info = get_api_integration_info_for_policy(snowflake_cursor, api_integration_name)
    except Exception as e:
        logger.exception('Problem running SQL statements: ' + str(e))
        responseData = {'Failed': 'Unable to execute SQL statements in Snowflake'}
        sendResponse(event, context, FAILED, responseData)
        return
    finally:
        if 'snowflake_cursor' in vars():
            snowflake_cursor.close()
        if 'snowflake_connection' in vars():
            snowflake_connection.close()

    # Update IAM role to add Snowflake information
    logger.info("Updating IAM Role")
    storage_integration_policy_str = create_policy_string(storage_integration_info)
    api_integration_policy_str = create_policy_string(api_integration_info)

    try:
        update_assume_role_policy(storage_integration_policy_str, auto_ml_role_name)
        update_assume_role_policy(api_integration_policy_str, api_gateway_role_name)
    except Exception as e:
        logger.exception('Problem updating assume role policy: ' + str(e))
        responseData = {'Failed': 'There was a problem updating the assume role policies'}
        sendResponse(event, context, FAILED, responseData)
        return

    responseData = {'Success': 'Snowflake resources created.'}
    sendResponse(event, context, SUCCESS, responseData)
    logger.info("Success")