#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.


import psycopg2
import json
from timeit import default_timer as timer
import logging
import boto3
from botocore.exceptions import ClientError

import logan_db_config


def record_queue_depth_lambda_handler(event, context):

    my_payload = run()

    return {
        "statusCode": 200,
        "body": json.dumps({
            "message": "hello world",
            "payload": my_payload,
        }),
    }


def get_queue_depth(sqs_queue_url):

    sqs_client = boto3.client('sqs', region_name='us-east-2')
    try:
        msg = sqs_client.get_queue_attributes(
            QueueUrl=sqs_queue_url,
            AttributeNames=['ApproximateNumberOfMessages'],
        )
    except ClientError as e:
        logging.error(e)
        return None
    return msg['Attributes']['ApproximateNumberOfMessages']


def run():

    print("Now connecting to database...")

    conn = psycopg2.connect(
        host=logan_db_config.db_hostname,
        database=logan_db_config.db_name,
        user=logan_db_config.db_username,
        password=logan_db_config.db_password)

    with conn.cursor() as cur:

        cur.execute('SET SESSION lock_timeout = 3000;')
        cur.execute('SET SESSION statement_timeout = %d;' % (1000*60*3))  # 3 minutes

        print("Work begins now...")

        start = timer()

        queue_depth = get_queue_depth(logan_db_config.sqs_queue_url)

        cur.execute("INSERT INTO lambda_logging.sqs_queue_depth_history (queue_depth) VALUES (%s) RETURNING inserted_at;", (queue_depth,))
        insertion_timestamp = cur.fetchone()[0]
        conn.commit()

        print("Inserted queue depth at %s..." % insertion_timestamp)

        end = timer()

        execution_seconds = end - start
        print("Completed in", execution_seconds, "seconds")

        return queue_depth


if __name__ == "__main__":

    run()
