postgresql check support via momoko

README.md

@@ -13,6 +13,7 @@ This will check the following locations for service state:
   * if `tcp`: will attempt to connect to port `<port>` on localhost. `<query>` is currently ignored
   * if `spool`: will only check the spool state
   * if `mysql` and the `mysql_username` and `mysql_password` are set, will do a login and quit on the requested mysql port; `<query>` is ignored and no logical database is selected.
+  * if `username` and the `username_username` and `username_password` are set, will do a login, SELECT 1, and quit on the requested username port; `<query>` is the logical database name
 
 When it does query the actual service check endpoint, **hacheck** MAY cache the value of that query for some amount of time
 
@@ -36,6 +37,8 @@ Imagine you want to take down the server `web01` for maintenance. Just SSH to it
 * `log_path`: Either the string `"stdout"`, the string `"stderr"`, or a fully-qualified path to a file to write logs to. Uses a [WatchedFileHandler](http://docs.python.org/2/library/logging.handlers.html#watchedfilehandler) and ought to play nicely with logrotate
 * `mysql_username`: username to use when logging into mysql for checks
 * `mysql_password`: password to use when logging into mysql for checks
+* `postgresql_username`: username to use when logging into postgresql for checks
+* `postgresql_password`: password to use when logging into postgresql for checks
 
 ### License
 

hacheck/checker.py

@@ -160,3 +160,42 @@ def timed_out(duration):
         raise tornado.gen.Return((500, 'MySQL sez %s' % response))
     yield conn.quit()
     raise tornado.gen.Return((200, 'MySQL connect response: %s' % response))
+
+
+@cache.cached
+@tornado.gen.coroutine
+def check_postgresql(service_name, port, query, io_loop, query_params, headers):
+    dbname = query or 'uber'
+
+    username = config.config.get('postgresql_username', None)
+    password = config.config.get('postgresql_password', None)
+    if username is None or password is None:
+        raise tornado.gen.Return((500, 'No PostgreSQL username/password in config file'))
+
+    try:
+        import momoko.connection
+    except ImportError:
+        raise tornado.gen.Return((500, 'momoko module not available'))
+
+    dsn = "host=127.0.0.1 port=%d user=%s password=%s connect_timeout=%d dbname=%s" % (
+        port, username, password, 2, dbname
+    )
+
+    c = momoko.connection.Connection()
+    (conn, error), _ = yield tornado.gen.Task(c.connect, dsn)
+
+    if error:
+        raise tornado.gen.Return((503, 'Error connecting to postgresql as user %s: %r' % (
+            username, error
+        )))
+
+    (cursor, error), _ = yield tornado.gen.Task(conn.execute, 'SELECT 2')
+
+    if error:
+        raise tornado.gen.Return((503, 'SELECT 1 returned error: %r' % error))
+
+    # is this synchronous?
+    response = cursor.fetchall()
+    conn.close()
+    success = response == [(1,)]
+    raise tornado.gen.Return((200 if success else 503, 'PING response: %r' % response))

hacheck/config.py

@@ -7,6 +7,8 @@
     'log_path': (str, 'stderr'),
     'mysql_username': (str, None),
     'mysql_password': (str, None),
+    'postgresql_username': (str, None),
+    'postgresql_password': (str, None),
 }
 
 

hacheck/handlers.py

@@ -47,21 +47,28 @@ def get(self, service_name, port, query):
             last_message = ""
             querystr = self.request.query
             for this_checker in self.CHECKERS:
-                code, message = yield this_checker(
-                    service_name,
-                    port,
-                    query,
-                    io_loop=tornado.ioloop.IOLoop.current(),
-                    query_params=querystr,
-                    headers=self.request.headers,
-                )
-                last_message = message
-                if code > 200:
-                    if code in tornado.httputil.responses:
-                        self.set_status(code)
-                    else:
-                        self.set_status(503)
-                    self.write(message)
+                try:
+                    code, message = yield this_checker(
+                        service_name,
+                        port,
+                        query,
+                        io_loop=tornado.ioloop.IOLoop.current(),
+                        query_params=querystr,
+                        headers=self.request.headers,
+                    )
+                    last_message = message
+                    if code > 200:
+                        if code in tornado.httputil.responses:
+                            self.set_status(code)
+                        else:
+                            self.set_status(503)
+                        self.write(message)
+                        self.finish()
+                        break
+                except Exception as e:
+                    log.exception('Error running checker %s' % this_checker)
+                    self.set_status(500)
+                    self.write(str(e))
                     self.finish()
                     break
             else:
@@ -84,3 +91,7 @@ class TCPServiceHandler(BaseServiceHandler):
 
 class MySQLServiceHandler(BaseServiceHandler):
     CHECKERS = [checker.check_spool, checker.check_mysql]
+
+
+class PostgreSQLServiceHandler(BaseServiceHandler):
+    CHECKERS = [checker.check_spool, checker.check_postgresql]

hacheck/main.py

@@ -27,6 +27,7 @@ def get_app():
         (r'/http/([a-zA-Z0-9_-]+)/([0-9]+)/(.*)', handlers.HTTPServiceHandler),
         (r'/tcp/([a-zA-Z0-9_-]+)/([0-9]+)/?(.*)', handlers.TCPServiceHandler),
         (r'/mysql/([a-zA-Z0-9_-]+)/([0-9]+)/?(.*)', handlers.MySQLServiceHandler),
+        (r'/postgresql/([a-zA-Z0-9_-]+)/([0-9]+)/?(.*)', handlers.PostgreSQLServiceHandler),
         (r'/spool/([a-zA-Z0-9_-]+)/([0-9]+)/?(.*)', handlers.SpoolServiceHandler),
         (r'/recent', handlers.ListRecentHandler),
         (r'/status', handlers.StatusHandler),

requirements.txt

@@ -2,3 +2,4 @@ tornado>=3.0.1,<4.1
 futures
 PyYAML>=3.0
 six>=1.4.0
+Momoko>=1.1.5