app.py

"""
Lambda function for Python Driver testing

Creates the client that is cached for all requests, subscribes to
relevant events, and forces the connection pool to get populated.
"""
import json
import os

from pymongo import MongoClient
from pymongo.monitoring import (
    CommandListener,
    ConnectionPoolListener,
    ServerHeartbeatListener,
)

open_connections = 0
heartbeat_count = 0
total_heartbeat_duration = 0
total_commands = 0
total_command_duration = 0


class CommandHandler(CommandListener):
    def started(self, event):
        print("command started", event)

    def succeeded(self, event):
        global total_commands, total_command_duration
        total_commands += 1
        total_command_duration += event.duration_micros / 1e6
        print("command succeeded", event)

    def failed(self, event):
        global total_commands, total_command_duration
        total_commands += 1
        total_command_duration += event.duration_micros / 1e6
        print("command failed", event)


class ServerHeartbeatHandler(ServerHeartbeatListener):
    def started(self, event):
        print("server heartbeat started", event)

    def succeeded(self, event):
        global heartbeat_count, total_heartbeat_duration
        heartbeat_count += 1
        total_heartbeat_duration += event.duration
        print("server heartbeat succeeded", event)

    def failed(self, event):
        global heartbeat_count, total_heartbeat_duration
        heartbeat_count += 1
        total_heartbeat_duration += event.duration
        print("server heartbeat  failed", event)


class ConnectionHandler(ConnectionPoolListener):
    def connection_created(self, event):
        global open_connections
        open_connections += 1
        print("connection created")

    def connection_ready(self, event):
        pass

    def connection_closed(self, event):
        global open_connections
        open_connections -= 1
        print("connection closed")

    def connection_check_out_started(self, event):
        pass

    def connection_check_out_failed(self, event):
        pass

    def connection_checked_out(self, event):
        pass

    def connection_checked_in(self, event):
        pass

    def pool_created(self, event):
        pass

    def pool_ready(self, event):
        pass

    def pool_cleared(self, event):
        pass

    def pool_closed(self, event):
        pass


listeners = [CommandHandler(), ServerHeartbeatHandler(), ConnectionHandler()]
print("Creating client")
client = MongoClient(os.environ["MONGODB_URI"], event_listeners=listeners)


# Populate the connection pool.
print("Connecting")
client.lambdaTest.list_collections()
print("Connected")


# Create the response to send back.
def create_response():
    return dict(
        averageCommandDuration=total_command_duration / total_commands,
        averageHeartbeatDuration=total_heartbeat_duration / heartbeat_count,
        openConnections=open_connections,
        heartbeatCount=heartbeat_count,
    )


# Reset the numbers.
def reset():
    global open_connections, heartbeat_count, total_heartbeat_duration, total_commands, total_command_duration
    open_connections = 0
    heartbeat_count = 0
    total_heartbeat_duration = 0
    total_commands = 0
    total_command_duration = 0


def lambda_handler(event, context):
    """
    The handler function itself performs an insert/delete and returns the
    id of the document in play.
    """
    print("initializing")
    db = client.lambdaTest
    collection = db.test
    result = collection.insert_one({"n": 1})
    collection.delete_one({"_id": result.inserted_id})
    # Create the response and then reset the numbers.
    response = json.dumps(create_response())
    reset()
    print("finished!")

    return dict(statusCode=200, body=response)