Merge branch 'master' of github.com:CyberDem0n/governor into features/refactoring

Conflicts:
	helpers/postgresql.py

governor.py

@@ -1,6 +1,10 @@
 #!/usr/bin/env python
 
-import sys, os, yaml, time, urllib2, atexit
+import sys
+import yaml
+import time
+import urllib2
+import atexit
 import logging
 
 from helpers.etcd import Etcd
@@ -31,6 +35,8 @@ postgresql = Postgresql(config["postgresql"], aws_host_address)
 ha = Ha(postgresql, etcd)
 
 # stop postgresql on script exit
+
+
 def stop_postgresql():
     postgresql.stop()
 atexit.register(stop_postgresql)
@@ -68,7 +74,7 @@ if postgresql.data_directory_empty():
             else:
                 time.sleep(5)
 else:
-    postgresql.write_recovery_conf({"address": "postgres://169.0.0.1:5432"})
+    postgresql.write_recovery_conf(None)
     postgresql.start()
 
 while True:
@@ -76,9 +82,10 @@ while True:
 
     # create replication slots
     if postgresql.is_leader():
-        for node in etcd.get_client_path("/members?recursive=true")["node"]["nodes"]:
-            member = node["key"].split('/')[-1]
-            if member != postgresql.name:
-                postgresql.query("DO LANGUAGE plpgsql $$DECLARE somevar VARCHAR; BEGIN SELECT slot_name INTO somevar FROM pg_replication_slots WHERE slot_name = '%(slot)s' LIMIT 1; IF NOT FOUND THEN PERFORM pg_create_physical_replication_slot('%(slot)s'); END IF; END$$;" % {"slot": member})
+        for member in etcd.members():
+            if member['hostname'] != postgresql.name:
+                postgresql.query("""SELECT pg_create_physical_replication_slot(%s)
+                                     WHERE NOT EXISTS (SELECT 1 FROM pg_replication_slots
+                                     WHERE slot_name = %s)""", member['hostname'], member['hostname'])
 
     time.sleep(config["loop_wait"])

helpers/errors.py

@@ -1,13 +1,15 @@
-class CurrentLeaderError(Exception):
+class EtcdError(Exception):
+
     def __init__(self, value):
         self.value = value
 
     def __str__(self):
         return repr(self.value)
 
-class HealthiestMemberError(Exception):
-    def __init__(self, value):
-        self.value = value
 
-    def __str__(self):
-        return repr(self.value)
+class CurrentLeaderError(EtcdError):
+    pass
+
+
+class HealthiestMemberError(EtcdError):
+    pass

helpers/etcd.py

@@ -1,15 +1,19 @@
-import urllib2, json, os, time
+import urllib2
+import json
+import time
 import logging
+
+from helpers.errors import CurrentLeaderError
 from urllib import urlencode
-import helpers.errors
 
 logger = logging.getLogger(__name__)
 
+
 class Etcd:
+
     def __init__(self, config):
-        self.scope = config["scope"]
-        self.host = config["host"]
-        self.ttl = config["ttl"]
+        self.ttl = config['ttl']
+        self.base_client_url = 'http://{host}/v2/keys/service/{scope}'.format(**config)
 
     def get_client_path(self, path, max_attempts=1):
         attempts = 0
@@ -22,7 +26,7 @@ class Etcd:
             except (urllib2.HTTPError, urllib2.URLError) as e:
                 attempts += 1
                 if attempts < max_attempts:
-                    logger.info("Failed to return %s, trying again. (%s of %s)" % (path, attempts, max_attempts))
+                    logger.info('Failed to return %s, trying again. (%s of %s)', path, attempts, max_attempts)
                     time.sleep(3)
                 else:
                     raise e
@@ -38,18 +42,18 @@ class Etcd:
         opener.open(request)
 
     def client_url(self, path):
-        return "http://%s/v2/keys/service/%s%s" % (self.host, self.scope, path)
+        return self.base_client_url + path
 
     def current_leader(self):
         try:
-            hostname = self.get_client_path("/leader")["node"]["value"]
-            address = self.get_client_path("/members/%s" % hostname)["node"]["value"]
+            hostname = self.get_client_path('/leader')['node']['value']
+            address = self.get_client_path('/members/' + hostname)['node']['value']
 
-            return {"hostname": hostname, "address": address}
+            return {'hostname': hostname, 'address': address}
         except urllib2.HTTPError as e:
             if e.code == 404:
                 return None
-            raise helpers.errors.CurrentLeaderError("Etcd is not responding properly")
+            raise CurrentLeaderError("Etcd is not responding properly")
 
     def members(self):
         try:
@@ -63,20 +67,20 @@ class Etcd:
         except urllib2.HTTPError as e:
             if e.code == 404:
                 return None
-            raise helpers.errors.CurrentLeaderError("Etcd is not responding properly")
+            raise CurrentLeaderError("Etcd is not responding properly")
 
     def touch_member(self, member, connection_string):
-        self.put_client_path("/members/%s" % member, {"value": connection_string})
+        self.put_client_path('/members/' + member, {"value": connection_string})
 
     def take_leader(self, value):
-        return self.put_client_path("/leader", {"value": value, "ttl": self.ttl}) == None
+        return self.put_client_path("/leader", {"value": value, "ttl": self.ttl}) is None
 
     def attempt_to_acquire_leader(self, value):
         try:
-            return self.put_client_path("/leader", {"value": value, "ttl": self.ttl, "prevExist": False}) == None
+            return self.put_client_path("/leader", {"value": value, "ttl": self.ttl, "prevExist": False}) is None
         except urllib2.HTTPError as e:
             if e.code == 412:
-                logger.info("Could not take out TTL lock: %s" % e)
+                logger.info('Could not take out TTL lock: %s', e)
             return False
 
     def update_leader(self, value):
@@ -98,15 +102,15 @@ class Etcd:
             return False
 
     def am_i_leader(self, value):
-        #try:
-           reponse = self.get_client_path("/leader")
-           logger.info("Lock owner: %s; I am %s" % (reponse["node"]["value"], value))
-           return reponse["node"]["value"] == value
-        #except Exception as e:
-            #return False
+        # try:
+            reponse = self.get_client_path("/leader")
+            logger.info('Lock owner: %s; I am %s', reponse["node"]["value"], value)
+            return reponse["node"]["value"] == value
+        # except Exception as e:
+            # return False
 
     def race(self, path, value):
         try:
-            return self.put_client_path(path, {"prevExist": False, "value": value}) == None
+            return self.put_client_path(path, {"prevExist": False, "value": value}) is None
         except urllib2.HTTPError:
             return False

helpers/ha.py

@@ -1,10 +1,9 @@
-import sys, time, re, urllib2, json, psycopg2
-import logging
-from base64 import b64decode
-
-import helpers.errors
-
 import inspect
+import logging
+import time
+
+from helpers.errors import CurrentLeaderError, HealthiestMemberError
+from psycopg2 import OperationalError
 
 logger = logging.getLogger(__name__)
 
@@ -13,7 +12,9 @@ def lineno():
     """Returns the current line number in our program."""
     return inspect.currentframe().f_back.f_lineno
 
+
 class Ha:
+
     def __init__(self, state_handler, etcd):
         self.state_handler = state_handler
         self.etcd = etcd
@@ -31,7 +32,7 @@ class Ha:
         return self.etcd.am_i_leader(self.state_handler.name)
 
     def fetch_current_leader(self):
-            return self.etcd.current_leader()
+        return self.etcd.current_leader()
 
     def run_cycle(self):
         try:
@@ -42,7 +43,6 @@ class Ha:
                             if not self.state_handler.is_leader():
                                 self.state_handler.promote()
                                 return "promoted self to leader by acquiring session lock"
-
                             return "acquired session lock as a leader"
                         else:
                             if self.state_handler.is_leader():
@@ -58,7 +58,6 @@ class Ha:
                         else:
                             self.state_handler.follow_the_leader(self.fetch_current_leader())
                             return "following a different leader because i am not the healthiest node"
-
                 else:
                     if self.has_lock():
                         self.update_lock()
@@ -81,11 +80,11 @@ class Ha:
                     self.state_handler.start()
                     return "postgresql was stopped.  starting again."
                 return "no action.  not healthy enough to do anything."
-        except helpers.errors.CurrentLeaderError:
+        except CurrentLeaderError:
             logger.error("failed to fetch current leader from etcd")
-        except psycopg2.OperationalError:
+        except OperationalError:
             logger.error("Error communicating with Postgresql.  Will try again.")
-        except helpers.errors.HealthiestMemberError:
+        except HealthiestMemberError:
             logger.error("failed to determine healthiest member fromt etcd")
 
     def run(self):

helpers/postgresql.py

@@ -1,33 +1,48 @@
-import os, psycopg2, re, time
 import logging
-
-from urlparse import urlparse
+import os
+import psycopg2
+import re
+import time
+import urlparse
 
 
 logger = logging.getLogger(__name__)
 
 
+def parseurl(url):
+    r = urlparse.urlparse(url)
+    return {
+        'hostname': r.hostname,
+        'port': r.port or 5432,
+        'username': r.username,
+        'password': r.password,
+    }
+
+
 class Postgresql:
 
     def __init__(self, config, aws_host_address=None):
-        self.name = config["name"]
-        self.host, self.port = config["listen"].split(":")
-        self.data_dir = config["data_dir"]
-        self.replication = config["replication"]
-        self.superuser = config.get('superuser')
-        self.admin = config.get('admin')
+        self.name = config['name']
+        self.host, self.port = config['listen'].split(':')
+        self.data_dir = config['data_dir']
+        self.replication = config['replication']
+        self.superuser = config['superuser']
+        self.admin = config['admin']
+        self.recovery_conf = os.path.join(self.data_dir, 'recovery.conf')
+        self._pg_ctl = 'pg_ctl -w -D ' + self.data_dir
 
         self.config = config
 
-        self.cursor_holder = None
         connection_host = aws_host_address or self.host
-        self.connection_string = "postgres://%s:%s@%s:%s/postgres" % (self.replication["username"], self.replication["password"], connection_host, self.port)
+        self.connection_string = 'postgres://{username}:{password}@{host}:{port}/postgres'.format(
+            host=connection_host, port=self.port, **self.replication)
 
         self.conn = None
+        self.cursor_holder = None
 
     def cursor(self):
         if not self.cursor_holder:
-            self.conn = psycopg2.connect("postgres://%s:%s/postgres" % (self.host, self.port))
+            self.conn = psycopg2.connect('postgres://{}/postgres'.format(self.config['listen']))
             self.conn.autocommit = True
             self.cursor_holder = self.conn.cursor()
 
@@ -36,14 +51,14 @@ class Postgresql:
     def disconnect(self):
         try:
             self.conn.close()
-        except Exception as e:
-            logger.error("Error disconnecting: %s" % e)
+        except:
+            logger.exception('Error disconnecting')
 
-    def query(self, sql):
+    def query(self, sql, *params):
         max_attempts = 0
         while True:
             try:
-                self.cursor().execute(sql)
+                self.cursor().execute(sql, params)
                 break
             except psycopg2.OperationalError as e:
                 if self.conn:
@@ -59,7 +74,7 @@ class Postgresql:
         return not os.path.exists(self.data_dir) or os.listdir(self.data_dir) == []
 
     def initialize(self):
-        if os.system("initdb -D %s" % self.data_dir) == 0:
+        if os.system(self._pg_ctl + ' initdb') == 0:
             self.write_pg_hba()
 
             return True
@@ -67,71 +82,68 @@ class Postgresql:
         return False
 
     def sync_from_leader(self, leader):
-        leader = urlparse(leader["address"])
+        r = parseurl(leader['address'])
 
-        f = open("./pgpass", "w")
-        f.write("%(hostname)s:%(port)s:*:%(username)s:%(password)s\n" %
-                {"hostname": leader.hostname, "port": leader.port, "username": leader.username, "password": leader.password})
-        f.close()
+        pgpass = 'pgpass'
+        with open(pgpass, 'w') as f:
+            os.fchmod(f.fileno(), 0600)
+            f.write('{hostname}:{port}:*:{username}:{password}\n'.format(**r))
 
-        os.system("chmod 600 pgpass")
-
-        return os.system("PGPASSFILE=pgpass pg_basebackup -R -D %(data_dir)s --host=%(host)s --port=%(port)s -U %(username)s" %
-                {"data_dir": self.data_dir, "host": leader.hostname, "port": leader.port, "username": leader.username}) == 0
+        return os.system('PGPASSFILE={pgpass} pg_basebackup -R -D {data_dir} --host={hostname} --port={port} -U {username}'.format(
+            pgpass=pgpass, data_dir=self.data_dir, **r)) == 0
 
     def is_leader(self):
-        return not self.query("SELECT pg_is_in_recovery();").fetchone()[0]
+        return not self.query('SELECT pg_is_in_recovery()').fetchone()[0]
 
     def is_running(self):
-        return os.system("pg_ctl status -D %s > /dev/null" % self.data_dir) == 0
+        return os.system(self._pg_ctl + ' status > /dev/null') == 0
 
     def start(self):
         if self.is_running():
-            logger.error("Cannot start PostgreSQL because one is already running.")
+            logger.error('Cannot start PostgreSQL because one is already running.')
             return False
 
-        pid_path = "%s/postmaster.pid" % self.data_dir
+        pid_path = os.path.join(self.data_dir, 'postmaster.pid')
         if os.path.exists(pid_path):
             os.remove(pid_path)
-            logger.info("Removed %s" % pid_path)
+            logger.info('Removed %s', pid_path)
 
-        command_code = os.system("postgres -D %s %s &" % (self.data_dir, self.server_options()))
-        time.sleep(5)
-        return command_code != 0
+        return os.system(self._pg_ctl + ' start -o "{}"'.format(self.server_options())) == 0
 
     def stop(self):
-        return os.system("pg_ctl stop -w -D %s -m fast -w" % self.data_dir) != 0
+        return os.system(self._pg_ctl + ' stop') != 0
 
     def reload(self):
-        return os.system("pg_ctl reload -w -D %s" % self.data_dir) == 0
+        return os.system(self._pg_ctl + ' reload') == 0
 
     def restart(self):
-        return os.system("pg_ctl restart -w -D %s -m fast" % self.data_dir) == 0
+        return os.system(self._pg_ctl + ' restart -m fast') == 0
 
     def server_options(self):
-        options = "-c listen_addresses=%s -c port=%s" % (self.host, self.port)
-        for setting, value in self.config["parameters"].iteritems():
-            options += " -c \"%s=%s\"" % (setting, value)
+        options = '--listen_addresses={} --port={}'.format(self.host, self.port)
+        for setting, value in self.config['parameters'].iteritems():
+            options += " --{}='{}'".format(setting, value)
         return options
 
     def is_healthy(self):
         if not self.is_running():
-            logger.warning("Postgresql is not running.")
+            logger.warning('Postgresql is not running.')
             return False
 
         return True
 
     def is_healthiest_node(self, members):
         for member in members:
-            if member["hostname"] == self.name:
+            if member['hostname'] == self.name:
                 continue
             try:
-                member_conn = psycopg2.connect(member["address"])
+                member_conn = psycopg2.connect(member['address'])
                 member_conn.autocommit = True
                 member_cursor = member_conn.cursor()
-                member_cursor.execute("SELECT '%s'::pg_lsn - pg_last_xlog_replay_location() AS bytes;" % self.xlog_position())
+                member_cursor.execute(
+                    'SELECT %s::pg_lsn - pg_last_xlog_replay_location() AS bytes', (self.xlog_position(), ))
                 xlog_diff = member_cursor.fetchone()[0]
-                logger.info([self.name, member["hostname"], xlog_diff])
+                logger.info([self.name, member['hostname'], xlog_diff])
                 if xlog_diff < 0:
                     member_cursor.close()
                     return False
@@ -146,55 +158,72 @@ class Postgresql:
         return member
 
     def write_pg_hba(self):
-        f = open("%s/pg_hba.conf" % self.data_dir, "a")
-        f.write("host replication %(username)s %(network)s md5" %
-                {"username": self.replication["username"], "network": self.replication["network"]})
-        # allow TCP connections from the host's own address
-        f.write("\nhost postgres postgres samehost trust\n")
-        # allow TCP connections from the rest of the world with a password
-        f.write("\nhost all all 0.0.0.0/0 md5\n")
-        f.close()
+        with open(os.path.join(self.data_dir, 'pg_hba.conf'), 'a') as f:
+            f.write('host replication {username} {network} md5'.format(**self.replication))
+            # allow TCP connections from the host's own address
+            f.write("\nhost postgres postgres samehost trust\n")
+            # allow TCP connections from the rest of the world with a password
+            f.write("\nhost all all 0.0.0.0/0 md5\n")
+
+    @staticmethod
+    def primary_conninfo(leader_url):
+        r = parseurl(leader_url)
+        return 'user={username} password={password} host={hostname} port={port} sslmode=prefer sslcompression=1'.format(**r)
+
+    def check_recovery_conf(self, leader_hash):
+        if not os.path.isfile(self.recovery_conf):
+            return False
+
+        pattern = leader_hash and 'address' in leader_hash and self.primary_conninfo(leader_hash['address'])
+
+        with open(self.recovery_conf, 'r') as f:
+            for line in f:
+                if line.startswith('primary_conninfo'):
+                    if not pattern:
+                        return False
+                    return pattern in line
+
+        return not pattern
 
     def write_recovery_conf(self, leader_hash):
-        leader = urlparse(leader_hash["address"])
-
-        f = open("%s/recovery.conf" % self.data_dir, "w")
-        f.write("""
-standby_mode = 'on'
-primary_slot_name = '%(recovery_slot)s'
-primary_conninfo = 'user=%(user)s password=%(password)s host=%(hostname)s port=%(port)s sslmode=prefer sslcompression=1'
+        with open(self.recovery_conf, 'w') as f:
+            f.write("""standby_mode = 'on'
 recovery_target_timeline = 'latest'
-""" % {"recovery_slot": self.name, "user": leader.username, "password": leader.password, "hostname": leader.hostname, "port": leader.port})
-        if "recovery_conf" in self.config:
-            for name, value in self.config["recovery_conf"].iteritems():
-                f.write("%s = '%s'" % (name, value))
-        f.close()
+""")
+            if leader_hash and 'address' in leader_hash:
+                f.write("""
+primary_slot_name = '{}'
+primary_conninfo = '{}'
+""".format(self.name, self.primary_conninfo(leader_hash['address'])))
+                for name, value in self.config.get('recovery_conf', {}).iteritems():
+                    f.write("{} = '{}'\n".format(name, value))
 
     def follow_the_leader(self, leader_hash):
-        leader = urlparse(leader_hash["address"])
-        if os.system("grep 'host=%(hostname)s port=%(port)s' %(data_dir)s/recovery.conf > /dev/null" % {"hostname": leader.hostname, "port": leader.port, "data_dir": self.data_dir}) != 0:
-            self.write_recovery_conf(leader_hash)
-            self.restart()
-        return True
-
-    def promote(self):
-        return os.system("pg_ctl promote -w -D %s" % self.data_dir) == 0
-
-    def demote(self, leader):
-        self.write_recovery_conf(leader)
+        if self.check_recovery_conf(leader_hash):
+            return
+        self.write_recovery_conf(leader_hash)
         self.restart()
 
+    def promote(self):
+        return os.system(self._pg_ctl + ' promote') == 0
+
+    def demote(self, leader):
+        self.follow_the_leader(leader)
+
     def create_replication_user(self):
-        self.query("CREATE USER \"%s\" WITH REPLICATION ENCRYPTED PASSWORD '%s';" % (self.replication["username"], self.replication["password"]))
+        self.query('CREATE USER "{}" WITH REPLICATION ENCRYPTED PASSWORD %s'.format(
+            self.replication['username']), self.replication['password'])
 
     def create_connection_users(self):
         if self.superuser:
             if 'username' in self.superuser:
-                self.query("CREATE ROLE \"{0}\" LOGIN SUPERUSER PASSWORD '{1}';".format(self.superuser["username"], self.superuser["password"]))
+                self.query("CREATE ROLE \"{0}\" LOGIN SUPERUSER PASSWORD '{1}';".format(
+                    self.superuser["username"], self.superuser["password"]))
             else:
                 self.query("ALTER ROLE postgres PASSWORD '{0}';".format(self.superuser['password']))
         if self.admin:
-            self.query("CREATE ROLE \"{0}\" LOGIN CREATEDB CREATEROLE PASSWORD '{1}';".format(self.admin["username"], self.admin["password"]))
+            self.query("CREATE ROLE \"{0}\" LOGIN CREATEDB CREATEROLE PASSWORD '{1}';".format(
+                self.admin["username"], self.admin["password"]))
 
     def xlog_position(self):
-        return self.query("SELECT pg_last_xlog_replay_location();").fetchone()[0]
+        return self.query('SELECT pg_last_xlog_replay_location()').fetchone()[0]