diff --git a/pymemcache/client/base.py b/pymemcache/client/base.py index bac39dd..a8de859 100644 --- a/pymemcache/client/base.py +++ b/pymemcache/client/base.py @@ -1030,6 +1030,15 @@ def flush_all(self, delay: int = 0, noreply: Optional[bool] = None) -> bool: return True return results[0] == b"OK" + def auto_discover(self): + 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. diff --git a/pymemcache/client/hash.py b/pymemcache/client/hash.py index c0574ce..6cd5953 100644 --- a/pymemcache/client/hash.py +++ b/pymemcache/client/hash.py @@ -48,6 +48,7 @@ def __init__( default_noreply=True, encoding="ascii", tls_context=None, + enable_auto_discovery=False ): """ Constructor. @@ -55,6 +56,7 @@ def __init__( Args: servers: list() of tuple(hostname, port) or string containing a UNIX socket path. + If enable_auto_discovery is set, just a configuration endpoint would suffice hasher: optional class three functions ``get_node``, ``add_node``, and ``remove_node`` defaults to Rendezvous (HRW) hash. @@ -70,12 +72,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 +116,14 @@ 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 enable_auto_discovery: + servers = self._auto_discover(server) + + 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 +347,14 @@ 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, configuration_endpoint, tls_context): + host, port = configuration_endpoint.split(':') + server = (host, port) + _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)