diff --git a/.travis.yml b/.travis.yml new file mode 100644 index 00000000..a0659429 --- /dev/null +++ b/.travis.yml @@ -0,0 +1,12 @@ +language: python +python: + - "2.7" + - "3.3" + - "3.4" +install: + - pip install -r requirements.txt + - pip install coveralls +script: + - python setup.py test +after_success: + - coveralls diff --git a/governor.py b/governor.py index d80df64b..66967510 100755 --- a/governor.py +++ b/governor.py @@ -84,7 +84,8 @@ def main(): 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_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 diff --git a/helpers/etcd.py b/helpers/etcd.py index 29552a2b..4ea9ba12 100644 --- a/helpers/etcd.py +++ b/helpers/etcd.py @@ -47,7 +47,7 @@ class Etcd: def put_client_path(self, path, **data): try: response = requests.put(self.client_url(path), data=data) - return response.status_code in [200, 201] + return response.status_code in [200, 201, 202, 204] except: logger.exception('PUT %s data=%s', path, data) return False @@ -55,7 +55,7 @@ class Etcd: def delete_client_path(self, path): try: response = requests.delete(self.client_url(path)) - return response.status_code == 204 + return response.status_code in [200, 202, 204] except: logger.exception('DELETE %s', path) return False diff --git a/helpers/postgresql.py b/helpers/postgresql.py index 14bfbeb7..e1041d71 100644 --- a/helpers/postgresql.py +++ b/helpers/postgresql.py @@ -1,6 +1,8 @@ import logging import os import psycopg2 +import shutil +import subprocess import sys import time @@ -42,7 +44,13 @@ class Postgresql: self.superuser = config['superuser'] self.admin = config['admin'] 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._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 @@ -87,12 +95,9 @@ class Postgresql: return not os.path.exists(self.data_dir) or os.listdir(self.data_dir) == [] def initialize(self): - if os.system(self._pg_ctl + ' initdb -o --encoding=UTF8') == 0: - self.write_pg_hba() - - return True - - return False + ret = os.system(self._pg_ctl + ' initdb -o --encoding=UTF8') == 0 + ret and self.write_pg_hba() + return ret def sync_from_leader(self, leader): r = parseurl(leader.address) @@ -104,11 +109,100 @@ class Postgresql: try: os.environ['PGPASSFILE'] = pgpass - return os.system('pg_basebackup -R -D {data_dir} --host={hostname} --port={port} -U {username}'.format( - data_dir=self.data_dir, **r)) == 0 + return self.create_replica(leader.address, r) == 0 finally: os.environ.pop('PGPASSFILE') + def create_replica(self, master_connurl, master_connection): + """ creates a new replica using either pg_basebackup or WAL-E """ + if self.should_use_s3_to_create_replica(master_connurl): + result = self.create_replica_with_s3() + # if restore from the backup on S3 failed - try with the pg_basebackup + if result == 0: + return result + 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( + data_dir=self.data_dir, **master_connection)) + + def create_replica_with_s3(self): + if not self.wal_e or not self.wal_e_path: + return 1 + + ret = os.system(self.wal_e_path + ' backup-fetch {} LATEST'.format(self.data_dir)) + self.restore_configuration_files() + return ret + + def should_use_s3_to_create_replica(self, master_connurl): + """ 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 + + threshold_megabytes = self.wal_e.get('threshold_megabytes', 10240) + threshold_backup_size_percentage = self.wal_e.get('threshold_backup_size_percentage', 30) + + try: + latest_backup = subprocess.check_output(self.wal_e_path.split() + ['backup-list', '--detail', 'LATEST']) + # name last_modified expanded_size_bytes wal_segment_backup_start wal_segment_offset_backup_start wal_segment_backup_stop wal_segment_offset_backup_stop + # base_00000001000000000000007F_00000040 2015-05-18T10:13:25.000Z + # 20310671 00000001000000000000007F 00000040 + # 00000001000000000000007F 00000240 + backup_strings = latest_backup.splitlines() if latest_backup else () + if len(backup_strings) != 2: + return False + + names = backup_strings[0].split() + vals = backup_strings[1].split() + if (len(names) != len(vals)) or (len(names) != 7): + return False + + backup_info = dict(zip(names, vals)) + except subprocess.CalledProcessError as e: + logger.error("could not query wal-e latest backup: {}".format(e)) + return False + + try: + backup_size = backup_info['expanded_size_bytes'] + backup_start_segment = backup_info['wal_segment_backup_start'] + backup_start_offset = backup_info['wal_segment_offset_backup_start'] + except Exception as e: + logger.error("unable to get some of S3 backup parameters: {}".format(e)) + return False + + # WAL filename is XXXXXXXXYYYYYYYY000000ZZ, where X - timeline, Y - LSN logical log file, + # ZZ - 2 high digits of LSN offset. The rest of the offset is the provided decimal offset, + # that we have to convert to hex and 'prepend' to the high offset digits. + + lsn_segment = backup_start_segment[8:16] + # first 2 characters of the result are 0x and the last one is L + lsn_offset = hex((long(backup_start_segment[16:32], 16) << 24) + long(backup_start_offset))[2:-1] + + # construct the LSN from the segment and offset + backup_start_lsn = '{}/{}'.format(lsn_segment, lsn_offset) + + conn = None + cursor = None + 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.autocommit = True + cursor = conn.cursor() + cursor.execute("SELECT pg_xlog_location_diff(pg_current_xlog_location(), %s)", (backup_start_lsn,)) + diff_in_bytes = long(cursor.fetchone()[0]) + except psycopg2.Error as e: + logger.error('could not determine difference with the master location: {}'.format(e)) + return False + finally: + cursor and cursor.close() + conn and conn.close() + + # if the size of the accumulated WAL segments is more than a certan percentage of the backup size + # or exceeds the pre-determined size - pg_basebackup is chosen instead. + return (diff_in_bytes < long(threshold_megabytes) * 1048576) and\ + (diff_in_bytes < long(backup_size) * float(threshold_backup_size_percentage) / 100) + def is_leader(self): return not self.query('SELECT pg_is_in_recovery()').fetchone()[0] @@ -128,6 +222,7 @@ class Postgresql: ret = os.system(self._pg_ctl + ' start -o "{}"'.format(self.server_options())) == 0 ret and self.load_replication_slots() + self.save_configuration_files() return ret def stop(self): @@ -222,6 +317,22 @@ primary_conninfo = '{}' self.write_recovery_conf(leader) self.restart() + def save_configuration_files(self): + """ + copy postgresql.conf to postgresql.conf.backup to preserve it in the WAL-e backup. + see http://comments.gmane.org/gmane.comp.db.postgresql.wal-e/239 + """ + for f in self.configuration_to_save: + shutil.copy(f, f + '.backup') + + def restore_configuration_files(self): + """ restore a previously saved postgresql.conf """ + try: + for f in self.configuration_to_save: + shutil.copy(f + '.backup', f) + except Exception as e: + logger.error("unable to restore configuration from WAL-E backup: {}".format(e)) + def promote(self): return os.system(self._pg_ctl + ' promote') == 0 diff --git a/helpers/statuspage.py b/helpers/statuspage.py index 2819beeb..4e97f17a 100644 --- a/helpers/statuspage.py +++ b/helpers/statuspage.py @@ -1,14 +1,19 @@ #!/usr/bin/env python # -*- coding: utf-8 -*- -from BaseHTTPServer import BaseHTTPRequestHandler, HTTPServer 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' + content_type = 'text/plain' if self.path == '/pg_master': if not self.pg_is_in_recovery(): response, content = 200, 'I am currently a master' @@ -67,7 +72,7 @@ if __name__ == '__main__': logging.basicConfig(format='%(levelname)-6s %(asctime)s - %(message)s', level=logging.DEBUG) logging.debug('Starting as a standalone application') - # # Create a dummy configuration to be able to use the Postgresql class + # Create a dummy configuration to be able to use the Postgresql class from postgresql import Postgresql postgres_config = { 'name': 'dummy', diff --git a/postgres0.yml b/postgres0.yml index da630242..263cd693 100644 --- a/postgres0.yml +++ b/postgres0.yml @@ -19,6 +19,10 @@ postgresql: admin: username: admin password: admin + wal_e: + env_dir: /home/postgres/etc/wal-e.d/env + threshold_megabytes: 10240 + threshold_backup_size_percentage: 30 #recovery_conf: #restore_command: cp ../wal_archive/%f %p parameters: diff --git a/postgres1.yml b/postgres1.yml index e8f958b5..6d0082b2 100644 --- a/postgres1.yml +++ b/postgres1.yml @@ -21,6 +21,10 @@ postgresql: password: admin #recovery_conf: #restore_command: cp ../wal_archive/%f %p + wal_e: + env_dir: /home/postgres/etc/wal-e.d/env + threshold_megabytes: 10240 + threshold_backup_size_percentage: 30 parameters: archive_mode: "on" wal_level: hot_standby diff --git a/setup.py b/setup.py new file mode 100644 index 00000000..cafa4c23 --- /dev/null +++ b/setup.py @@ -0,0 +1,135 @@ +#!/usr/bin/env python + +""" + Setup file for governor +""" + +import sys +import os +import inspect + +import setuptools +from setuptools.command.test import test as TestCommand +from setuptools import setup + +if sys.version_info < (2, 7, 0): + sys.stderr.write('FATAL: governor needs to be run with Python 2.7+\n') + sys.exit(1) + +__location__ = os.path.join(os.getcwd(), os.path.dirname(inspect.getfile(inspect.currentframe()))) + + +NAME = 'governor' +MAIN_PACKAGE = 'governor.py' +HELPERS = 'helpers' +VERSION = '0.1' +DESCRIPTION = 'A Template for PostgreSQL HA with etcd' +LICENSE = 'The MIT License' + +COVERAGE_XML = True +COVERAGE_HTML = False +JUNIT_XML = True + +# Add here all kinds of additional classifiers as defined under +# https://pypi.python.org/pypi?%3Aaction=list_classifiers +CLASSIFIERS = [ + 'Development Status :: 4 - Beta', + 'Environment :: Console', + 'Intended Audience :: Developers', + 'Intended Audience :: System Administrators', + 'License :: OSI Approved :: The MIT License', + 'Operating System :: POSIX :: Linux', + 'Programming Language :: Python', + 'Programming Language :: Python :: 2.7', + 'Programming Language :: Python :: 3.3', + 'Programming Language :: Python :: 3.4', + 'Programming Language :: Python :: Implementation :: CPython', +] + + +class PyTest(TestCommand): + + user_options = [('cov=', None, 'Run coverage'), ('cov-xml=', None, 'Generate junit xml report'), ('cov-html=', + None, 'Generate junit html report'), ('junitxml=', None, 'Generate xml of test results')] + + def initialize_options(self): + TestCommand.initialize_options(self) + self.cov_xml = False + self.cov_html = False + self.junitxml = None + + def finalize_options(self): + TestCommand.finalize_options(self) + if self.cov_xml or self.cov_html: + self.cov = ['--cov', MAIN_PACKAGE, '--cov', HELPERS, '--cov-report', 'term-missing'] + if self.cov_xml: + self.cov.extend(['--cov-report', 'xml']) + if self.cov_html: + self.cov.extend(['--cov-report', 'html']) + if self.junitxml is not None: + self.junitxml = ['--junitxml', self.junitxml] + + def run_tests(self): + try: + import pytest + except: + raise RuntimeError('py.test is not installed, run: pip install pytest') + params = {'args': self.test_args} + if self.cov: + params['args'] += self.cov + params['plugins'] = ['cov'] + if self.junitxml: + params['args'] += self.junitxml + params['args'] += ['--doctest-modules', HELPERS, '-s'] + errno = pytest.main(**params) + sys.exit(errno) + + +def get_install_requirements(path): + content = open(os.path.join(__location__, path)).read() + return [req for req in content.split('\n') if req != ''] + + +def read(fname): + return open(os.path.join(__location__, fname)).read() + + +def setup_package(): + # Assemble additional setup commands + cmdclass = {} + cmdclass['test'] = PyTest + + # Some helper variables + version = os.getenv('GO_PIPELINE_LABEL', VERSION) + + install_reqs = get_install_requirements('requirements.txt') + + command_options = {'test': {'test_suite': ('setup.py', 'tests')}} + if JUNIT_XML: + command_options['test']['junitxml'] = 'setup.py', 'junit.xml' + if COVERAGE_XML: + command_options['test']['cov_xml'] = 'setup.py', True + if COVERAGE_HTML: + command_options['test']['cov_html'] = 'setup.py', True + + setup( + name=NAME, + version=version, + description=DESCRIPTION, + license=LICENSE, + keywords='etcd governor postgresql postgres ha', + long_description=read('README.md'), + classifiers=CLASSIFIERS, + test_suite='tests', + packages=setuptools.find_packages(exclude=['tests', 'tests.*']), + package_data={MAIN_PACKAGE: ["*.json"]}, + install_requires=install_reqs, + setup_requires=['six', 'flake8'], + cmdclass=cmdclass, + tests_require=['pytest-cov', 'pytest'], + command_options=command_options, + ) + + +if __name__ == '__main__': + setup_package() diff --git a/tests/test_etcd.py b/tests/test_etcd.py new file mode 100644 index 00000000..c358e89e --- /dev/null +++ b/tests/test_etcd.py @@ -0,0 +1,82 @@ +import unittest +import requests +import time +import json + +from helpers.etcd import Cluster, Etcd +from helpers.errors import EtcdError, CurrentLeaderError + + +class MockResponse: + + def __init__(self): + self.status_code = 200 + self.content = '{}' + + def json(self): + return json.loads(self.content) + + +def requests_get(url, **kwargs): + if url.startswith('http://local'): + raise Exception() + 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 + return response + + +def requests_put(url, **kwargs): + if url.startswith('http://local'): + raise Exception() + response = MockResponse() + response.status_code = 201 + return response + + +def requests_delete(url): + if url.startswith('http://local'): + raise Exception() + response = MockResponse() + response.status_code = 204 + return response + + +def time_sleep(_): + pass + + +class TestEtcd(unittest.TestCase): + + def __init__(self, method_name='runTest'): + self.setUp = self.set_up + super(TestEtcd, self).__init__(method_name) + + def set_up(self): + requests.get = requests_get + requests.put = requests_put + requests.delete = requests_delete + time.sleep = time_sleep + self.etcd = Etcd({'ttl': 30, 'host': 'localhost', 'scope': 'test'}) + + def test_get_client_path(self): + self.assertRaises(Exception, self.etcd.get_client_path, '', 2) + + def test_put_client_path(self): + self.assertFalse(self.etcd.put_client_path('')) + + def test_delete_client_path(self): + self.assertFalse(self.etcd.delete_client_path('')) + + def test_get_cluster(self): + self.assertRaises(EtcdError, self.etcd.get_cluster) + self.etcd.base_client_url = self.etcd.base_client_url.replace('local', 'remote') + cluster = self.etcd.get_cluster() + self.assertIsInstance(cluster, Cluster) + self.etcd.base_client_url = self.etcd.base_client_url.replace('remote', 'other') + self.etcd.get_cluster() + + def test_current_leader(self): + self.assertRaises(CurrentLeaderError, self.etcd.current_leader) diff --git a/tests/test_governor.py b/tests/test_governor.py new file mode 100644 index 00000000..05fc5a1e --- /dev/null +++ b/tests/test_governor.py @@ -0,0 +1,69 @@ +import os +import psycopg2 +import unittest +import requests +import sys +import time +import yaml + +from governor import Governor, main, sigchld_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 + + +def nop(*args, **kwargs): + pass + + +def os_waitpid(a, b): + return (0, 0) + + +class TestGovernor(unittest.TestCase): + + def __init__(self, method_name='runTest'): + self.setUp = self.set_up + self.tearDown = self.tear_down + super(TestGovernor, self).__init__(method_name) + + def set_up(self): + 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 + + def tear_down(self): + Postgresql.write_pg_hba = self.write_pg_hba + Postgresql.write_recovery_conf = self.write_recovery_conf + + def test_governor_main(self): + sys.argv = ['governor.py', 'postgres0.yml'] + main() + + def test_governor_initialize(self): + with open('postgres0.yml', 'r') as f: + config = yaml.load(f) + g = Governor(config) + g.postgresql.should_use_s3_to_create_replica = false + g.etcd.base_client_url = 'http://remote' + g.etcd.client_url + g.postgresql.data_directory_empty = true + g.etcd.race = true + g.initialize() + g.etcd.race = false + g.initialize() + g.postgresql.data_directory_empty = false + g.initialize() + + def test_sigchld_handler(self): + sigchld_handler(None, None) + os.waitpid = os_waitpid + sigchld_handler(None, None) diff --git a/tests/test_ha.py b/tests/test_ha.py new file mode 100644 index 00000000..741b674f --- /dev/null +++ b/tests/test_ha.py @@ -0,0 +1,121 @@ +import unittest +import requests + +from helpers.etcd import Etcd +from helpers.ha import Ha +from test_etcd import requests_get, requests_put, requests_delete + + +def true(*args, **kwargs): + return True + + +def false(*args, **kwargs): + return False + + +class MockPostgresql: + + def __init__(self): + self.name = 'postgresql0' + + def is_healthy(self): + return True + + def write_recovery_conf(self, _): + return True + + def start(self): + return True + + def is_healthiest_node(self, members): + return True + + def is_leader(self): + return True + + def promote(self): + return True + + def demote(self, _): + return True + + def follow_the_leader(self, _): + return True + + def create_replication_slots(self, _): + return True + + def last_operation(self): + return 0 + + +class TestHa(unittest.TestCase): + + def __init__(self, method_name='runTest'): + self.setUp = self.set_up + super(TestHa, self).__init__(method_name) + + def set_up(self): + requests.get = requests_get + requests.put = requests_put + requests.delete = requests_delete + self.p = MockPostgresql() + self.e = Etcd({'ttl': 30, 'host': 'remotehost', 'scope': 'test'}) + self.ha = Ha(self.p, self.e) + + def test_start_as_slave(self): + self.p.is_healthy = false + self.assertEquals(self.ha.run_cycle(), 'started as a secondary') + + def test_start_as_readonly(self): + self.p.is_leader = self.p.is_healthy = false + self.ha.has_lock = true + self.assertEquals(self.ha.run_cycle(), 'promoted self to leader because i had the session lock') + + def test_acquire_lock_as_master(self): + self.ha.is_unlocked = true + self.assertEquals(self.ha.run_cycle(), 'acquired session lock as a leader') + + def test_promoted_by_acquiring_lock(self): + self.ha.is_unlocked = true + self.p.is_leader = false + self.assertEquals(self.ha.run_cycle(), 'promoted self to leader by acquiring session lock') + + def test_demote_after_failing_to_obtain_lock(self): + self.ha.is_unlocked = true + self.ha.acquire_lock = false + self.assertEquals(self.ha.run_cycle(), 'demoted self due after trying and failing to obtain lock') + + def test_follow_new_leader_after_failing_to_obtain_lock(self): + self.ha.is_unlocked = true + self.ha.acquire_lock = false + self.p.is_leader = false + self.assertEquals(self.ha.run_cycle(), 'following new leader after trying and failing to obtain lock') + + def test_demote_because_not_healthiest(self): + self.ha.is_unlocked = true + self.p.is_healthiest_node = false + self.assertEquals(self.ha.run_cycle(), 'demoting self because i am not the healthiest node') + + def test_follow_new_leader_because_not_healthiest(self): + self.ha.is_unlocked = true + self.p.is_healthiest_node = false + self.p.is_leader = false + self.assertEquals(self.ha.run_cycle(), 'following a different leader because i am not the healthiest node') + + def test_promote_because_have_lock(self): + self.ha.has_lock = true + self.p.is_leader = false + self.assertEquals(self.ha.run_cycle(), 'promoted self to leader because i had the session lock') + + def test_leader_with_lock(self): + self.ha.has_lock = true + self.assertEquals(self.ha.run_cycle(), 'no action. i am the leader with the lock') + + def test_demote_because_not_having_lock(self): + self.assertEquals(self.ha.run_cycle(), 'demoting self because i do not have the lock and i was a leader') + + def test_follow_the_leader(self): + self.p.is_leader = false + self.assertEquals(self.ha.run_cycle(), 'no action. i am a secondary and i am following a leader') diff --git a/tests/test_postgresql.py b/tests/test_postgresql.py new file mode 100644 index 00000000..7fa5dcf7 --- /dev/null +++ b/tests/test_postgresql.py @@ -0,0 +1,158 @@ +import os +import psycopg2 +import unittest +import shutil + +from helpers.etcd import Cluster, Member +from helpers.postgresql import Postgresql + + +def nop(*args, **kwargs): + pass + + +def os_system(cmd): + return 0 + + +def false(*args, **kwargs): + return False + + +class MockCursor: + + def __init__(self): + self.current = 0 + self.results = [] + + def execute(self, sql, *params): + if sql.startswith('blabla'): + raise psycopg2.OperationalError() + 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,)] + elif sql.startswith('SELECT pg_is_in_recovery()'): + self.results = [(False, )] + else: + self.results = [] + + def fetchone(self): + return self.results[0] + + def close(self): + pass + + def __iter__(self): + for i in self.results: + yield i + + +class MockConnect: + + def __init__(self): + self.autocommit = False + + def cursor(self): + return MockCursor() + + def close(self): + if not self.autocommit: + raise psycopg2.OperationalError() + + +def psycopg2_connect(*args, **kwargs): + + return MockConnect() + + +def is_running(): + return False + + +class TestPostgresql(unittest.TestCase): + + def __init__(self, method_name='runTest'): + self.setUp = self.set_up + self.tearDown = self.tear_down + super(TestPostgresql, self).__init__(method_name) + + def set_up(self): + os.system = os_system + shutil.copy = nop + self.p = Postgresql({'name': 'test0', 'data_dir': 'data/test0', 'listen': '127.0.0.1, 127.0.0.2:5432', + 'connect_address': '127.0.0.2:5432', 'superuser': {'password': ''}, + 'admin': {'username': 'admin', 'password': 'admin'}, 'replication': { + 'username': 'replicator', 'password': 'rep-pass', 'network': '127.0.0.1/32'}, + 'parameters': {'foo': 'bar'}, 'recovery_conf': {'foo': 'bar'}}) + psycopg2.connect = psycopg2_connect + if not os.path.exists(self.p.data_dir): + os.makedirs(self.p.data_dir) + self.leader = Member('leader', 'postgres://replicator:rep-pass@127.0.0.1:5434/postgres', 28) + + def tear_down(self): + shutil.rmtree('data') + + def test_data_directory_empty(self): + self.assertTrue(self.p.data_directory_empty()) + + def test_initialize(self): + self.assertTrue(self.p.initialize()) + self.assertTrue(os.path.exists(os.path.join(self.p.data_dir, 'pg_hba.conf'))) + + def test_start(self): + self.assertFalse(self.p.start()) + self.p.is_running = is_running + with open(os.path.join(self.p.data_dir, 'postmaster.pid'), 'w'): + pass + self.assertTrue(self.p.start()) + + def test_sync_from_leader(self): + self.assertTrue(self.p.sync_from_leader(self.leader)) + + def test_follow_the_leader(self): + self.p.demote(self.leader) + self.p.follow_the_leader(None) + self.p.demote(self.leader) + self.p.follow_the_leader(self.leader) + self.p.follow_the_leader(Member('leader', 'postgres://replicator:rep-pass@127.0.0.1:5435/postgres', 28)) + + def test_create_replication_slots(self): + self.p.start() + self.p.create_replication_slots('qaz') + + def test_query(self): + self.p.query('select 1') + self.p.conn.autocommit = 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]) + self.assertTrue(self.p.is_healthiest_node(cluster)) + self.p.config['maximum_lag_on_failover'] = -1 + self.assertFalse(self.p.is_healthiest_node(cluster)) + + def test_is_leader(self): + self.assertTrue(self.p.is_leader()) + + def test_reload(self): + self.assertTrue(self.p.reload()) + + def test_is_healthy(self): + self.assertTrue(self.p.is_healthy()) + self.p.is_running = is_running + self.assertFalse(self.p.is_healthy()) + + def test_promote(self): + self.assertTrue(self.p.promote()) + + def test_last_operation(self): + self.assertEquals(self.p.last_operation(), 0) diff --git a/tox.ini b/tox.ini new file mode 100644 index 00000000..aa079ec5 --- /dev/null +++ b/tox.ini @@ -0,0 +1,2 @@ +[flake8] +max-line-length=120