From 8184efbf459d900cefaa47c271862ba25c0c3a03 Mon Sep 17 00:00:00 2001
From: Oleksii Kliukin <oleksii.kliukin@zalando.de>
Date: Wed, 29 Apr 2015 12:54:29 +0200
Subject: [PATCH] Add an option to use AWS instance address in a PostgreSQL
 connection string.

When PostgreSQL is running inside the docker container, it should provide
the address of the host it's running on and not the address from inside
of the container to the outside world.
---
 governor.py           | 14 +++++++++++++-
 helpers/postgresql.py |  5 +++--
 postgres0.yml         |  1 +
 postgres1.yml         |  1 +
 4 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/governor.py b/governor.py
index 57591274..a6717faf 100755
--- a/governor.py
+++ b/governor.py
@@ -7,6 +7,7 @@ from helpers.etcd import Etcd
 from helpers.postgresql import Postgresql
 from helpers.ha import Ha
 
+INSTANCE_METADATA_URL = "http://169.254.169.254/latest/meta-data/"
 
 logging.basicConfig(format='%(asctime)s %(levelname)s: %(message)s', level=logging.INFO)
 
@@ -14,8 +15,19 @@ f = open(sys.argv[1], "r")
 config = yaml.load(f.read())
 f.close()
 
+if config.get('aws_use_host_address', False):
+    # get host address of the AWS host via a call to
+    # http://169.254.169.254/latest/meta-data/local-ipv4
+    try:
+        aws_host_address = urllib2.urlopen(INSTANCE_METADATA_URL+"/local-ipv4").read()
+    except (urllib2.HTTPError, urllib2.URLError) as e:
+        logging.error("Error retrieiving IPv4 address from AWS instance: {0}".format(e))
+        aws_host_address = None
+else:
+    aws_host_address = None
+
 etcd = Etcd(config["etcd"])
-postgresql = Postgresql(config["postgresql"])
+postgresql = Postgresql(config["postgresql"], aws_host_address)
 ha = Ha(postgresql, etcd)
 
 # stop postgresql on script exit
diff --git a/helpers/postgresql.py b/helpers/postgresql.py
index 517f656b..ecc81b02 100644
--- a/helpers/postgresql.py
+++ b/helpers/postgresql.py
@@ -9,7 +9,7 @@ logger = logging.getLogger(__name__)
 
 class Postgresql:
 
-    def __init__(self, config):
+    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"]
@@ -18,7 +18,8 @@ class Postgresql:
         self.config = config
 
         self.cursor_holder = None
-        self.connection_string = "postgres://%s:%s@%s:%s/postgres" % (self.replication["username"], self.replication["password"], self.host, self.port)
+        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.conn = None
 
diff --git a/postgres0.yml b/postgres0.yml
index 78389db6..66d0b9eb 100644
--- a/postgres0.yml
+++ b/postgres0.yml
@@ -1,4 +1,5 @@
 loop_wait: 10
+aws_use_host_address: "on"
 etcd:
   scope: batman
   ttl: 30
diff --git a/postgres1.yml b/postgres1.yml
index 29373471..1332291b 100644
--- a/postgres1.yml
+++ b/postgres1.yml
@@ -1,4 +1,5 @@
 loop_wait: 10
+aws_use_host_address: "on"
 etcd:
   scope: batman
   ttl: 30