check_connections.py

"""Check all required network connections
This can be called in AWS or in local docker-compose to check
that all settings work for connecting to database, memcached and rabbit.
Returns:
    exits with error unless all connections are OK
"""
import os
import socket
import subprocess
import sys

from django.conf import settings
from django.core.cache import cache

import paths  # noqa
from globalapp.tasks import ping

database = settings.DATABASES['default']
DATABASE_HOST = database['HOST']
DATABASE_PORT = database['PORT']
POSTGRES_USER = database['USER']
POSTGRES_DB = database['NAME']
RABBIT_PORT = 5672


def line(caption):
    print("")
    hr = "*" * 35
    print("%s %s %s" % (hr, caption, hr)),
    print("")


def check_socket(host, port):
    line(host)
    print("CHECKING DNS {}".format(host))
    err = None
    try:
        p = subprocess.Popen(["dig", host], env=os.environ, stdout=subprocess.PIPE)
        out, err = p.communicate()
    except Exception as error:
        print("dig ERROR {} {}".format(error, err))
    else:
        # ANSWER: 0
        if "ANSWER: 0" in out:
            print("!!! DNS RECORD NOT FOUND")
            return False
        else:
            print("dig OK")
        print(out)

    print("CHECKING SOCKET host={} port={}".format(host, port))
    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        s.settimeout(0.5)
        s.connect((host, int(port)))
    except Exception as error:
        print("!!! SOCKET ERROR {} {} {}".format(host, port, error))
        return False
    else:
        s.close()
    print "SOCKET OK {}:{}".format(host, port)
    return True


def check_db():
    line("check_db")
    command = ["psql", "-h", DATABASE_HOST, '-U', POSTGRES_USER, "--dbname", POSTGRES_DB, "-w", "-c", "\\l"]
    # print(command)
    env = os.environ.copy()
    # print("database", database)
    env['PGPASSWORD'] = database['PASSWORD']
    err = None
    try:
        p = subprocess.Popen(command, env=env, stdout=subprocess.PIPE)
        out, err = p.communicate()
    except Exception as error:
        print("POSTGRES connect ERROR: {} {}".format(error, err))
        return False
    else:
        print(out)
        if "FATAL" in out:
            print("POSTGRES login failed")
            return False
        print("POSTGRES connect OK")
        return True


def check_rabbit():
    line("check_rabbit")
    # connect to rabbit
    try:
        ping.apply_async()
    except Exception as error:
        print("RABBIT connect failure: {}".format(error))
        return False
    else:
        print("RABBIT connect OK")
        return True


def check_cache():
    line("check_cache")
    # connect to cache
    key = "__TEST_GET_SET__"
    value = "OK"
    cache.set(key, value)
    got = cache.get(key)
    ok = got == value
    print("CACHE can set and get: {}".format(ok))
    return ok


def main():
    checks = [
        (DATABASE_HOST, DATABASE_PORT),
        (settings.RABBITMQ_HOST, RABBIT_PORT),
    ]
    checks.extend(
        loc.split(':')
        for loc in settings.CACHES['default']["LOCATION"]
    )

    all_ok = [
        check_socket(h, p)
        for h, p in checks
    ]

    all_ok.extend([
        check_db(),
        check_rabbit(),
        check_cache()
    ])

    if all(all_ok):
        print("\nALL CONNECTIONS OK")
    else:
        print("\nERROR: CONNECTION FAILURES")
        sys.exit(1)


if __name__ == "__main__":
    main()