Find and load dcs class implementation dynamically

patroni/__init__.py

@@ -5,7 +5,7 @@ import time
 import yaml
 
 from patroni.api import RestApiServer
-from patroni.exceptions import PatroniException
+from patroni.dcs import get_dcs
 from patroni.ha import Ha
 from patroni.postgresql import Postgresql
 from patroni.utils import reap_children, set_ignore_sigterm, setup_signal_handlers
@@ -22,7 +22,7 @@ class Patroni(object):
         self.tags = {tag: value for tag, value in config.get('tags', {}).items()
                      if tag not in ('clonefrom', 'nofailover', 'noloadbalance') or value}
         self.postgresql = Postgresql(config['postgresql'])
-        self.dcs = self.get_dcs(self.postgresql.name, config)
+        self.dcs = get_dcs(self.postgresql.name, config)
         self.version = __version__
         self.api = RestApiServer(self, config['restapi'])
         self.ha = Ha(self)
@@ -40,19 +40,6 @@ class Patroni(object):
     def replicatefrom(self):
         return self.tags.get('replicatefrom')
 
-    @staticmethod
-    def get_dcs(name, config):
-        if 'etcd' in config:
-            from patroni.dcs.etcd import Etcd
-            return Etcd(name, config['etcd'])
-        if 'zookeeper' in config:
-            from patroni.dcs.zookeeper import ZooKeeper
-            return ZooKeeper(name, config['zookeeper'])
-        if 'consul' in config:
-            from patroni.dcs.consul import Consul
-            return Consul(name, config['consul'])
-        raise PatroniException('Can not find suitable configuration of distributed configuration store')
-
     def schedule_next_run(self):
         self.next_run += self.nap_time
         current_time = time.time()

patroni/ctl.py

@@ -16,7 +16,8 @@ import tzlocal
 import yaml
 
 from click import ClickException
-from patroni import Patroni, PatroniException
+from patroni.dcs import get_dcs as _get_dcs
+from patroni.exceptions import PatroniException
 from patroni.postgresql import parseurl
 from prettytable import PrettyTable
 from six.moves.urllib_parse import urlparse
@@ -93,10 +94,9 @@ def ctl(ctx):
 
 
 def get_dcs(config, scope):
-    for k in set(DCS_DEFAULTS.keys()) & set(config.keys()):
-        config[k].setdefault('scope', scope)
+    config.setdefault('scope', scope)
     try:
-        return Patroni.get_dcs(scope, config)
+        return _get_dcs(scope, config)
     except PatroniException as e:
         raise PatroniCtlException(str(e))
 

patroni/dcs/__init__.py

@@ -1,9 +1,13 @@
 import abc
 import dateutil
+import importlib
+import inspect
 import json
+import os
 import six
 
 from collections import namedtuple
+from patroni.exceptions import PatroniException
 from random import randint
 from six.moves.urllib_parse import urlparse, urlunparse, parse_qsl
 from threading import Event, Lock
@@ -26,6 +30,26 @@ def parse_connection_string(value):
     return conn_url, api_url
 
 
+def get_dcs(node_name, config):
+    available_implementations = []
+    for name in os.listdir(os.path.dirname(__file__)):
+        if name.endswith('.py') and not name.startswith('__'):  # find module
+            module = importlib.import_module(__package__ + '.' + name[:-3])
+            for name in dir(module):  # iterate through module content
+                if not name.startswith('__'):  # skip internal stuff
+                    value = getattr(module, name)
+                    name = name.lower()
+                    # try to find implementation of AbstractDCS interface
+                    if inspect.isclass(value) and issubclass(value, AbstractDCS):
+                        available_implementations.append(name)
+                        if name in config:  # which has configuration section in the config file
+                            # propagate some parameters
+                            config[name].update({p: config[p] for p in ('namespace', 'scope', 'ttl') if p in config})
+                            return value(node_name, config[name])
+    raise PatroniException("""Can not find suitable configuration of distributed configuration store
+Available implementations: """ + ', '.join(available_implementations))
+
+
 class Member(namedtuple('Member', 'index,name,session,data')):
 
     """Immutable object (namedtuple) which represents single member of PostgreSQL cluster.

tests/test_ctl.py

@@ -146,8 +146,6 @@ class TestCtl(unittest.TestCase):
 
     def test_get_dcs(self):
         self.assertRaises(PatroniCtlException, get_dcs, {'dummy': {}}, 'dummy')
-        with patch('patroni.Patroni.get_dcs', Mock(return_value=self.e)):
-            assert get_dcs({'etcd': {'host': 'none'}}, 'dummy').client_path('') == '/service/test/'
 
     @patch('psycopg2.connect', psycopg2_connect)
     @patch('patroni.ctl.query_member', Mock(return_value=([['mock column']], None)))

tests/test_ha.py

@@ -4,11 +4,10 @@ import datetime
 import pytz
 
 from mock import Mock, MagicMock, patch
-from patroni.dcs import Cluster, Failover, Leader, Member
+from patroni.dcs import Cluster, Failover, Leader, Member, get_dcs
 from patroni.exceptions import DCSError, PostgresException
 from patroni.ha import Ha
 from patroni.postgresql import Postgresql
-from patroni import Patroni
 from test_etcd import socket_getaddrinfo, etcd_read, etcd_write, requests_get
 
 
@@ -94,7 +93,7 @@ class TestHa(unittest.TestCase):
             self.p.set_role('replica')
             self.p.check_replication_lag = true
             self.p.can_create_replica_without_replication_connection = MagicMock(return_value=False)
-            self.e = Patroni.get_dcs('foo', {'etcd': {'ttl': 30, 'host': 'ok:2379', 'scope': 'test'}})
+            self.e = get_dcs('foo', {'etcd': {'ttl': 30, 'host': 'ok:2379', 'scope': 'test'}})
             self.ha = Ha(MockPatroni(self.p, self.e))
             self.ha._async_executor.run_async = run_async
             self.ha.old_cluster = self.e.get_cluster()

tests/test_patroni.py

@@ -8,13 +8,10 @@ import yaml
 from mock import Mock, patch
 from patroni.api import RestApiServer
 from patroni.async_executor import AsyncExecutor
-from patroni.dcs.consul import Consul
-from patroni.dcs.zookeeper import ZooKeeper
-from patroni import Patroni, PatroniException, main as _main
+from patroni import Patroni, main as _main
 from six.moves import BaseHTTPServer
 from test_etcd import SleepException, etcd_read, etcd_write
 from test_postgresql import Postgresql, psycopg2_connect
-from test_zookeeper import MockKazooClient
 
 
 @patch('time.sleep', Mock())
@@ -38,13 +35,6 @@ class TestPatroni(unittest.TestCase):
                 config = yaml.load(f)
                 self.p = Patroni(config)
 
-    @patch('patroni.dcs.zookeeper.KazooClient', MockKazooClient())
-    @patch.object(Consul, 'create_or_restore_session', Mock())
-    def test_get_dcs(self):
-        self.assertIsInstance(self.p.get_dcs('', {'zookeeper': {'scope': '', 'hosts': ''}}), ZooKeeper)
-        self.assertIsInstance(self.p.get_dcs('', {'consul': {'scope': '', 'hosts': '127.0.0.1:1'}}), Consul)
-        self.assertRaises(PatroniException, self.p.get_dcs, '', {})
-
     @patch('time.sleep', Mock(side_effect=SleepException))
     @patch.object(etcd.Client, 'delete', Mock())
     @patch.object(etcd.Client, 'machines')