diff --git a/README.md b/README.md index f90478ff..2582eb37 100644 --- a/README.md +++ b/README.md @@ -73,6 +73,12 @@ For an example file, see `postgres0.yml`. Below is an explanation of settings: * *username*: replication username, user will be created during initialization * *password*: replication password, user will be created during initialization * *network*: network setting for replication in pg_hba.conf + * *callbacks* callback scripts to run on certain actions. Patroni will pass current action, role and cluster name. See scripts/aws.py as an example on how to write them. + * *on_start*: a script to run when the cluster starts + * *on_stop*: a script to run when the cluster stops + * *on_restart*: a script to run when the cluster restarts + * *on_reload*: a script to run when configuration reload is triggered + * *on_role_change*: a script to run when the cluster is being promoted or demoted * *superuser* * *password*: password for postgres user. It would be set during initialization * *admin*: diff --git a/helpers/postgresql.py b/helpers/postgresql.py index bbf5ffe9..84a94346 100644 --- a/helpers/postgresql.py +++ b/helpers/postgresql.py @@ -1,6 +1,7 @@ import logging import os import psycopg2 +import shlex import shutil import subprocess import sys @@ -15,6 +16,12 @@ else: logger = logging.getLogger(__name__) +ACTION_ON_START = "on_start" +ACTION_ON_STOP = "on_stop" +ACTION_ON_RESTART = "on_restart" +ACTION_ON_RELOAD = "on_reload" +ACTION_ON_ROLE_CHANGE = "on_role_change" + def parseurl(url): r = urlparse(url) @@ -35,14 +42,16 @@ def parseurl(url): class Postgresql: - def __init__(self, config, on_change_callback=None): + def __init__(self, config): self.config = config self.name = config['name'] + self.scope = config['scope'] self.listen_addresses, self.port = config['listen'].split(':') self.data_dir = config['data_dir'] self.replication = config['replication'] self.superuser = config['superuser'] self.admin = config['admin'] + self.callback = config.get('callbacks', {}) 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')) @@ -65,7 +74,6 @@ class Postgresql: self._connection = None self._cursor_holder = None self.members = [] # list of already existing replication slots - self.on_change_callback = on_change_callback def get_local_address(self): listen_addresses = self.listen_addresses.split(',') @@ -230,9 +238,9 @@ class Postgresql: 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): + def is_leader(self, check_only=False): ret = not self.query('SELECT pg_is_in_recovery()').fetchone()[0] - if ret and self.is_promoted: + if ret and self.is_promoted and not check_only: self.delete_trigger_file() self.is_promoted = False return ret @@ -240,6 +248,26 @@ class Postgresql: def is_running(self): return subprocess.call(' '.join(self._pg_ctl) + ' status > /dev/null', shell=True) == 0 + def call_nowait(self, cb_name, is_leader=None): + """ pick a callback command and call it without waiting for it to finish """ + if not self.callback or cb_name not in self.callback: + return False + cmd = self.callback[cb_name] + if is_leader is None: + try: + is_leader = self.is_leader(check_only=True) + except psycopg2.OperationalError as e: + logger.warning("unable to perform {0} action, cannot obtain the cluster role: {1}".format(cb_name, e)) + return False + scope = self.scope + try: + role = "master" if is_leader else "replica" + subprocess.Popen(shlex.split(os.path.abspath(cmd))+[cb_name, role, scope]) + except Exception as e: + logger.warning("callback {0} {1} {2} {3} failed: {4}".format(os.path.abspath(cmd), cb_name, role, scope, e)) + return False + return True + def start(self): if self.is_running(): self.load_replication_slots() @@ -253,18 +281,37 @@ class Postgresql: ret = subprocess.call(self._pg_ctl + ['start', '-o', self.server_options()]) == 0 ret and self.load_replication_slots() self.save_configuration_files() - if self.on_change_callback: - self.on_change_callback('replica' if os.path.exists(self.recovery_conf) else 'master') + if ret and ACTION_ON_START in self.callback: + self.call_nowait(ACTION_ON_START) return ret def stop(self): - return subprocess.call(self._pg_ctl + ['stop', '-m', 'fast']) != 0 + try: + is_leader = self.is_leader(check_only=True) + except: + is_leader = None + pass + ret = subprocess.call(self._pg_ctl + ['stop', '-m', 'fast']) + if ret == 0 and ACTION_ON_STOP in self.callback: + self.call_nowait(ACTION_ON_STOP, is_leader=is_leader) + return ret == 0 def reload(self): - return subprocess.call(self._pg_ctl + ['reload']) == 0 + ret = subprocess.call(self._pg_ctl + ['reload']) + if ret == 0 and ACTION_ON_RELOAD in self.callback: + self.call_nowait(ACTION_ON_RELOAD) + return ret == 0 def restart(self): - return subprocess.call(self._pg_ctl + ['restart', '-m', 'fast']) == 0 + try: + is_leader = self.is_leader(check_only=True) + except: + is_leader = None + pass + ret = subprocess.call(self._pg_ctl + ['restart', '-m', 'fast']) + if ret == 0 and ACTION_ON_RESTART in self.callback: + self.call_nowait(ACTION_ON_RESTART, is_leader=is_leader) + return ret == 0 def server_options(self): options = "--listen_addresses='{}' --port={}".format(self.listen_addresses, self.port) @@ -351,8 +398,8 @@ primary_conninfo = '{}' if not self.check_recovery_conf(leader): self.write_recovery_conf(leader) self.restart() - if self.on_change_callback: - self.on_change_callback('replica') + if ACTION_ON_ROLE_CHANGE in self.callback: + self.call_nowait(ACTION_ON_ROLE_CHANGE) def save_configuration_files(self): """ @@ -372,8 +419,8 @@ primary_conninfo = '{}' def promote(self): self.is_promoted = subprocess.call(self._pg_ctl + ['promote']) == 0 - if self.on_change_callback: - self.on_change_callback('master') + if self.is_promoted and ACTION_ON_ROLE_CHANGE in self.callback: + self.call_nowait(ACTION_ON_ROLE_CHANGE) return self.is_promoted def demote(self, leader): diff --git a/patroni.py b/patroni.py index d90e5c64..6f06a056 100755 --- a/patroni.py +++ b/patroni.py @@ -6,7 +6,6 @@ import time import yaml from helpers.api import RestApiServer -from helpers.aws import AWSConnection from helpers.etcd import Etcd from helpers.ha import Ha from helpers.postgresql import Postgresql @@ -18,7 +17,6 @@ class Patroni: def __init__(self, config): self.nap_time = config['loop_wait'] - self.aws = AWSConnection(config) self.postgresql = Postgresql(config['postgresql']) self.ha = Ha(self.postgresql, self.get_dcs(self.postgresql.name, config)) host, port = config['restapi']['listen'].split(':') diff --git a/postgres0.yml b/postgres0.yml index 7b8527dc..293ebbc5 100644 --- a/postgres0.yml +++ b/postgres0.yml @@ -1,15 +1,16 @@ ttl: &ttl 30 loop_wait: &loop_wait 10 +scope: &scope batman restapi: listen: 127.0.0.1:8008 connect_address: 127.0.0.1:8008 etcd: - scope: batman + scope: *scope ttl: *ttl host: 127.0.0.1:4001 #discovery_srv: my-etcd.domain #zookeeper: -# scope: batman +# scope: *scope # session_timeout: *ttl # reconnect_timeout: *loop_wait # hosts: @@ -24,6 +25,7 @@ etcd: # - host3 postgresql: name: postgresql0 + scope: *scope listen: 127.0.0.1:5432 connect_address: 127.0.0.1:5432 data_dir: data/postgresql0 diff --git a/postgres1.yml b/postgres1.yml index 5be3fc25..6ef6b1c9 100644 --- a/postgres1.yml +++ b/postgres1.yml @@ -1,15 +1,16 @@ ttl: &ttl 30 loop_wait: &loop_wait 10 +scope: &scope batman restapi: listen: 127.0.0.1:8009 connect_address: 127.0.0.1:8009 etcd: - scope: batman + scope: *scope ttl: *ttl host: 127.0.0.1:4001 #discovery_srv: my-etcd.domain #zookeeper: -# scope: batman +# scope: *scope # session_timeout: *ttl # reconnect_timeout: *loop_wait # hosts: @@ -24,6 +25,7 @@ etcd: # - host3 postgresql: name: postgresql1 + scope: *scope listen: 127.0.0.1:5433 connect_address: 127.0.0.1:5433 data_dir: data/postgresql1 diff --git a/scripts/__init__.py b/scripts/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/helpers/aws.py b/scripts/aws.py old mode 100644 new mode 100755 similarity index 82% rename from helpers/aws.py rename to scripts/aws.py index 352e3617..b172a8c3 --- a/helpers/aws.py +++ b/scripts/aws.py @@ -1,24 +1,18 @@ +#!/usr/bin/python + import logging import requests from requests.exceptions import RequestException -import yaml +import sys import boto.ec2 logger = logging.getLogger(__name__) class AWSConnection: - - def __init__(self, config): + def __init__(self, cluster_name): self.available = False - self.config = config - - if 'cluster_name' in config: - self.cluster_name = config.get('cluster_name') - elif 'etcd' in config and isinstance(config['etcd'], dict): - self.cluster_name = config['etcd'].get('scope', 'unknown') - else: - self.cluster_name = 'unknown' + self.cluster_name = cluster_name if cluster_name is not None else 'unknown' try: # get the instance id r = requests.get('http://169.254.169.254/latest/dynamic/instance-identity/document', timeout=0.1) @@ -27,7 +21,7 @@ class AWSConnection: return if r.ok: try: - content = yaml.load(r.content) + content = r.json() self.instance_id = content['instanceId'] self.region = content['region'] except Exception as e: @@ -69,3 +63,10 @@ class AWSConnection: def on_role_change(self, new_role): ret = self._tag_ec2(new_role) return self._tag_ebs(new_role) and ret + + +if __name__ == '__main__': + if len(sys.argv) == 4 and sys.argv[1] in ('on_start', 'on_stop', 'on_role_change'): + AWSConnection(cluster_name=sys.argv[3]).on_role_change(sys.argv[2]) + else: + sys.exit("Usage: {0} action role name".format(sys.argv[0])) diff --git a/setup.py b/setup.py index 5ac4a89f..62d05a69 100644 --- a/setup.py +++ b/setup.py @@ -22,6 +22,7 @@ __location__ = os.path.join(os.getcwd(), os.path.dirname(inspect.getfile(inspect NAME = 'patroni' MAIN_PACKAGE = 'patroni.py' HELPERS = 'helpers' +SCRIPTS = 'scripts' VERSION = '0.1' DESCRIPTION = 'A Template for PostgreSQL HA with etcd' LICENSE = 'The MIT License' @@ -61,7 +62,8 @@ class PyTest(TestCommand): 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'] + self.cov = ['--cov', MAIN_PACKAGE, '--cov', HELPERS, '--cov', SCRIPTS, '--cov-report', + 'term-missing'] if self.cov_xml: self.cov.extend(['--cov-report', 'xml']) if self.cov_html: @@ -80,7 +82,7 @@ class PyTest(TestCommand): params['plugins'] = ['cov'] if self.junitxml: params['args'] += self.junitxml - params['args'] += ['--doctest-modules', HELPERS, '-s'] + params['args'] += ['--doctest-modules', HELPERS, '--doctest-modules', SCRIPTS, '-s'] errno = pytest.main(**params) sys.exit(errno) diff --git a/tests/test_aws.py b/tests/test_aws.py index 63b50734..84d495fe 100644 --- a/tests/test_aws.py +++ b/tests/test_aws.py @@ -2,9 +2,8 @@ import unittest import requests import boto.ec2 from collections import namedtuple -from helpers.aws import AWSConnection +from scripts.aws import AWSConnection from requests.exceptions import RequestException -import yaml class MockEc2Connection: @@ -24,6 +23,16 @@ class MockEc2Connection: return True +class MockResponse: + + def __init__(self, content): + self.content = content + self.ok = True + + def json(self): + return self.content + + class TestAWSConnection(unittest.TestCase): def __init__(self, method_name='runTest'): @@ -32,8 +41,8 @@ class TestAWSConnection(unittest.TestCase): def set_error(self): self.error = True - def set_ok(self): - self.error = False + def set_json_error(self): + self.json_error = True def boto_ec2_connect_to_region(self, region): return MockEc2Connection(self.error) @@ -43,50 +52,18 @@ class TestAWSConnection(unittest.TestCase): raise RequestException("foo") result = namedtuple('Request', 'ok content') result.ok = True - if url.split('/')[-1] == 'document': - result.content = '{\n "instanceId" : "012345",\n "region" : "eu-west-1"\n}' + if url.split('/')[-1] == 'document' and not self.json_error: + result = {"instanceId": "012345", "region": "eu-west-1"} else: - result.content = 'foo' - return result + result = 'foo' + return MockResponse(result) def setUp(self): self.error = False + self.json_error = False requests.get = self.requests_get boto.ec2.connect_to_region = self.boto_ec2_connect_to_region - self.config_string = """ -loop_wait: 10 -restapi: - listen: 0.0.0.0:8008 - connect_address: 127.0.0.1:5432 -etcd: - scope: test - ttl: 30 - host: 127.0.0.1:8080 -postgresql: - name: postgresql_foo - listen: 0.0.0.0:5432 - connect_address: 127.0.0.1:5432 - data_dir: /home/postgres/pgdata/data - replication: - username: standby - password: standby - network: 0.0.0.0/0 - superuser: - password: zalando - admin: - username: admin - password: admin - parameters: - archive_mode: "on" - wal_level: hot_standby - max_wal_senders: 5 - wal_keep_segments: 8 - archive_timeout: 1800s - max_replication_slots: 5 - hot_standby: "on" - ssl: "on" -""" - self.conn = AWSConnection(yaml.load(self.config_string)) + self.conn = AWSConnection('test') def test_aws_available(self): self.assertTrue(self.conn.aws_available()) @@ -98,11 +75,16 @@ postgresql: def test_non_aws(self): self.set_error() - conn = AWSConnection(yaml.load(self.config_string)) + conn = AWSConnection('test') self.assertFalse(conn.aws_available()) self.assertFalse(conn._tag_ebs('master')) self.assertFalse(conn._tag_ec2('master')) + def test_aws_bizare_response(self): + self.set_json_error() + conn = AWSConnection('test') + self.assertFalse(conn.aws_available()) + def test_aws_tag_ebs_error(self): self.set_error() self.assertFalse(self.conn._tag_ebs("master")) diff --git a/tests/test_postgresql.py b/tests/test_postgresql.py index 30041675..2ee6e32a 100644 --- a/tests/test_postgresql.py +++ b/tests/test_postgresql.py @@ -110,15 +110,19 @@ class TestPostgresql(unittest.TestCase): def set_up(self): subprocess.call = subprocess_call shutil.copy = nop - self.p = Postgresql({'name': 'test0', 'data_dir': 'data/test0', 'listen': '127.0.0.1, *:5432', - 'connect_address': '127.0.0.2:5432', + self.p = Postgresql({'name': 'test0', 'scope': 'batman', 'data_dir': 'data/test0', + 'listen': '127.0.0.1, *:5432', 'connect_address': '127.0.0.2:5432', 'pg_hba': ['hostssl all all 0.0.0.0/0 md5', 'host all all 0.0.0.0/0 md5'], - 'superuser': {'password': ''}, 'admin': {'username': 'admin', 'password': 'admin'}, + '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'}}, - on_change_callback=lambda state: True) + 'parameters': {'foo': 'bar'}, 'recovery_conf': {'foo': 'bar'}, + 'callbacks': {'on_start': '/usr/bin/true', 'on_stop': '/usr/bin/true', + 'on_restart': '/usr/bin/true', 'on_role_change': '/bin/true', + 'on_reload': '/usr/bin/true' + }}) psycopg2.connect = psycopg2_connect if not os.path.exists(self.p.data_dir): os.makedirs(self.p.data_dir) @@ -129,6 +133,9 @@ class TestPostgresql(unittest.TestCase): def tear_down(self): shutil.rmtree('data') + def mock_query(self, p): + raise psycopg2.OperationalError("not supported") + def test_data_directory_empty(self): self.assertTrue(self.p.data_directory_empty()) @@ -136,12 +143,13 @@ class TestPostgresql(unittest.TestCase): self.assertTrue(self.p.initialize()) self.assertTrue(os.path.exists(os.path.join(self.p.data_dir, 'pg_hba.conf'))) - def test_start(self): + def test_start_stop(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()) + self.assertTrue(self.p.stop()) def test_sync_from_leader(self): self.assertTrue(self.p.sync_from_leader(self.leader)) @@ -202,3 +210,11 @@ class TestPostgresql(unittest.TestCase): def test_last_operation(self): self.assertEquals(self.p.last_operation(), '0') + + def test_non_existing_callback(self): + self.assertFalse(self.p.call_nowait('foobar')) + + def test_is_leader_exception(self): + self.p.start() + self.p.query = self.mock_query + self.assertTrue(self.p.stop())