Merge pull request jplana#136 from jplana/srv_record_corrected

lavagetto · lavagetto · commit 42aa30c8c80f · 2015-11-01T13:16:07.000+01:00
Add SRV record
diff --git a/.travis.yml b/.travis.yml
@@ -1,10 +1,10 @@
 language: python
 python:
   - "2.7"
-  - "3.3"
+  - "3.5"
 
 before_install:
-  - ./build_etcd.sh v2.0.10
+  - ./build_etcd.sh v2.2.0
   - pip install --upgrade setuptools
 
 # command to install dependencies
diff --git a/README.rst b/README.rst
@@ -41,6 +41,8 @@ Create a client object
     client = etcd.Client(port=4002)
     client = etcd.Client(host='127.0.0.1', port=4003)
     client = etcd.Client(host='127.0.0.1', port=4003, allow_redirect=False) # wont let you run sensitive commands on non-leader machines, default is true
+    # If you have defined a SRV record for _etcd._tcp.example.com pointing to the clients
+    client = etcd.Client(srv_domain='example.com', protocol="https")
     # create a client against https://api.example.com:443/etcd
     client = etcd.Client(host='api.example.com', protocol='https', port=443, version_prefix='/etcd')
 Write a key
diff --git a/buildout.cfg b/buildout.cfg
@@ -6,6 +6,7 @@ develop = .
 eggs = 
      urllib3==1.7.1
      pyOpenSSL==0.13.1
+     ${deps:extraeggs}
 
 [python]
 recipe = zc.recipe.egg
@@ -21,3 +22,12 @@ eggs = ${python:eggs}
 recipe = collective.recipe.sphinxbuilder
 source = ${buildout:directory}/docs-source
 build = ${buildout:directory}/docs
+
+
+[deps:python2]
+extraeggs =
+        dnspython==1.12.0
+
+[deps:python3]
+extraeggs =
+        dnspython3==1.12.0
diff --git a/setup.py b/setup.py
@@ -8,8 +8,15 @@
 
 version = '0.4.2'
 
+# Dnspython is two different packages depending on python version
+if sys.version_info.major == 2:
+    dns = 'dnspython'
+else:
+    dns = 'dnspython3'
+
 install_requires = [
-    'urllib3>=1.7.1'
+    'urllib3>=1.7.1',
+    dns
 ]
 
 test_requires = [
diff --git a/src/etcd/client.py b/src/etcd/client.py
@@ -18,6 +18,7 @@
 import urllib3.util
 import json
 import ssl
+import dns.resolver
 import etcd
 
 try:
@@ -46,6 +47,7 @@ def __init__(
             self,
             host='127.0.0.1',
             port=4001,
+            srv_domain=None,
             version_prefix='/v2',
             read_timeout=60,
             allow_redirect=True,
@@ -67,6 +69,8 @@ def __init__(
 
             port (int):  Port used to connect to etcd.
 
+            srv_domain (str): Domain to search the SRV record for cluster autodiscovery.
+
             version_prefix (str): Url or version prefix in etcd url (default=/v2).
 
             read_timeout (int):  max seconds to wait for a read.
@@ -98,8 +102,15 @@ def __init__(
                                       by host. By default this will use up to 10
                                       connections.
         """
-        _log.debug("New etcd client created for %s:%s%s",
-                  host, port, version_prefix)
+
+        # If a DNS record is provided, use it to get the hosts list
+        if srv_domain is not None:
+            try:
+                host = self._discover(srv_domain)
+            except Exception as e:
+                _log.error("Could not discover the etcd hosts from %s: %s",
+                           srv_domain, e)
+
         self._protocol = protocol
 
         def uri(protocol, host, port):
@@ -153,6 +164,8 @@ def uri(protocol, host, port):
 
         self.http = urllib3.PoolManager(num_pools=10, **kw)
 
+        _log.debug("New etcd client created for %s", self.base_uri)
+
         if self._allow_reconnect:
             # we need the set of servers in the cluster in order to try
             # reconnecting upon error. The cluster members will be
@@ -174,6 +187,18 @@ def uri(protocol, host, port):
             _log.debug("Machines cache initialised to %s",
                        self._machines_cache)
 
+    def _discover(self, domain):
+        srv_name = "_etcd._tcp.{}".format(domain)
+        answers = dns.resolver.query(srv_name, 'SRV')
+        hosts = []
+        for answer in answers:
+            hosts.append(
+                (answer.target.to_text(omit_final_dot=True), answer.port))
+        _log.debug("Found %s", hosts)
+        if not len(hosts):
+            raise ValueError("The SRV record is present but no host were found")
+        return tuple(hosts)
+
     @property
     def base_uri(self):
         """URI used by the client to connect to etcd."""
diff --git a/src/etcd/tests/unit/__init__.py b/src/etcd/tests/unit/__init__.py
@@ -27,8 +27,7 @@ def _prepare_response(self, s, d, cluster_id=None):
 
     def _mock_api(self, status, d, cluster_id=None):
         resp = self._prepare_response(status, d, cluster_id=cluster_id)
-        self.client.api_execute = mock.create_autospec(
-            self.client.api_execute, return_value=resp)
+        self.client.api_execute = mock.MagicMock(return_value=resp)
 
     def _mock_exception(self, exc, msg):
         self.client.api_execute = mock.Mock(side_effect=exc(msg))
diff --git a/src/etcd/tests/unit/test_client.py b/src/etcd/tests/unit/test_client.py
@@ -1,5 +1,12 @@
 import unittest
 import etcd
+import dns.name
+import dns.rdtypes.IN.SRV
+import dns.resolver
+try:
+    import mock
+except ImportError:
+    from unittest import mock
 
 
 class TestClient(unittest.TestCase):
@@ -97,3 +104,25 @@ def test_allow_reconnect(self):
             allow_reconnect=True,
             use_proxies=True,
         )
+
+    def test_discover(self):
+        """Tests discovery."""
+        answers = []
+        for i in range(1,3):
+            r = mock.create_autospec(dns.rdtypes.IN.SRV.SRV)
+            r.port = 2379
+            try:
+                method = dns.name.from_unicode
+            except AttributeError:
+                method = dns.name.from_text
+            r.target = method(u'etcd{}.example.com'.format(i))
+            answers.append(r)
+        dns.resolver.query = mock.create_autospec(dns.resolver.query, return_value=answers)
+        self.machines = etcd.Client.machines
+        etcd.Client.machines = mock.create_autospec(etcd.Client.machines, return_value=[u'https://etcd2.example.com:2379'])
+        c = etcd.Client(srv_domain="example.com", allow_reconnect=True, protocol="https")
+        etcd.Client.machines = self.machines
+        self.assertEquals(c.host, u'etcd1.example.com')
+        self.assertEquals(c.port, 2379)
+        self.assertEquals(c._machines_cache,
+                          [u'https://etcd2.example.com:2379'])
diff --git a/src/etcd/tests/unit/test_lock.py b/src/etcd/tests/unit/test_lock.py
@@ -1,5 +1,8 @@
 import etcd
-import mock
+try:
+    import mock
+except ImportError:
+    from unittest import mock
 from etcd.tests.unit import TestClientApiBase
 
 
@@ -92,8 +95,7 @@ def test_acquired(self):
         self.assertTrue(self.locker._acquired())
         self.assertTrue(self.locker.is_taken)
         retval = ('/_locks/test_lock/1', '/_locks/test_lock/4')
-        self.locker._get_locker = mock.create_autospec(
-            self.locker._get_locker, return_value=retval)
+        self.locker._get_locker = mock.MagicMock(return_value=retval)
         self.assertFalse(self.locker._acquired(blocking=False))
         self.assertFalse(self.locker.is_taken)
         d = {
diff --git a/src/etcd/tests/unit/test_request.py b/src/etcd/tests/unit/test_request.py
@@ -397,12 +397,9 @@ def setUp(self):
     def _mock_api(self, status, d, cluster_id=None):
         resp = self._prepare_response(status, d)
         resp.getheader.return_value = cluster_id or "abcdef1234"
-        self.client.http.request_encode_body = mock.create_autospec(
-            self.client.http.request_encode_body, return_value=resp
-        )
-        self.client.http.request = mock.create_autospec(
-            self.client.http.request, return_value=resp
-        )
+        self.client.http.request_encode_body = mock.MagicMock(
+            return_value=resp)
+        self.client.http.request = mock.MagicMock(return_value=resp)
 
     def _mock_error(self, error_code, msg, cause, method='PUT', fields=None,
                     cluster_id=None):
