From 499061918d075cd663b129755d1574e29181e8bc Mon Sep 17 00:00:00 2001
From: Alexander Kukushkin <alexander.kukushkin@zalando.de>
Date: Thu, 21 Apr 2016 16:58:44 +0200
Subject: [PATCH] Implement noloadbalance support

Mostly this tag is necessary to give a hint to load balancer
auto-configuration tool that node should not be included into
LB configuration.
In addition to that Patroni also should not return status_code=200
for a health check if the tag is present and value is not `False`.
---
 patroni/__init__.py   | 4 ++++
 patroni/api.py        | 6 +++---
 tests/test_api.py     | 1 +
 tests/test_patroni.py | 4 ++++
 4 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/patroni/__init__.py b/patroni/__init__.py
index 916c0d3d..630fa672 100644
--- a/patroni/__init__.py
+++ b/patroni/__init__.py
@@ -28,6 +28,10 @@ class Patroni(object):
         self.ha = Ha(self)
         self.next_run = time.time()
 
+    @property
+    def noloadbalance(self):
+        return self.tags.get('noloadbalance', False)
+
     @property
     def nofailover(self):
         return self.tags.get('nofailover', False)
diff --git a/patroni/api.py b/patroni/api.py
index f5bf9ce1..81f7a78b 100644
--- a/patroni/api.py
+++ b/patroni/api.py
@@ -74,12 +74,12 @@ class RestApiHandler(BaseHTTPRequestHandler):
                 status_code = 503
             elif response['role'] == 'master':  # running as master but without leader lock!!!!
                 status_code = 503
-            elif response['role'] in path:
-                status_code = 200
+            elif response['role'] in path:  # response['role'] != 'master'
+                status_code = 503 if patroni.noloadbalance else 200
             else:
                 status_code = 503
         elif 'role' in response and response['role'] in path:
-            status_code = 200
+            status_code = 503 if response['role'] != 'master' and patroni.noloadbalance else 200
         elif patroni.ha.restart_scheduled() and patroni.postgresql.role == 'master' and 'master' in path:
             # exceptional case for master node when the postgres is being restarted via API
             status_code = 200
diff --git a/tests/test_api.py b/tests/test_api.py
index a19da05b..308ddf8e 100644
--- a/tests/test_api.py
+++ b/tests/test_api.py
@@ -54,6 +54,7 @@ class MockPatroni(object):
     dcs = Mock()
     tags = {}
     version = '0.00'
+    noloadbalance = Mock(return_value=False)
 
 
 class MockRequest(object):
diff --git a/tests/test_patroni.py b/tests/test_patroni.py
index 8c6a9718..d46d23ed 100644
--- a/tests/test_patroni.py
+++ b/tests/test_patroni.py
@@ -77,6 +77,10 @@ class TestPatroni(unittest.TestCase):
         self.p.next_run = time.time() - self.p.nap_time - 1
         self.p.schedule_next_run()
 
+    def test_noloadbalance(self):
+        self.p.tags['noloadbalance'] = True
+        self.assertTrue(self.p.noloadbalance)
+
     def test_nofailover(self):
         self.p.tags['nofailover'] = True
         self.assertTrue(self.p.nofailover)