From a9f86aa1952bc99825a84c40b12a3d00d88517a6 Mon Sep 17 00:00:00 2001
From: Alexander Kukushkin <cyberdemn@gmail.com>
Date: Fri, 15 Jan 2021 14:30:48 +0100
Subject: [PATCH] Add compatibility with python-consul2 (#1812)

the good old python-consul is not maintained for a few years in a row, therefore someone forked under a different name, but package files are installed into the same location as for the old.

The API of both modules is mostly compatible therefore it wasn't hard to add the support of both modules in Patroni.

Taking into account that python-consul is not a direct requirement for Patroni, but extra, now the end-user has a choice what to install.

Close https://github.com/zalando/patroni/issues/1810
---
 patroni/dcs/consul.py | 25 +++++++++++++++++--------
 tests/test_consul.py  |  5 +++--
 2 files changed, 20 insertions(+), 10 deletions(-)

diff --git a/patroni/dcs/consul.py b/patroni/dcs/consul.py
index ba6799af..a9acb08b 100644
--- a/patroni/dcs/consul.py
+++ b/patroni/dcs/consul.py
@@ -8,6 +8,7 @@ import ssl
 import time
 import urllib3
 
+from collections import namedtuple
 from consul import ConsulException, NotFound, base
 from urllib3.exceptions import HTTPError
 from six.moves.urllib.parse import urlencode, urlparse, quote
@@ -36,6 +37,9 @@ class InvalidSession(ConsulException):
     """invalid session"""
 
 
+Response = namedtuple('Response', 'code,headers,body,content')
+
+
 class HTTPClient(object):
 
     def __init__(self, host='127.0.0.1', port=8500, token=None, scheme='http', verify=True, cert=None, ca_cert=None):
@@ -71,16 +75,17 @@ class HTTPClient(object):
 
     @staticmethod
     def response(response):
-        data = response.data.decode('utf-8')
+        content = response.data
+        body = content.decode('utf-8')
         if response.status == 500:
-            msg = '{0} {1}'.format(response.status, data)
-            if data.startswith('Invalid Session TTL'):
+            msg = '{0} {1}'.format(response.status, body)
+            if body.startswith('Invalid Session TTL'):
                 raise InvalidSessionTTL(msg)
-            elif data.startswith('invalid session'):
+            elif body.startswith('invalid session'):
                 raise InvalidSession(msg)
             else:
                 raise ConsulInternalError(msg)
-        return base.Response(response.status, response.headers, data)
+        return Response(response.status, response.headers, body, content)
 
     def uri(self, path, params=None):
         return '{0}{1}{2}'.format(self.base_uri, path, params and '?' + urlencode(params) or '')
@@ -89,7 +94,7 @@ class HTTPClient(object):
         if method not in ('get', 'post', 'put', 'delete'):
             raise AttributeError("HTTPClient instance has no attribute '{0}'".format(method))
 
-        def wrapper(callback, path, params=None, data=''):
+        def wrapper(callback, path, params=None, data='', headers=None):
             # python-consul doesn't allow to specify ttl smaller then 10 seconds
             # because session_ttl_min defaults to 10s, so we have to do this ugly dirty hack...
             if method == 'put' and path == '/v1/session/create':
@@ -110,8 +115,9 @@ class HTTPClient(object):
                 kwargs['timeout'] = timeout + max(timeout/15.0, 1)
             else:
                 kwargs['timeout'] = self._read_timeout
+            kwargs['headers'] = (headers or {}).copy()
+            kwargs['headers'].update(urllib3.make_headers(user_agent=USER_AGENT))
             token = params.pop('token', self.token) if isinstance(params, dict) else self.token
-            kwargs['headers'] = urllib3.make_headers(user_agent=USER_AGENT)
             if token:
                 kwargs['headers']['X-Consul-Token'] = token
             return callback(self.response(self.http.request(method.upper(), self.uri(path, params), **kwargs)))
@@ -126,7 +132,7 @@ class ConsulClient(base.Consul):
         self.token = kwargs.get('token')
         super(ConsulClient, self).__init__(*args, **kwargs)
 
-    def connect(self, *args, **kwargs):
+    def http_connect(self, *args, **kwargs):
         kwargs.update(dict(zip(['host', 'port', 'scheme', 'verify'], args)))
         if self._cert:
             kwargs['cert'] = self._cert
@@ -136,6 +142,9 @@ class ConsulClient(base.Consul):
             kwargs['token'] = self.token
         return HTTPClient(**kwargs)
 
+    def connect(self, *args, **kwargs):
+        return self.http_connect(*args, **kwargs)
+
     def reload_config(self, config):
         self.http.token = self.token = config.get('token')
         self.consistency = config.get('consistency', 'default')
diff --git a/tests/test_consul.py b/tests/test_consul.py
index 6d0c4732..a38a2869 100644
--- a/tests/test_consul.py
+++ b/tests/test_consul.py
@@ -4,7 +4,7 @@ import unittest
 from consul import ConsulException, NotFound
 from mock import Mock, patch
 from patroni.dcs.consul import AbstractDCS, Cluster, Consul, ConsulInternalError, \
-                                ConsulError, HTTPClient, InvalidSessionTTL, InvalidSession
+                                ConsulError, ConsulClient, HTTPClient, InvalidSessionTTL, InvalidSession
 from . import SleepException
 
 
@@ -41,7 +41,8 @@ def kv_get(self, key, **kwargs):
 class TestHTTPClient(unittest.TestCase):
 
     def setUp(self):
-        self.client = HTTPClient('127.0.0.1', '8500', 'http', False)
+        c = ConsulClient()
+        self.client = c.http
         self.client.http.request = Mock()
 
     def test_get(self):