Skip to content

Instantly share code, notes, and snippets.

@crucialfelix
Created January 11, 2020 00:05
Show Gist options
  • Save crucialfelix/806e11ce0368e40922352e185e9dc668 to your computer and use it in GitHub Desktop.
Save crucialfelix/806e11ce0368e40922352e185e9dc668 to your computer and use it in GitHub Desktop.
"""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()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment