diff --git a/governor.py b/governor.py index bf75efaa..25647355 100755 --- a/governor.py +++ b/governor.py @@ -4,14 +4,13 @@ import logging import os import signal import sys -import threading import time import yaml +from helpers.api import RestApiServer from helpers.etcd import Etcd from helpers.postgresql import Postgresql from helpers.ha import Ha -from helpers.statuspage import getHTTPServer def sigterm_handler(signo, stack_frame): @@ -37,8 +36,8 @@ class Governor: self.postgresql = Postgresql(config['postgresql']) self.ha = Ha(self.postgresql, self.etcd) - def touch_member(self): - return self.etcd.touch_member(self.postgresql.name, self.postgresql.connection_string) + def touch_member(self, ttl=None): + return self.etcd.touch_member(self.postgresql.name, self.postgresql.connection_string, ttl) def initialize(self): # wait for etcd to be available @@ -74,6 +73,10 @@ class Governor: def main(): + logging.basicConfig(format='%(asctime)s %(levelname)s: %(message)s', level=logging.INFO) + signal.signal(signal.SIGTERM, sigterm_handler) + signal.signal(signal.SIGCHLD, sigchld_handler) + if len(sys.argv) < 2 or not os.path.isfile(sys.argv[1]): print('Usage: {} config.yml'.format(sys.argv[0])) return @@ -82,25 +85,17 @@ def main(): config = yaml.load(f) governor = Governor(config) - - # Start the http_server to serve a simple healthcheck - http_server = getHTTPServer(governor.postgresql, http_port=config.get( - 'healtcheck_port', 8008), listen_address='0.0.0.0') - http_thread = threading.Thread(target=http_server.serve_forever, args=()) - http_thread.daemon = True - - governor.initialize() - http_thread.start() - try: + governor.initialize() + # Start the http_server to serve a simple healthcheck + host, port = config['restapi']['listen'].split(':') + RestApiServer(governor, host, int(port)).start() governor.run() finally: + governor.touch_member(300) # schedule member removal governor.postgresql.stop() governor.etcd.delete_leader(governor.postgresql.name) if __name__ == '__main__': - logging.basicConfig(format='%(asctime)s %(levelname)s: %(message)s', level=logging.INFO) - signal.signal(signal.SIGTERM, sigterm_handler) - signal.signal(signal.SIGCHLD, sigchld_handler) main() diff --git a/helpers/api.py b/helpers/api.py new file mode 100644 index 00000000..f0d47e10 --- /dev/null +++ b/helpers/api.py @@ -0,0 +1,70 @@ +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): + logging.exception('get_postgresql_status') + 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: + self._cursor_holder = self.governor.postgresql.connection().cursor() + return self._cursor_holder diff --git a/helpers/errors.py b/helpers/errors.py index dd1bc698..b86bd0f0 100644 --- a/helpers/errors.py +++ b/helpers/errors.py @@ -9,7 +9,3 @@ class EtcdError(Exception): class CurrentLeaderError(EtcdError): pass - - -class HealthiestMemberError(EtcdError): - pass diff --git a/helpers/etcd.py b/helpers/etcd.py index e808f8d4..6da6a3bc 100644 --- a/helpers/etcd.py +++ b/helpers/etcd.py @@ -2,16 +2,16 @@ import logging import requests import time +from requests.exceptions import RequestException from collections import namedtuple from helpers.errors import CurrentLeaderError, EtcdError logger = logging.getLogger(__name__) - Member = namedtuple('Member', 'hostname,address,ttl') -class Cluster(namedtuple('Cluster', 'leader,last_leader_operation,members')): +class Cluster(namedtuple('Cluster', 'initialize,leader,last_leader_operation,members')): def is_unlocked(self): return not (self.leader and self.leader.hostname) @@ -35,7 +35,7 @@ class Etcd: response = requests.get(self.client_url(path)) if response.status_code == 200: break - except Exception as e: + except RequestException as e: logger.exception('get_client_path') ex = e @@ -54,15 +54,15 @@ class Etcd: try: response = requests.put(self.client_url(path), data=data) return response.status_code in [200, 201, 202, 204] - except: + except RequestException: logger.exception('PUT %s data=%s', path, data) - return False + raise EtcdError('Etcd is not responding properly') def delete_client_path(self, path): try: response = requests.delete(self.client_url(path)) return response.status_code in [200, 202, 204] - except: + except RequestException: logger.exception('DELETE %s', path) return False @@ -87,6 +87,8 @@ class Etcd: try: response, status_code = self.get_client_path('?recursive=true') if status_code == 200: + node = self.find_node(response['node'], '/initialize') + initialize = True if node else False # get list of members node = self.find_node(response['node'], '/members') or {'nodes': []} members = [Member(n['key'].split('/')[-1], n['value'], n.get('ttl', None)) for n in node['nodes']] @@ -108,11 +110,11 @@ class Etcd: leader = m break if not leader: - leader = Member(leader['value'], None, None) + leader = Member(node['value'], None, None) - return Cluster(leader, last_leader_operation, members) + return Cluster(initialize, leader, last_leader_operation, members) elif status_code == 404: - return Cluster(None, None, []) + return Cluster(False, None, None, []) except: logger.exception('get_cluster') @@ -122,27 +124,43 @@ class Etcd: try: cluster = self.get_cluster() return None if cluster.is_unlocked() else cluster.leader - except: - raise CurrentLeaderError("Etcd is not responding properly") + except EtcdError: + raise CurrentLeaderError('Etcd is not responding properly') - def touch_member(self, member, connection_string): - return self.put_client_path('/members/' + member, value=connection_string, ttl=self.member_ttl) + def touch_member(self, member, connection_string, ttl=None): + try: + return self.put_client_path('/members/' + member, value=connection_string, ttl=ttl or self.member_ttl) + except EtcdError: + return False def take_leader(self, value): - return self.put_client_path('/leader', value=value, ttl=self.ttl) + try: + return self.put_client_path('/leader', value=value, ttl=self.ttl) + except EtcdError: + return False def attempt_to_acquire_leader(self, value): - ret = self.put_client_path('/leader', value=value, ttl=self.ttl, prevExist=False) - ret or logger.info('Could not take out TTL lock') - return ret + try: + ret = self.put_client_path('/leader', value=value, ttl=self.ttl, prevExist=False) + ret or logger.info('Could not take out TTL lock') + return ret + except EtcdError: + return False def update_leader(self, state_handler): - ret = self.put_client_path('/leader', value=state_handler.name, ttl=self.ttl, prevValue=state_handler.name) - ret and self.put_client_path('/optime/leader', value=state_handler.last_operation()) - return ret + if self.put_client_path('/leader', value=state_handler.name, ttl=self.ttl, prevValue=state_handler.name): + try: + self.put_client_path('/optime/leader', value=state_handler.last_operation()) + except EtcdError: + pass + return True + return False def race(self, path, value): - return self.put_client_path(path, value=value, prevExist=False) + try: + return self.put_client_path(path, value=value, prevExist=False) + except EtcdError: + return False def delete_member(self, member): return self.delete_client_path('/members/' + member) diff --git a/helpers/ha.py b/helpers/ha.py index 74c6a892..29979363 100644 --- a/helpers/ha.py +++ b/helpers/ha.py @@ -1,7 +1,7 @@ import logging -from helpers.errors import EtcdError, HealthiestMemberError -from psycopg2 import OperationalError +from helpers.errors import EtcdError +from psycopg2 import InterfaceError, OperationalError logger = logging.getLogger(__name__) @@ -92,7 +92,5 @@ class Ha: if self.state_handler.is_leader(): self.state_handler.demote(None) return 'demoted self because etcd is not accessible and i was a leader' - except OperationalError: + except (InterfaceError, OperationalError): logger.error('Error communicating with Postgresql. Will try again') - except HealthiestMemberError: - logger.error('failed to determine healthiest member fromt etcd') diff --git a/helpers/postgresql.py b/helpers/postgresql.py index 75ce0998..de75f8b1 100644 --- a/helpers/postgresql.py +++ b/helpers/postgresql.py @@ -19,26 +19,27 @@ logger = logging.getLogger(__name__) def parseurl(url): r = urlparse(url) - return { - 'hostname': r.hostname, + ret = { + 'host': r.hostname, 'port': r.port or 5432, - 'username': r.username, - 'password': r.password, + 'database': r.path[1:], + 'fallback_application_name': 'Governor', + 'connect_timeout': 3, + 'options': '-c statement_timeout=2000', } + if r.username: + ret['user'] = r.username + if r.password: + ret['password'] = r.password + return ret class Postgresql: def __init__(self, config): + self.config = config self.name = config['name'] self.listen_addresses, self.port = config['listen'].split(':') - self.libpq_parameters = { - 'host': self.listen_addresses.split(',')[0].strip(), - 'port': self.port, - 'fallback_application_name': 'Governor', - 'connect_timeout': 5, - 'options': '-c statement_timeout=2000' - } self.data_dir = config['data_dir'] self.replication = config['replication'] self.superuser = config['superuser'] @@ -46,51 +47,62 @@ class Postgresql: self.recovery_conf = os.path.join(self.data_dir, 'recovery.conf') self.configuration_to_save = (os.path.join(self.data_dir, 'pg_hba.conf'), os.path.join(self.data_dir, 'postgresql.conf')) + self.pid_path = os.path.join(self.data_dir, 'postmaster.pid') self._pg_ctl = 'pg_ctl -w -D ' + self.data_dir self.wal_e = config.get('wal_e', None) if self.wal_e: self.wal_e_path = 'envdir {} wal-e --aws-instance-profile '.\ format(self.wal_e.get('env_dir', '/home/postgres/etc/wal-e.d/env')) - self.config = config - - connect_address = config.get('connect_address', config['listen']) or config['listen'] + self.local_address = self.get_local_address() + connect_address = config.get('connect_address', None) or self.local_address self.connection_string = 'postgres://{username}:{password}@{connect_address}/postgres'.format( connect_address=connect_address, **self.replication) - self.conn = None - self.cursor_holder = None + self._connection = None + self._cursor_holder = None self.members = [] # list of already existing replication slots - def cursor(self): - if not self.cursor_holder: - self.conn = psycopg2.connect(**self.libpq_parameters) - self.conn.autocommit = True - self.cursor_holder = self.conn.cursor() + def get_local_address(self): + # TODO: try to get unix_socket_directory from postmaster.pid + return self.listen_addresses.split(',')[0].strip() + ':' + self.port - return self.cursor_holder + def connection(self): + if not self._connection or self._connection.closed != 0: + r = parseurl('postgres://{}/postgres'.format(self.local_address)) + self._connection = psycopg2.connect(**r) + self._connection.autocommit = True + return self._connection + + def _cursor(self): + if not self._cursor_holder or self._cursor_holder.closed: + self._cursor_holder = self.connection().cursor() + return self._cursor_holder def disconnect(self): - try: - self.conn.close() - except: - logger.exception('Error disconnecting') + self._connection and self._connection.close() + self._connection = self._cursor_holder = None def query(self, sql, *params): max_attempts = 0 while True: + ex = None try: - self.cursor().execute(sql, params) - break + cursor = self._cursor() + cursor.execute(sql, params) + return cursor + except psycopg2.InterfaceError as e: + ex = e except psycopg2.OperationalError as e: - if self.conn: - self.disconnect() - self.cursor_holder = None - if max_attempts > 4: + if self._connection and self._connection.closed == 0: raise e + ex = e + if ex: + self.disconnect() max_attempts += 1 + if max_attempts >= 3: + raise ex time.sleep(5) - return self.cursor() def data_directory_empty(self): return not os.path.exists(self.data_dir) or os.listdir(self.data_dir) == [] @@ -106,17 +118,17 @@ class Postgresql: pgpass = 'pgpass' with open(pgpass, 'w') as f: os.fchmod(f.fileno(), 0o600) - f.write('{hostname}:{port}:*:{username}:{password}\n'.format(**r)) + f.write('{host}:{port}:*:{user}:{password}\n'.format(**r)) try: os.environ['PGPASSFILE'] = pgpass - return self.create_replica(leader.address, r) == 0 + return self.create_replica(r) == 0 finally: os.environ.pop('PGPASSFILE') - def create_replica(self, master_connurl, master_connection): + def create_replica(self, master_connection): """ creates a new replica using either pg_basebackup or WAL-E """ - if self.should_use_s3_to_create_replica(master_connurl): + if self.should_use_s3_to_create_replica(master_connection): result = self.create_replica_with_s3() # if restore from the backup on S3 failed - try with the pg_basebackup if result == 0: @@ -124,7 +136,7 @@ class Postgresql: return self.create_replica_with_pg_basebackup(master_connection) def create_replica_with_pg_basebackup(self, master_connection): - return os.system('pg_basebackup -R -D {data_dir} --host={hostname} --port={port} -U {username}'.format( + return os.system('pg_basebackup -R -D {data_dir} --host={host} --port={port} -U {user}'.format( data_dir=self.data_dir, **master_connection)) def create_replica_with_s3(self): @@ -135,7 +147,7 @@ class Postgresql: self.restore_configuration_files() return ret - def should_use_s3_to_create_replica(self, master_connurl): + def should_use_s3_to_create_replica(self, master_connection): """ determine whether it makes sense to use S3 and not pg_basebackup """ if not self.wal_e or not self.wal_e_path: return False @@ -187,7 +199,7 @@ class Postgresql: diff_in_bytes = long(backup_size) try: # get the difference in bytes between the current WAL location and the backup start offset - conn = psycopg2.connect(master_connurl) + conn = psycopg2.connect(master_connection) conn.autocommit = True cursor = conn.cursor() cursor.execute("SELECT pg_xlog_location_diff(pg_current_xlog_location(), %s)", (backup_start_lsn,)) @@ -216,10 +228,9 @@ class Postgresql: logger.error('Cannot start PostgreSQL because one is already running.') return False - pid_path = os.path.join(self.data_dir, 'postmaster.pid') - if os.path.exists(pid_path): - os.remove(pid_path) - logger.info('Removed %s', pid_path) + if os.path.exists(self.pid_path): + os.remove(self.pid_path) + logger.info('Removed %s', self.pid_path) ret = os.system(self._pg_ctl + ' start -o "{}"'.format(self.server_options())) == 0 ret and self.load_replication_slots() @@ -248,6 +259,9 @@ class Postgresql: return True def is_healthiest_node(self, cluster): + if self.is_leader(): + return True + if cluster.last_leader_operation - self.xlog_position() > self.config.get('maximum_lag_on_failover', 0): return False @@ -255,18 +269,19 @@ class Postgresql: if member.hostname == self.name: continue try: - member_conn = psycopg2.connect(member.address) + member_conn = psycopg2.connect(parseurl(member.address)) member_conn.autocommit = True member_cursor = member_conn.cursor() member_cursor.execute( - "SELECT %s - (pg_last_xlog_replay_location() - '0/0000000'::pg_lsn)", (self.xlog_position(), )) - xlog_diff = member_cursor.fetchone()[0] - logger.info([self.name, member.hostname, xlog_diff]) + "SELECT pg_is_in_recovery(), %s - (pg_last_xlog_replay_location() - '0/0000000'::pg_lsn)", + (self.xlog_position(), )) + row = member_cursor.fetchone() member_cursor.close() member_conn.close() - if xlog_diff < 0: + logger.error([self.name, member.hostname, row]) + if not row[0] or row[1] < 0: return False - except psycopg2.OperationalError: + except psycopg2.Error: continue return True @@ -283,7 +298,7 @@ class Postgresql: @staticmethod def primary_conninfo(leader_url): r = parseurl(leader_url) - return 'user={username} password={password} host={hostname} port={port} sslmode=prefer sslcompression=1'.format(**r) + return 'user={user} password={password} host={host} port={port} sslmode=prefer sslcompression=1'.format(**r) def check_recovery_conf(self, leader): if not os.path.isfile(self.recovery_conf): @@ -357,7 +372,9 @@ primary_conninfo = '{}' self.admin['username']), self.admin['password']) def xlog_position(self): - return self.query("SELECT pg_last_xlog_replay_location() - '0/0000000'::pg_lsn").fetchone()[0] + return self.query("""SELECT CASE WHEN pg_is_in_recovery() + THEN pg_last_xlog_replay_location() - '0/0000000'::pg_lsn + ELSE pg_current_xlog_location() - '0/00000'::pg_lsn END""").fetchone()[0] def load_replication_slots(self): cursor = self.query("SELECT slot_name FROM pg_replication_slots WHERE slot_type='physical'") @@ -378,4 +395,4 @@ primary_conninfo = '{}' self.members = members def last_operation(self): - return self.query("SELECT pg_current_xlog_location() - '0/00000'::pg_lsn").fetchone()[0] + return self.xlog_position() diff --git a/helpers/statuspage.py b/helpers/statuspage.py deleted file mode 100644 index 6fcd9462..00000000 --- a/helpers/statuspage.py +++ /dev/null @@ -1,65 +0,0 @@ -#!/usr/bin/env python -# -*- coding: utf-8 -*- - -import json -import sys - -if sys.hexversion >= 0x03000000: - from http.server import BaseHTTPRequestHandler, HTTPServer -else: - from BaseHTTPServer import BaseHTTPRequestHandler, HTTPServer - - -class StatusPage(BaseHTTPRequestHandler): - - def do_GET(self): - content_type = 'text/plain' - if self.path == '/pg_master': - if not self.pg_is_in_recovery(): - response, content = 200, 'I am currently a master' - else: - response, content = 503, 'I am not a master' - elif self.path == '/pg_slave': - if self.pg_is_in_recovery(): - response, content = 200, 'I am currently a slave' - else: - response, content = 503, 'I am not a slave' - elif self.path == '/pg_status': - response, content = 200, self.pg_status() - content_type = 'application/json' - else: - response, content = 404, 'Page not found' - - self.send_response(response) - self.send_header('Content-Type', content_type) - self.end_headers() - self.wfile.write(content.encode('utf-8')) - - def pg_is_in_recovery(self): - cursor = self.server.postgresql.cursor() - cursor.execute('SELECT pg_is_in_recovery()') - res = cursor.fetchone() - return res[0] - - def pg_status(self): - cursor = self.server.postgresql.cursor() - cursor.execute(""" - SELECT pg_is_in_recovery(), - to_char(pg_last_xact_replay_timestamp(), 'YYYY-MM-DD HH24:MI:SS.MS TZ'), - extract(epoch from now() - pg_last_xact_replay_timestamp()), - inet_server_addr(), - inet_server_port(), - to_char(pg_postmaster_start_time(), 'YYYY-MM-DD HH24:MI:SS.MS TZ') - """) - res = cursor.fetchone() - status = {'role': ('master' if not res[0] else 'slave'), 'recovery': {'last_transaction_timestamp': res[1]}, - 'server': {'hostaddr': res[3], 'port': res[4], 'start_time': res[5]}} - - return json.dumps(status) - - -def getHTTPServer(postgresql, http_port=8081, listen_address='0.0.0.0'): - server = HTTPServer((listen_address, http_port), StatusPage) - server.postgresql = postgresql - - return server diff --git a/postgres0.yml b/postgres0.yml index 263cd693..8d56715b 100644 --- a/postgres0.yml +++ b/postgres0.yml @@ -1,5 +1,6 @@ loop_wait: 10 -healthcheck_port: 8008 +restapi: + listen: 127.0.0.1:8008 etcd: scope: batman ttl: 30 diff --git a/postgres1.yml b/postgres1.yml index 6d0082b2..6f08741f 100644 --- a/postgres1.yml +++ b/postgres1.yml @@ -1,5 +1,6 @@ loop_wait: 10 -healthcheck_port: 8009 +restapi: + listen: 127.0.0.1:8009 etcd: scope: batman ttl: 30 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) diff --git a/tests/test_etcd.py b/tests/test_etcd.py index c358e89e..8b3d60b2 100644 --- a/tests/test_etcd.py +++ b/tests/test_etcd.py @@ -17,28 +17,39 @@ class MockResponse: return json.loads(self.content) +class MockPostgresql: + name = '' + + def last_operation(self): + return 0 + + def requests_get(url, **kwargs): if url.startswith('http://local'): - raise Exception() + raise requests.exceptions.RequestException() response = MockResponse() if url.startswith('http://remote'): response.content = '{"action":"get","node":{"key":"/service/batman5","dir":true,"nodes":[{"key":"/service/batman5/initialize","value":"postgresql0","modifiedIndex":1582,"createdIndex":1582},{"key":"/service/batman5/leader","value":"postgresql1","expiration":"2015-05-15T09:11:00.037397538Z","ttl":21,"modifiedIndex":20728,"createdIndex":20434},{"key":"/service/batman5/optime","dir":true,"nodes":[{"key":"/service/batman5/optime/leader","value":"2164261704","modifiedIndex":20729,"createdIndex":20729}],"modifiedIndex":20437,"createdIndex":20437},{"key":"/service/batman5/members","dir":true,"nodes":[{"key":"/service/batman5/members/postgresql1","value":"postgres://replicator:rep-pass@127.0.0.1:5434/postgres","expiration":"2015-05-15T09:10:59.949384522Z","ttl":21,"modifiedIndex":20727,"createdIndex":20727},{"key":"/service/batman5/members/postgresql0","value":"postgres://replicator:rep-pass@127.0.0.1:5433/postgres","expiration":"2015-05-15T09:11:09.611860899Z","ttl":30,"modifiedIndex":20730,"createdIndex":20730}],"modifiedIndex":1581,"createdIndex":1581}],"modifiedIndex":1581,"createdIndex":1581}}' elif url.startswith('http://other'): response.status_code = 404 + elif url.startswith('http://noleader'): + response.content = '{"action":"get","node":{"key":"/service/batman5","dir":true,"nodes":[{"key":"/service/batman5/initialize","value":"postgresql0","modifiedIndex":1582,"createdIndex":1582},{"key":"/service/batman5/leader","value":"postgresql1","expiration":"2015-05-15T09:11:00.037397538Z","ttl":21,"modifiedIndex":20728,"createdIndex":20434},{"key":"/service/batman5/optime","dir":true,"nodes":[{"key":"/service/batman5/optime/leader","value":"2164261704","modifiedIndex":20729,"createdIndex":20729}],"modifiedIndex":20437,"createdIndex":20437},{"key":"/service/batman5/members","dir":true,"nodes":[{"key":"/service/batman5/members/postgresql0","value":"postgres://replicator:rep-pass@127.0.0.1:5433/postgres","expiration":"2015-05-15T09:11:09.611860899Z","ttl":30,"modifiedIndex":20730,"createdIndex":20730}],"modifiedIndex":1581,"createdIndex":1581}],"modifiedIndex":1581,"createdIndex":1581}}' return response def requests_put(url, **kwargs): - if url.startswith('http://local'): - raise Exception() + if url.startswith('http://local') or '/optime/leader' in url: + raise requests.exceptions.RequestException() response = MockResponse() response.status_code = 201 + if url.startswith('http://other'): + response.status_code = 404 return response def requests_delete(url): if url.startswith('http://local'): - raise Exception() + raise requests.exceptions.RequestException() response = MockResponse() response.status_code = 204 return response @@ -65,7 +76,7 @@ class TestEtcd(unittest.TestCase): self.assertRaises(Exception, self.etcd.get_client_path, '', 2) def test_put_client_path(self): - self.assertFalse(self.etcd.put_client_path('')) + self.assertRaises(EtcdError, self.etcd.put_client_path, '') def test_delete_client_path(self): self.assertFalse(self.etcd.delete_client_path('')) @@ -77,6 +88,29 @@ class TestEtcd(unittest.TestCase): self.assertIsInstance(cluster, Cluster) self.etcd.base_client_url = self.etcd.base_client_url.replace('remote', 'other') self.etcd.get_cluster() + self.etcd.base_client_url = self.etcd.base_client_url.replace('other', 'noleader') + self.etcd.get_cluster() def test_current_leader(self): self.assertRaises(CurrentLeaderError, self.etcd.current_leader) + + def test_touch_member(self): + self.assertFalse(self.etcd.touch_member('', '')) + + def test_take_leader(self): + self.assertFalse(self.etcd.take_leader('')) + + def test_attempt_to_acquire_leader(self): + self.assertFalse(self.etcd.attempt_to_acquire_leader('')) + + def test_update_leader(self): + self.etcd.base_client_url = self.etcd.base_client_url.replace('local', 'remote') + self.assertTrue(self.etcd.update_leader(MockPostgresql())) + self.etcd.base_client_url = self.etcd.base_client_url.replace('remote', 'other') + self.assertFalse(self.etcd.update_leader(MockPostgresql())) + + def test_race(self): + self.assertFalse(self.etcd.race('', '')) + + def test_delete_member(self): + self.assertFalse(self.etcd.delete_member('')) diff --git a/tests/test_governor.py b/tests/test_governor.py index 05fc5a1e..3d939670 100644 --- a/tests/test_governor.py +++ b/tests/test_governor.py @@ -6,11 +6,16 @@ import sys import time import yaml -from governor import Governor, main, sigchld_handler +from governor import Governor, main, sigchld_handler, sigterm_handler from test_ha import true, false from test_postgresql import Postgresql, os_system, psycopg2_connect from test_etcd import requests_get, requests_put, requests_delete +if sys.hexversion >= 0x03000000: + import http.server as BaseHTTPServer +else: + import BaseHTTPServer + def nop(*args, **kwargs): pass @@ -20,6 +25,10 @@ def os_waitpid(a, b): return (0, 0) +def time_sleep(_): + raise Exception() + + class TestGovernor(unittest.TestCase): def __init__(self, method_name='runTest'): @@ -28,25 +37,37 @@ class TestGovernor(unittest.TestCase): super(TestGovernor, self).__init__(method_name) def set_up(self): + self.touched = False os.system = os_system psycopg2.connect = psycopg2_connect requests.get = requests_get requests.put = requests_put requests.delete = requests_delete time.sleep = nop - Governor.run = nop self.write_pg_hba = Postgresql.write_pg_hba self.write_recovery_conf = Postgresql.write_recovery_conf Postgresql.write_pg_hba = nop Postgresql.write_recovery_conf = nop + BaseHTTPServer.HTTPServer.__init__ = nop def tear_down(self): Postgresql.write_pg_hba = self.write_pg_hba Postgresql.write_recovery_conf = self.write_recovery_conf + def test_sigterm_handler(self): + self.assertRaises(SystemExit, sigterm_handler, None, None) + def test_governor_main(self): - sys.argv = ['governor.py', 'postgres0.yml'] main() + sys.argv = ['governor.py', 'postgres0.yml'] + time.sleep = time_sleep + self.assertRaises(Exception, main) + + def touch_member(self): + if not self.touched: + self.touched = True + return False + return True def test_governor_initialize(self): with open('postgres0.yml', 'r') as f: @@ -61,7 +82,12 @@ class TestGovernor(unittest.TestCase): g.etcd.race = false g.initialize() g.postgresql.data_directory_empty = false + g.touch_member = self.touch_member g.initialize() + g.postgresql.data_directory_empty = true + time.sleep = time_sleep + g.postgresql.sync_from_leader = false + self.assertRaises(Exception, g.initialize) def test_sigchld_handler(self): sigchld_handler(None, None) diff --git a/tests/test_ha.py b/tests/test_ha.py index cabdf7d5..65d7bd6a 100644 --- a/tests/test_ha.py +++ b/tests/test_ha.py @@ -72,7 +72,8 @@ class TestHa(unittest.TestCase): self.p = MockPostgresql() self.e = Etcd({'ttl': 30, 'host': 'remotehost', 'scope': 'test'}) self.ha = Ha(self.p, self.e) - self.ha.cluster = Cluster(None, None, []) + self.ha.load_cluster_from_etcd() + self.ha.cluster = Cluster(False, None, None, []) self.ha.load_cluster_from_etcd = nop def test_start_as_slave(self): diff --git a/tests/test_postgresql.py b/tests/test_postgresql.py index 6e70dd89..473b9243 100644 --- a/tests/test_postgresql.py +++ b/tests/test_postgresql.py @@ -22,37 +22,36 @@ def false(*args, **kwargs): return False +def xlog_position(): + return 1 + + class MockCursor: - def __init__(self, server): + def __init__(self): + self.closed = False self.current = 0 self.results = [] - self.server = server def execute(self, sql, *params): if sql.startswith('blabla'): raise psycopg2.OperationalError() + elif sql.startswith('InterfaceError'): + raise psycopg2.InterfaceError() elif sql.startswith('SELECT slot_name'): self.results = [('blabla',), ('foobar',)] elif sql.startswith('SELECT pg_current_xlog_location()'): - self.results = [(0, )] - elif sql.startswith('SELECT %s - (pg_last_xlog_replay_location()'): - self.results = [(0, )] - elif sql.startswith('SELECT pg_last_xlog_replay_location()'): - self.results = [(0, )] + self.results = [(0,)] + elif sql.startswith('SELECT pg_is_in_recovery(), %s'): + if params[0][0] != 0: + raise psycopg2.OperationalError() + self.results = [(False, 0)] + elif sql.startswith('SELECT CASE WHEN pg_is_in_recovery()'): + self.results = [(0,)] elif sql.startswith('SELECT pg_is_in_recovery()'): - self.results = [( - self.server.mock_values['mock_recovery'], - None, - None, - None, - None, - None, - None, - None, - None, - None, - )] + self.results = [(False, )] + elif sql.startswith('SELECT to_char(pg_postmaster_start_time'): + self.results = [('', True, '', '', '', False)] else: self.results = [( None, @@ -82,14 +81,13 @@ class MockConnect: def __init__(self): self.autocommit = False - self.mock_values = {'mock_recovery': False} + self.closed = 0 def cursor(self): - return MockCursor(self) + return MockCursor() def close(self): - if not self.autocommit: - raise psycopg2.OperationalError() + pass def psycopg2_connect(*args, **kwargs): @@ -160,17 +158,25 @@ class TestPostgresql(unittest.TestCase): def test_query(self): self.p.query('select 1') - self.p.conn.autocommit = False + self.assertRaises(psycopg2.InterfaceError, self.p.query, 'InterfaceError') + self.assertRaises(psycopg2.OperationalError, self.p.query, 'blabla') + self.p._connection.closed = 2 + self.assertRaises(psycopg2.OperationalError, self.p.query, 'blabla') + self.p._connection.closed = 2 + self.p.disconnect = false self.assertRaises(psycopg2.OperationalError, self.p.query, 'blabla') - self.p.query('select %s', 1) def test_is_healthiest_node(self): leader = Member('leader', 'postgres://replicator:rep-pass@127.0.0.1:5435/postgres', 28) me = Member('test0', 'postgres://replicator:rep-pass@127.0.0.1:5434/postgres', 28) other = Member('test1', 'postgres://replicator:rep-pass@127.0.0.1:5433/postgres', 28) - cluster = Cluster(leader, 0, [leader, me, other]) + cluster = Cluster(True, leader, 0, [me, other, leader]) self.assertTrue(self.p.is_healthiest_node(cluster)) - self.p.config['maximum_lag_on_failover'] = -1 + self.p.is_leader = false + self.assertFalse(self.p.is_healthiest_node(cluster)) + self.p.xlog_position = xlog_position + self.assertTrue(self.p.is_healthiest_node(cluster)) + self.p.config['maximum_lag_on_failover'] = -2 self.assertFalse(self.p.is_healthiest_node(cluster)) def test_is_leader(self): diff --git a/tests/test_statuspage.py b/tests/test_statuspage.py deleted file mode 100644 index dcebd651..00000000 --- a/tests/test_statuspage.py +++ /dev/null @@ -1,38 +0,0 @@ -#!/usr/bin/env python -# -*- coding: utf-8 -*- - -import unittest -import sys - -from helpers.statuspage import StatusPage -from test_postgresql import MockConnect - -if sys.hexversion >= 0x03000000: - from io import BytesIO as IO -else: - from StringIO import StringIO as IO - - -class TestStatusPage(unittest.TestCase): - - def test_do_GET(self): - for mock_recovery in [True, False]: - for page in [b'GET /pg_master', b'GET /pg_slave', b'GET /pg_status', b'GET /not_found']: - self.http_server = MockServer(('0.0.0.0', 8888), StatusPage, page, mock_recovery) - - -class MockRequest(object): - - def __init__(self, path): - self.path = path - - def makefile(self, *args, **kwargs): - return IO(self.path) - - -class MockServer(object): - - def __init__(self, ip_port, Handler, path, mock_recovery=False): - self.postgresql = MockConnect() - self.postgresql.mock_values['mock_recovery'] = mock_recovery - Handler(MockRequest(path), ip_port, self)