A Simple PostgreSQL connection checker implemented as a HTTP server.

pg-connection-check-http.py

#!/usr/bin/env python

"""
A Simple PostgreSQL connection checker implemented as a HTTP server.

Helps with things like Amazon ELB with a slightly more nuanced check than a TCP
socket. You can pass in a DSN string, or inherit/use the environment variables
(see the psql man page).

See http://initd.org/psycopg/docs/module.html#psycopg2.connect for details on
constructing a DSN string.

Suggestion: don't run it as the postgres user or another superuser, since in
troubled times they may be able to connect when others can't.

Returns a 200 response if the connection is okay, a 500 otherwise with
the error message returned by libpq.
"""

import time
import BaseHTTPServer
from optparse import OptionParser

import psycopg2


class Handler(BaseHTTPServer.BaseHTTPRequestHandler):
    # default is to get connection properties from the environment
    dsn = ''
    timeout = 5

    def do_GET(s):
        dsn = '%s connect_timeout=%d' % (Handler.dsn, Handler.timeout)

        error = None
        try:
            conn = psycopg2.connect(dsn)
            conn.close()
        except Exception, e:
            error = unicode(e).strip()
            print error

        if error:
            s.send_response(500)
            s.send_header("Content-type", "text/plain; charset='utf-8")
            s.end_headers()
            s.wfile.write(error.encode('utf-8'))
        else:
            s.send_response(200)
            s.send_header("Content-type", "text/plain; charset='utf-8")
            s.end_headers()
            s.wfile.write("OK")


def main():
    usage = "Usage: %prog [options]\n" + __doc__

    parser = OptionParser(usage=usage)
    parser.add_option("-d", "--dsn", action="store", type="str",
                      dest="dsn", default=Handler.dsn,
                      help="Database connection string. Default '%default'")
    parser.add_option("-p", "--port", action="store", type="int",
                      dest="port", default=7432,
                      help="Port to listen on. Default %default")
    parser.add_option("-i", "--interface", action="store", type="str",
                      dest="interface", default='0.0.0.0',
                      help="Interface to listen on. Default %default")
    parser.add_option("-t", "--timeout", action="store", type="int",
                      dest="timeout", default=Handler.timeout,
                      help="Connection timeout. Default %default")

    (options, args) = parser.parse_args()

    if len(args):
        parser.error("I don't take arguments.")

    Handler.dsn = options.dsn
    Handler.timeout = options.timeout

    server_class = BaseHTTPServer.HTTPServer
    listen = (options.interface, options.port)
    httpd = server_class(listen, Handler)

    print time.asctime(), "Server Started - %s:%s" % listen
    try:
        httpd.serve_forever()
    except KeyboardInterrupt:
        pass
    httpd.server_close()
    print time.asctime(), "Server Stopped"


if __name__ == '__main__':
    main()