diff --git a/governor.py b/governor.py index d12c7093..27102d88 100755 --- a/governor.py +++ b/governor.py @@ -7,6 +7,7 @@ import sys import time import yaml +from helpers.api import RestApiServer from helpers.etcd import Etcd from helpers.postgresql import Postgresql from helpers.ha import Ha @@ -81,6 +82,7 @@ def main(): governor = Governor(config) try: governor.initialize() + RestApiServer(governor).start() governor.run() finally: governor.postgresql.stop() diff --git a/helpers/api.py b/helpers/api.py new file mode 100644 index 00000000..6269a83c --- /dev/null +++ b/helpers/api.py @@ -0,0 +1,69 @@ +import json +import logging +import psycopg2 +import sys + +from threading import Thread + +if sys.hexversion >= 0x03000000: + from http.server import BaseHTTPRequestHandler, HTTPServer +else: + from BaseHTTPServer import BaseHTTPRequestHandler, HTTPServer + + +logger = logging.getLogger(__name__) + + +class RestApiHandler(BaseHTTPRequestHandler): + + def do_GET(self): + try: + response = self.get_postgresql_status() + except (psycopg2.OperationalError, psycopg2.InterfaceError): + response = {'running': False} + + path = '/master' if self.path == '/' else self.path + status_code = 200 if response['running'] and response['role'] in path else 503 + + self.send_response(status_code) + self.send_header('Content-Type', 'application/json') + self.end_headers() + self.wfile.write(json.dumps(response).encode('utf-8')) + + def get_postgresql_status(self): + if not self.server.governor.postgresql.is_running(): + return {'running': False} + cursor = self.server.cursor() + cursor.execute("""SELECT to_char(pg_postmaster_start_time(), 'YYYY-MM-DD HH24:MI:SS.MS TZ'), + pg_is_in_recovery(), + pg_current_xlog_location(), + pg_last_xlog_receive_location(), + pg_last_xlog_replay_location(), + pg_is_in_recovery() AND pg_is_xlog_replay_paused()""") + row = cursor.fetchone() + return { + 'running': True, + 'postmaster_start_time': row[0], + 'role': 'slave' if row[1] else 'master', + 'xlog': ({ + 'received_location': row[3], + 'replayed_location': row[4], + 'paused': row[5]} if row[1] else { + 'location': row[2] + }) + } + + +class RestApiServer(HTTPServer, Thread): + + def __init__(self, governor, listen_address='0.0.0.0', listen_port=8080): + HTTPServer.__init__(self, (listen_address, listen_port), RestApiHandler) + Thread.__init__(self, target=self.serve_forever) + self.governor = governor + self._cursor_holder = None + self.daemon = True + + def cursor(self): + if not self._cursor_holder or self._cursor_holder.closed != 0: + self._cursor_holder = self.governor.postgresql.connection().cursor() + return self._cursor_holder diff --git a/tests/test_api.py b/tests/test_api.py new file mode 100644 index 00000000..6b117b48 --- /dev/null +++ b/tests/test_api.py @@ -0,0 +1,66 @@ +import psycopg2 +import sys +import unittest + +from helpers.api import RestApiHandler, RestApiServer +from test_postgresql import psycopg2_connect + +if sys.hexversion >= 0x03000000: + from io import BytesIO as IO +else: + from StringIO import StringIO as IO + + +def false(*args, **kwargs): + return False + + +def throws(*args, **kwargs): + raise psycopg2.OperationalError() + + +class MockPostgresql: + + def connection(self): + return psycopg2_connect() + + def is_running(self): + return True + + +class MockGovernor: + + def __init__(self): + self.postgresql = MockPostgresql() + + +class MockRequest: + + def __init__(self, path): + self.path = path + + def makefile(self, *args, **kwargs): + return IO(self.path) + + +class MockRestApiServer(RestApiServer): + + def __init__(self, Handler, path, *args): + self.governor = MockGovernor() + if len(args) > 0: + self.governor.postgresql.is_running = args[0] + self._cursor_holder = None + Handler(MockRequest(path), ('0.0.0.0', 8080), self) + + +class TestRestApiHandler(unittest.TestCase): + + def __init__(self, method_name='runTest'): + super(TestRestApiHandler, self).__init__(method_name) + + def test_do_GET(self): + MockRestApiServer(RestApiHandler, b'GET /') + MockRestApiServer(RestApiHandler, b'GET /', throws) + + def test_get_postgresql_status(self): + MockRestApiServer(RestApiHandler, b'GET /', false)