add code to set tags on AWS objects (instances and EBS storage), as well as an ability for the postgres state_handler to call a callback when the state is changed, and use the new AWS code in a callback.

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
+import 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,115 @@
+import logging
+import re
+import requests
+from requests.exceptions import RequestException
+import types
+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 type(config['etcd']) == types.DictType:
+            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/meta-data/instance-id', timeout=0.1)
+            if r.ok:
+                self.instance_id = r.content.strip()
+                r = requests.get('http://169.254.169.254/latest/meta-data/placement/availability-zone', timeout=0.1)
+                if r.ok:
+                    # get the region from the availability zone, i.e. eu-west-1 from eu-west-1c
+                    m = re.match(r'(\w+-\w+-\d+)[a-z]+', r.content)
+                    if m:
+                        self.region = m.group(1)
+                        self.available = True
+        except RequestException:
+            logger.info("cannot query AWS meta-data")
+            pass
+
+    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': self.cluster_name, 'Role': role, 'Instance': self.instance_id}
+        try:
+            conn = boto.ec2.connect_to_region(self.region)
+            # get all volumes attached to the current instance
+            volumes = conn.get_all_volumes(filter={'attachment.instance-id': self.instance_id})
+            if volumes:
+                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)
+            instances = conn.get_all_reservations(instance_ids=[self.instance_id])
+            if instances:
+                conn.create_tag([instances[0].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):
+        self._tag_ec2(new_role)
+        self._tag_ebs('spilo_' + self.cluster_name, new_role)
+
+if __name__ == '__main__':
+    import yaml
+    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"
+"""
+    awsconnection = AWSConnection(yaml.load(config_string))
+    print "AWS available: {}, Cluster_name: {}".format(awsconnection.available, awsconnection.cluster_name)
+    if awsconnection.available:
+        print "AWS Region: {}, Instance_id: {}".format(awsconnection.region, awsconnection.instance_id)

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):
         # TODO: try to get unix_socket_directory from postmaster.pid
@@ -342,6 +343,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):
         """
@@ -361,6 +364,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):