Question about python code structure, functions and calling functions

Question

I tried to create a function that does the following things. (dont know if its worth mentioning but this function is invoked by another function)

1. connects to aws resources using boto3
2. gets a the number of messages available in an sqs queue
3. counts the number of ec2 instances
4. evaluates a set of conditions based on the sqs queue and ec2 instances and either do nothing or write to an sns topic.

Basically i want to publish to a message to sns topic every time the sqs queue is high and the number of ec2 instances which are digesting these, is low.

import os
import boto3
import logging
import types

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

# Create session and clients
sts_client = boto3.client('sts')
sqs_client = boto3.client('sqs')
ec2_client = boto3.client('ec2')
sns_client = boto3.client('sns')

# Call the assume_role method of the STSConnection object and pass the role ARN and a role session name.
assumed_role_object = sts_client.assume_role(
    RoleArn=os.environ['ROLE_ARN'],
    RoleSessionName="AssumeRoleFromCloudOperations"
)

# From the response that contains the assumed role, get the temporary credentials
credentials = assumed_role_object['Credentials']

assumed_role_session = boto3.Session(
    aws_access_key_id=credentials['AccessKeyId'],
    aws_secret_access_key=credentials['SecretAccessKey'],
    aws_session_token=credentials['SessionToken']
)


# Check the queue size
def sqs():
    queue_size = sqs_client.get_queue_attributes(
        QueueUrl=os.environ['SQS_QUEUE_URL'],
        AttributeNames=['ApproximateNumberOfMessages']
    )
    messages = int(queue_size["Attributes"]["ApproximateNumberOfMessages"])
    return messages


# Count the number of active ec2 instances
def count_instances(ec2):
    total_instances = 0
    instances = ec2.instances.filter(Filters=[
        {
            'Instance State': 'instance-state-name',
            'Values': ['running'],
            'Name': 'tag:Name',
            'Values': ['NameOfInstance']
        },
    ])
    for _ in instances:
        total_instances += 1
    return total_instances

    print(f"Total number of active scan servers is: {total_instances}")


# Define the SNS Topic which will be integrated with OpsGenie
def sns():
    topicArn = os.environ['SNS_ARN']


# Evaluate the set of conditions
def evaluate_conditions(context, event):
    sqs()
    if messages > int(os.environ['AVG_QUEUE_SIZE']) and count_instances.total_instances > int(os.environ['AVG_NR_OF_EC2_SCAN_SERVERS']):
        print('False alert')
        logger.info()
    elif messages < int(os.environ['AVG_QUEUE_SIZE']) and count_instances.total_instances < int(os.environ['AVG_NR_OF_EC2_SCAN_SERVERS']):
        print('False alert')
        logger.info()
    elif messages < int(os.environ['AVG_QUEUE_SIZE']) and count_instances.total_instances > int(os.environ['AVG_NR_OF_EC2_SCAN_SERVERS']):
        print('False alert')
        logger.info()
    else:
        sns.publish(TopicArn=os.environ['SNS_ARN'],
                    Message='sameple message',
                    Subject='sample subject')
        print("Published to SNS Topic")

the handler is handler.evaluate_conditions

My question is how can i have some structure in this lambda function? When i run the function i get a naming error:

{
  "errorMessage": "name 'messages' is not defined",
  "errorType": "NameError",
  "stackTrace": [
    "  File \"/var/task/mdc_alert/handler.py\", line 67, in evaluate_conditions\n    if messages > int(os.environ['AVG_QUEUE_SIZE']) and count_instances.total_instances > int(\n"
  ]
}

So it seems that i cannot use the message variable in the evaluate_conditions() function. How can i make the "message" and "total_instances" variables usable in the evaluate_conditions() function?

I've written this function completely based on google searches, stackoverflow and boto3 docs since i don't have any experience with programming.

Is this structure any good, or does it need a complete overhaul? Do i need to change the order of the functions, or maybe create a class?

Answer 1

The immediate issue is that the messages variable is not defined. Your sqs function returns a value but since you're calling it in a void context you're not actually doing anything with that value. You can fix this by changing this line:

sqs()

to this one:

messages = sqs()

I also see some issues with the count_instances function. It's expecting to receive an ec2 variable but you're calling it incorrectly from evaluate_conditions. You could either pass it the ec2_client variable or just use the ec2_client variable directly from within the function.

I suggest renaming your functions to more accurately reflect their return values:

sqs -> sqs_msg_count
count_instances -> running_ec2_count

Making these changes will allow you to refactor evaluate_conditions to shorten the if-then lines, making your code overall easier to read and follow. If you took all these suggestions into account, your code might look something like this:

# Check the queue size
def sqs_msg_count():
    messages = sqs_client.get_queue_attributes(
        QueueUrl=os.environ['SQS_QUEUE_URL'],
        AttributeNames=['ApproximateNumberOfMessages']
    )
    return int(messages["Attributes"]["ApproximateNumberOfMessages"])

# Count the number of active ec2 instances
def running_instance_count():
    running_instances = ec2_client.instances.filter(Filters=[
        {
            'Instance State': 'instance-state-name',
            'Values': ['running'],
            'Name': 'tag:Name',
            'Values': ['NameOfInstance']
        },
    ])
    return len(running_instances)

# Evaluate the set of conditions
def evaluate_conditions(context, event):
    sqs_count = sqs_msg_count()
    sqs_average = int(os.environ['AVG_QUEUE_SIZE'])
    ec2_count = running_instance_count(ec2_client)
    ec2_average = int(os.environ['AVG_NR_OF_EC2_SCAN_SERVERS'])
    
    if sqs_count > sqs_average and ec2_count > ec2_average:
        print('False alert')
        logger.info()
    elif sqs_count < sqs_average and ec2_count < ec2_average:
        print('False alert')
        logger.info()
    elif sqs_count < sqs_average and ec2_count > ec2_average:
        print('False alert')
        logger.info()
    else:
        sns_client.publish(
            TopicArn=os.environ['SNS_ARN'],
            Message='sameple message',
            Subject='sample subject'
        )
        print("Published to SNS Topic")