Merge pull request #23 from zalando/feature/awstags

Set tags on EC2 instances and EBS storage

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

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'])

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

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):

requirements.txt

@@ -1,3 +1,4 @@
+boto
 PyYAML
 psycopg2
 requests

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"))

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:[email protected]: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
 
 

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)