diff --git a/Dockerfile b/Dockerfile index 42a90964..226c1564 100644 --- a/Dockerfile +++ b/Dockerfile @@ -10,7 +10,7 @@ RUN apt-get update -y RUN apt-get upgrade -y ENV PGVERSION 9.4 -RUN apt-get install python python-psycopg2 python-yaml python-requests postgresql-${PGVERSION} -y +RUN apt-get install python python-psycopg2 python-yaml python-requests python-boto postgresql-${PGVERSION} -y ENV PATH /usr/lib/postgresql/${PGVERSION}/bin:$PATH @@ -19,7 +19,7 @@ ADD governor.py /governor/governor.py ADD helpers /governor/helpers ADD postgres0.yml /governor/ -ENV ETCDVERSION 2.0.10 +ENV ETCDVERSION 2.0.11 RUN curl -L https://github.com/coreos/etcd/releases/download/v${ETCDVERSION}/etcd-v${ETCDVERSION}-linux-amd64.tar.gz | tar xz -C /bin --strip=1 --wildcards --no-anchored etcd etcdctl ## Setting up a simple script that will serve as an entrypoint diff --git a/governor.py b/governor.py index 193ad53e..d405a9d0 100755 --- a/governor.py +++ b/governor.py @@ -10,6 +10,7 @@ from helpers.etcd import Etcd from helpers.postgresql import Postgresql from helpers.ha import Ha from helpers.utils import setup_signal_handlers, sleep +from helpers.aws import AWSConnection class Governor: @@ -17,7 +18,8 @@ class Governor: def __init__(self, config): self.nap_time = config['loop_wait'] self.etcd = Etcd(config['etcd']) - self.postgresql = Postgresql(config['postgresql']) + self.aws = AWSConnection(config) + self.postgresql = Postgresql(config['postgresql'], self.aws.on_role_change) self.ha = Ha(self.postgresql, self.etcd) host, port = config['restapi']['listen'].split(':') self.api = RestApiServer(self, config['restapi']) diff --git a/helpers/aws.py b/helpers/aws.py new file mode 100644 index 00000000..b47e4e0f --- /dev/null +++ b/helpers/aws.py @@ -0,0 +1,71 @@ +import logging +import re +import requests +from requests.exceptions import RequestException +import yaml +import boto.ec2 + +logger = logging.getLogger(__name__) + + +class AWSConnection: + def __init__(self, config): + 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' + try: + # get the instance id + r = requests.get('http://169.254.169.254/latest/dynamic/instance-identity/document', timeout=0.1) + except RequestException: + logger.info("cannot query AWS meta-data") + return + if r.ok: + try: + content = yaml.load(r.content) + self.instance_id = content['instanceId'] + self.region = content['region'] + except Exception as e: + logger.info('unable to fetch instance id and region from AWS meta-data: {}'.format(e)) + return + self.available = True + + def aws_available(self): + return self.available + + def _tag_ebs(self, role): + """ set tags, carrying the cluster name, instance role and instance id for the EBS storage """ + if not self.available: + return False + + tags = {'Name': 'spilo_'+self.cluster_name, 'Role': role, 'Instance': self.instance_id} + try: + conn = boto.ec2.connect_to_region(self.region) + volumes = conn.get_all_volumes(filters={'attachment.instance-id': self.instance_id}) + conn.create_tags([v.id for v in volumes], tags) + except Exception as e: + logger.info('could not set tags for EBS storage devices attached: {}'.format(e)) + return False + return True + + def _tag_ec2(self, role): + """ tag the current EC2 instance with a cluster role """ + if not self.available: + return False + tags = {'Role': role} + try: + conn = boto.ec2.connect_to_region(self.region) + conn.create_tags([self.instance_id], tags) + except Exception as e: + logger.info("could not set tags for EC2 instance {}: {}".format(self.instance_id, e)) + return False + return True + + def on_role_change(self, new_role): + ret = self._tag_ec2(new_role) + return self._tag_ebs(new_role) and ret diff --git a/helpers/postgresql.py b/helpers/postgresql.py index fe170e30..bcf90c59 100644 --- a/helpers/postgresql.py +++ b/helpers/postgresql.py @@ -34,7 +34,7 @@ def parseurl(url): class Postgresql: - def __init__(self, config): + def __init__(self, config, on_change_callback=None): self.config = config self.name = config['name'] self.listen_addresses, self.port = config['listen'].split(':') @@ -64,6 +64,7 @@ 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(',') @@ -250,6 +251,8 @@ 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') return ret def stop(self): @@ -346,6 +349,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') def save_configuration_files(self): """ @@ -365,6 +370,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') return self.is_promoted def demote(self, leader): diff --git a/requirements.txt b/requirements.txt index 2185e5c6..d0f5e512 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,4 @@ +boto PyYAML psycopg2 requests diff --git a/tests/test_aws.py b/tests/test_aws.py new file mode 100644 index 00000000..63b50734 --- /dev/null +++ b/tests/test_aws.py @@ -0,0 +1,112 @@ +import unittest +import requests +import boto.ec2 +from collections import namedtuple +from helpers.aws import AWSConnection +from requests.exceptions import RequestException +import yaml + + +class MockEc2Connection: + + def __init__(self, error=False): + self.error = error + + def get_all_volumes(self, filters): + if self.error: + raise Exception("get_all_volumes") + oid = namedtuple('Volume', 'id') + return [oid(id='a'), oid(id='b')] + + def create_tags(self, objects, tags): + if self.error or len(objects) == 0: + raise Exception("create_tags") + return True + + +class TestAWSConnection(unittest.TestCase): + + def __init__(self, method_name='runTest'): + super(TestAWSConnection, self).__init__(method_name) + + def set_error(self): + self.error = True + + def set_ok(self): + self.error = False + + def boto_ec2_connect_to_region(self, region): + return MockEc2Connection(self.error) + + def requests_get(self, url, **kwargs): + if self.error: + 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}' + else: + result.content = 'foo' + return result + + def setUp(self): + self.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)) + + def test_aws_available(self): + self.assertTrue(self.conn.aws_available()) + + def test_on_role_change(self): + self.assertTrue(self.conn._tag_ebs('master')) + self.assertTrue(self.conn._tag_ec2('master')) + self.assertTrue(self.conn.on_role_change('master')) + + def test_non_aws(self): + self.set_error() + conn = AWSConnection(yaml.load(self.config_string)) + self.assertFalse(conn.aws_available()) + self.assertFalse(conn._tag_ebs('master')) + self.assertFalse(conn._tag_ec2('master')) + + def test_aws_tag_ebs_error(self): + self.set_error() + self.assertFalse(self.conn._tag_ebs("master")) + + def test_aws_tag_ec2_error(self): + self.set_error() + self.assertFalse(self.conn._tag_ec2("master")) diff --git a/tests/test_etcd.py b/tests/test_etcd.py index a7ab93ee..ec0a19bd 100644 --- a/tests/test_etcd.py +++ b/tests/test_etcd.py @@ -12,6 +12,7 @@ class MockResponse: def __init__(self): self.status_code = 200 self.content = '{}' + self.ok = True def json(self): return json.loads(self.content) @@ -34,6 +35,9 @@ def requests_get(url, **kwargs): 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?application_name=http://127.0.0.1:8008/governor","expiration":"2015-05-15T09:11:09.611860899Z","ttl":30,"modifiedIndex":20730,"createdIndex":20730}],"modifiedIndex":1581,"createdIndex":1581}],"modifiedIndex":1581,"createdIndex":1581}}' + else: + response.status_code = 404 + response.ok = False return response diff --git a/tests/test_postgresql.py b/tests/test_postgresql.py index e69e5703..87e207a4 100644 --- a/tests/test_postgresql.py +++ b/tests/test_postgresql.py @@ -117,7 +117,8 @@ class TestPostgresql(unittest.TestCase): 'replication': {'username': 'replicator', 'password': 'rep-pass', 'network': '127.0.0.1/32'}, - 'parameters': {'foo': 'bar'}, 'recovery_conf': {'foo': 'bar'}}) + 'parameters': {'foo': 'bar'}, 'recovery_conf': {'foo': 'bar'}}, + on_change_callback=lambda state: True) psycopg2.connect = psycopg2_connect if not os.path.exists(self.p.data_dir): os.makedirs(self.p.data_dir)