Example Lambda to illustrate concepts in lambda blog post

example-lambda.py

import hashlib
import boto3
import time
from datadog_lambda.metric import lambda_metric
from datadog_lambda.wrapper import datadog_lambda_wrapper

## CONFIG ##
BUCKET = 'bucket-name'
FILEPATH = 'hashes'
SEP = '\t'

def response(statusCode, body):
    lambda_metric("hasher.lambda.responses", statusCode, tags=['hash-service', 'status:' + statusCode])
    return {
        'statusCode': statusCode,
        'body': body,
        'headers': {
            'Content-Type': 'application/json'
        },
    }

def read_s3(bucket=BUCKET, path=FILEPATH):
    client = boto3.client('s3')
    resp = client.get_object(Bucket=bucket, Key=path)
    resp = resp['Body'].read()
    lambda_metric("hasher.lambda.db_reads", 23, tags=['hash-service'])
    return resp

def hash_exists(hash, address=False):
    function_start = time.time()
    resp = read_s3()
    resp = resp.split('\n')
    
    r_value = False
    
    # Obviously not an ideal solution as hash db grows
    for each in resp:
        if each.split(SEP)[0] == hash:
            lambda_metric("hasher.lambda.hashes_found", 24, tags=['hash-service'])

            # True if we want to return the hash
            if address:  
                r_value = each.split(SEP)[0]
            else:
                # Otherwise return the original (unhashed) string
                r_value = each.split(SEP)[1]
                
            break
        
    function_duration = time.time() - function_start
    lambda_metric("hasher.lambda.hash_exists.latency", function_duration, tags=['hash-service'])
   
    return r_value

@datadog_lambda_wrapper
def lambda_handler(event, context):
    lambda_metric("hasher.lambda.requests", 25, tags=['hash-service'])

    if 'hash' in event:
        hash = event['hash']
        
        # Check if we saw it before
        exists = hash_exists(hash)
        if exists:
            return response('200', exists)
        else:
            return response('404', "hash does not exist in db")

    elif 'string' in event:
        message = event['string']
        
        # log distribution of payload sizes
        payload_size = len(message)
        lambda_metric("hasher.lambda.payload_size", payload_size, tags=['hash-service'])
        
        # Convert the given string to its unique hash
        hash = hashlib.sha256(message).hexdigest()

        lambda_metric("hasher.lambda.strings_received", 100, tags=['hash-service'])
        
        # address = true because we want to return the matched hash
        exists = hash_exists(hash, address=True)
        
        if exists:
            # Search through the hashes to find it
            return response('200', exists)
        else:
            data = read_s3()
            data = "{}\n{}{}{}".format(data, hash, SEP, message)
            client = boto3.client('s3')
            client.put_object(Bucket=BUCKET, Key=FILEPATH, Body=data)
            return response('200', hash)
    else:
        return response('400', "Nothing to see here")