Memcache auto discovery

Author: Vedanta Jha

ChangeLog.rst

@@ -1,5 +1,9 @@
 Changelog
 =========
+New in version 4.0.1
+--------------------
+* Added auto discovery feature to HashClient
+
 New in version 4.0.0
 --------------------
 * Dropped Python 2 and 3.6 support

docs/getting_started.rst

@@ -118,6 +118,24 @@ follows:
    ``node3`` is added back into the hasher and will be retried for any future
    operations.
 
+
+Using a configuration node endpoint and auto discovery
+------------------------------------------------------
+This will use AWS elasticache auto discovery method to discover nodes by just 
+using Configuration node's endpoint
+
+.. code-block:: python
+
+    from pymemcache.client.hash import HashClient
+
+    client = HashClient('127.0.0.1:11211', enable_autodiscovery=True)
+    client.set('some_key', 'some value')
+    result = client.get('some_key')
+
+The client internally fetches all the nodes from the configuration nodes and sets up a connection with them,
+Refer AWS `doc`<https://docs.aws.amazon.com/AmazonElastiCache/latest/mem-ug/AutoDiscovery.html> for more information.
+
+
 Using the built-in retrying mechanism
 -------------------------------------
 The library comes with retry mechanisms that can be used to wrap all kinds of

pymemcache/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "4.0.0"
+__version__ = "4.0.1"
 
 from pymemcache.client.base import Client  # noqa
 from pymemcache.client.base import PooledClient  # noqa

pymemcache/client/base.py

@@ -1030,6 +1030,28 @@ def flush_all(self, delay: int = 0, noreply: Optional[bool] = None) -> bool:
             return True
         return results[0] == b"OK"
 
+    def auto_discover(self):
+        """
+        This is specific to AWS Elasticache
+
+        Returns list of hostname and ip address of the nodes
+        
+        The response received is as follows:
+        0: CONFIG cluster 0 134
+        1: configversion\r\n
+        2: hostname|ip-address|port hostname|ip-address|port ...\r\n
+        3:
+        4: END
+        5: blank
+        """
+        cmd = b"config get cluster"
+        data = self._misc_cmd([cmd], b"config get cluster", noreply=False)
+        lines = data.split(b'\n')
+        configs = [conf.split(b'|') for conf in lines[2].split(b' ')]
+        self.quit()
+        nodes = [(ip, int(port)) for host, ip, port in configs]
+        return nodes
+
     def quit(self) -> None:
         """
         The memcached "quit" command.
@@ -1356,7 +1378,19 @@ def __delitem__(self, key):
         self.delete(key, noreply=True)
 
 
-class PooledClient:
+class PooledClient:        # 0: CONFIG cluster 0 134        # 0: CONFIG cluster 0 134
+        # 1: configversion\r\n
+        # 2: hostname|ip-address|port hostname|ip-address|port ...\r\n
+        # 3:
+        # 4: END
+        # 5: blank
+
+        # 1: configversion\r\n
+        # 2: hostname|ip-address|port hostname|ip-address|port ...\r\n
+        # 3:
+        # 4: END
+        # 5: blank
+
     """A thread-safe pool of clients (with the same client api).
 
     Args:

pymemcache/client/hash.py

@@ -48,13 +48,16 @@ def __init__(
         default_noreply=True,
         encoding="ascii",
         tls_context=None,
+        enable_auto_discovery=False
     ):
         """
         Constructor.
 
         Args:
           servers: list() of tuple(hostname, port) or string containing a UNIX
                    socket path.
+                   Or If enable_auto_discovery is set, just tuple(hostname, port) or UNIX socket path string 
+                   of configuration node would suffice.
           hasher: optional class three functions ``get_node``, ``add_node``,
                   and ``remove_node``
                   defaults to Rendezvous (HRW) hash.
@@ -70,12 +73,14 @@ def __init__(
           dead_timeout (float): Time in seconds before attempting to add a node
                                 back in the pool.
           encoding: optional str, controls data encoding (defaults to 'ascii').
+          enable_auto_discovery (bool): If enabled, nodes would be discovered from the configuration endpoint.
 
         Further arguments are interpreted as for :py:class:`.Client`
         constructor.
         """
         self.clients = {}
         self.retry_attempts = retry_attempts
+        self.connect_timeout = connect_timeout
         self.retry_timeout = retry_timeout
         self.dead_timeout = dead_timeout
         self.use_pooling = use_pooling
@@ -112,11 +117,18 @@ def __init__(
                     "lock_generator": lock_generator,
                 }
             )
-
-        for server in servers:
-            self.add_server(normalize_server_spec(server))
         self.encoding = encoding
         self.tls_context = tls_context
+        if not isinstance(servers, list):
+            if not enable_auto_discovery:
+                raise ValueError(f"Auto Discovery should be enabled if configuration endpoint is used: {servers!r}")
+
+            if enable_auto_discovery and servers is not None:
+                # AutoDiscovery is enabled and a address of configuration node is provided
+                servers = self._auto_discover(normalize_server_spec(servers))
+        
+        for server in servers:
+            self.add_server(normalize_server_spec(server))
 
     def _make_client_key(self, server):
         if isinstance(server, (list, tuple)) and len(server) == 2:
@@ -340,6 +352,12 @@ def _set_many(self, client, values, *args, **kwargs):
         succeeded = [key for key in values if key not in failed]
         return succeeded, failed, None
 
+    def _auto_discover(self, server):
+            _class = PooledClient if self.use_pooling else self.client_class
+            client = _class(server)
+            nodes  = client.auto_discover()   
+            return nodes
+             
     def close(self):
         for client in self.clients.values():
             self._safely_run_func(client, client.close, False)

pymemcache/test/test_client.py

@@ -1219,6 +1219,20 @@ def test_send_end_token_types(self):
         assert client.raw_command("key", "\r\n") == b"REPLY"
         assert client.raw_command(b"key", b"\r\n") == b"REPLY"
 
+    def test_auto_discover(self):
+        mock_response = (
+            b"CONFIG cluster 0 134\r\n"
+            b"configversion\r\n"
+            b"hostname1|10.0.0.1|11211 hostname2|10.0.0.2|11211\r\n"
+            b"END\r\n"
+        )
+
+        client = self.make_client([mock_response])
+        nodes = client.auto_discover()
+
+        expected_nodes = [("10.0.0.1", 11211), ("10.0.0.2", 11211)]
+        assert nodes == expected_nodes
+        
 
 @pytest.mark.unit()
 class TestClientSocketConnect(unittest.TestCase):
@@ -1431,6 +1445,7 @@ class MyClient(Client):
         client = PooledClient(("host", 11211))
         client.client_class = MyClient
         assert isinstance(client.client_pool.get(), MyClient)
+        
 
 
 class TestPooledClientIdleTimeout(ClientTestMixin, unittest.TestCase):