1754 lines
60 KiB
Python
1754 lines
60 KiB
Python
# Copyright 2012 Pinterest.com
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import errno
|
|
import platform
|
|
import socket
|
|
from functools import partial
|
|
from ssl import SSLContext
|
|
from types import ModuleType
|
|
from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Union
|
|
|
|
from pymemcache import pool
|
|
from pymemcache.exceptions import (
|
|
MemcacheClientError,
|
|
MemcacheIllegalInputError,
|
|
MemcacheServerError,
|
|
MemcacheUnexpectedCloseError,
|
|
MemcacheUnknownCommandError,
|
|
MemcacheUnknownError,
|
|
)
|
|
from pymemcache.serde import LegacyWrappingSerde
|
|
|
|
RECV_SIZE = 4096
|
|
VALID_STORE_RESULTS = {
|
|
b"set": (b"STORED", b"NOT_STORED"),
|
|
b"add": (b"STORED", b"NOT_STORED"),
|
|
b"replace": (b"STORED", b"NOT_STORED"),
|
|
b"append": (b"STORED", b"NOT_STORED"),
|
|
b"prepend": (b"STORED", b"NOT_STORED"),
|
|
b"cas": (b"STORED", b"EXISTS", b"NOT_FOUND"),
|
|
}
|
|
|
|
SOCKET_KEEPALIVE_SUPPORTED_SYSTEM = {
|
|
"Linux",
|
|
}
|
|
|
|
STORE_RESULTS_VALUE = {
|
|
b"STORED": True,
|
|
b"NOT_STORED": False,
|
|
b"NOT_FOUND": None,
|
|
b"EXISTS": False,
|
|
}
|
|
|
|
ServerSpec = Union[Tuple[str, int], str]
|
|
Key = Union[bytes, str]
|
|
|
|
|
|
# Some of the values returned by the "stats" command
|
|
# need mapping into native Python types
|
|
def _parse_bool_int(value: bytes) -> bool:
|
|
return int(value) != 0
|
|
|
|
|
|
def _parse_bool_string_is_yes(value: bytes) -> bool:
|
|
return value == b"yes"
|
|
|
|
|
|
def _parse_float(value: bytes) -> float:
|
|
return float(value.replace(b":", b"."))
|
|
|
|
|
|
def _parse_hex(value: bytes) -> int:
|
|
return int(value, 8)
|
|
|
|
|
|
STAT_TYPES: Dict[bytes, Callable[[bytes], Any]] = {
|
|
# General stats
|
|
b"version": bytes,
|
|
b"rusage_user": _parse_float,
|
|
b"rusage_system": _parse_float,
|
|
b"hash_is_expanding": _parse_bool_int,
|
|
b"slab_reassign_running": _parse_bool_int,
|
|
# Settings stats
|
|
b"inter": bytes,
|
|
b"growth_factor": float,
|
|
b"stat_key_prefix": bytes,
|
|
b"umask": _parse_hex,
|
|
b"detail_enabled": _parse_bool_int,
|
|
b"cas_enabled": _parse_bool_int,
|
|
b"auth_enabled_sasl": _parse_bool_string_is_yes,
|
|
b"maxconns_fast": _parse_bool_int,
|
|
b"slab_reassign": _parse_bool_int,
|
|
b"slab_automove": _parse_bool_int,
|
|
}
|
|
|
|
# Common helper functions.
|
|
|
|
|
|
def check_key_helper(
|
|
key: Key, allow_unicode_keys: bool, key_prefix: bytes = b""
|
|
) -> bytes:
|
|
"""Checks key and add key_prefix."""
|
|
if allow_unicode_keys:
|
|
if isinstance(key, str):
|
|
key = key.encode("utf8")
|
|
elif isinstance(key, str):
|
|
try:
|
|
key = key.encode("ascii")
|
|
except (UnicodeEncodeError, UnicodeDecodeError):
|
|
raise MemcacheIllegalInputError("Non-ASCII key: %r" % key)
|
|
|
|
key = key_prefix + key
|
|
parts = key.split()
|
|
|
|
if len(key) > 250:
|
|
raise MemcacheIllegalInputError("Key is too long: %r" % key)
|
|
# second statement catches leading or trailing whitespace
|
|
elif len(parts) > 1 or (parts and parts[0] != key):
|
|
raise MemcacheIllegalInputError("Key contains whitespace: %r" % key)
|
|
elif b"\00" in key:
|
|
raise MemcacheIllegalInputError("Key contains null: %r" % key)
|
|
|
|
return key
|
|
|
|
|
|
def normalize_server_spec(server: ServerSpec) -> ServerSpec:
|
|
if isinstance(server, tuple):
|
|
return server
|
|
if not isinstance(server, str):
|
|
raise ValueError(f"Unsupported server specification: {server!r}")
|
|
if server.startswith("unix:"):
|
|
return server[5:]
|
|
if server.startswith("/"):
|
|
return server
|
|
if ":" not in server or server.endswith("]"):
|
|
host, port = server, 11211
|
|
else:
|
|
parts = server.rsplit(":", 1)
|
|
host, port = parts[0], int(parts[1])
|
|
if host.startswith("["):
|
|
host = host.strip("[]")
|
|
return (host, port)
|
|
|
|
|
|
class KeepaliveOpts:
|
|
"""
|
|
A configuration structure to define the socket keepalive.
|
|
|
|
This structure must be passed to a client. The client will configure
|
|
its socket keepalive by using the elements of the structure.
|
|
|
|
Args:
|
|
idle: The time (in seconds) the connection needs to remain idle
|
|
before TCP starts sending keepalive probes. Should be a positive
|
|
integer most greater than zero.
|
|
intvl: The time (in seconds) between individual keepalive probes.
|
|
Should be a positive integer most greater than zero.
|
|
cnt: The maximum number of keepalive probes TCP should send before
|
|
dropping the connection. Should be a positive integer most greater
|
|
than zero.
|
|
"""
|
|
|
|
__slots__ = ("idle", "intvl", "cnt")
|
|
|
|
def __init__(self, idle: int = 1, intvl: int = 1, cnt: int = 5) -> None:
|
|
if idle < 1:
|
|
raise ValueError("The idle parameter must be greater or equal to 1.")
|
|
self.idle = idle
|
|
if intvl < 1:
|
|
raise ValueError("The intvl parameter must be greater or equal to 1.")
|
|
self.intvl = intvl
|
|
if cnt < 1:
|
|
raise ValueError("The cnt parameter must be greater or equal to 1.")
|
|
self.cnt = cnt
|
|
|
|
|
|
class Client:
|
|
"""
|
|
A client for a single memcached server.
|
|
|
|
*Server Connection*
|
|
|
|
The ``server`` parameter controls how the client connects to the memcached
|
|
server. You can either use a (host, port) tuple for a TCP connection or a
|
|
string containing the path to a UNIX domain socket.
|
|
|
|
The ``connect_timeout`` and ``timeout`` parameters can be used to set
|
|
socket timeout values. By default, timeouts are disabled.
|
|
|
|
When the ``no_delay`` flag is set, the ``TCP_NODELAY`` socket option will
|
|
also be set. This only applies to TCP-based connections.
|
|
|
|
Lastly, the ``socket_module`` allows you to specify an alternate socket
|
|
implementation (such as `gevent.socket`_).
|
|
|
|
.. _gevent.socket: http://www.gevent.org/api/gevent.socket.html
|
|
|
|
*Keys and Values*
|
|
|
|
Keys must have a __str__() method which should return a str with no more
|
|
than 250 ASCII characters and no whitespace or control characters. Unicode
|
|
strings must be encoded (as UTF-8, for example) unless they consist only
|
|
of ASCII characters that are neither whitespace nor control characters.
|
|
|
|
Values must have a __str__() method to convert themselves to a byte
|
|
string. Unicode objects can be a problem since str() on a Unicode object
|
|
will attempt to encode it as ASCII (which will fail if the value contains
|
|
code points larger than U+127). You can fix this with a serializer or by
|
|
just calling encode on the string (using UTF-8, for instance).
|
|
|
|
If you intend to use anything but str as a value, it is a good idea to use
|
|
a serializer. The pymemcache.serde library has an already implemented
|
|
serializer which pickles and unpickles data.
|
|
|
|
*Serialization and Deserialization*
|
|
|
|
The constructor takes an optional object, the "serializer/deserializer"
|
|
("serde"), which is responsible for both serialization and deserialization
|
|
of objects. That object must satisfy the serializer interface by providing
|
|
two methods: `serialize` and `deserialize`. `serialize` takes two
|
|
arguments, a key and a value, and returns a tuple of two elements, the
|
|
serialized value, and an integer in the range 0-65535 (the "flags").
|
|
`deserialize` takes three parameters, a key, value, and flags, and returns
|
|
the deserialized value.
|
|
|
|
Here is an example using JSON for non-str values:
|
|
|
|
.. code-block:: python
|
|
|
|
class JSONSerde(object):
|
|
def serialize(self, key, value):
|
|
if isinstance(value, str):
|
|
return value, 1
|
|
return json.dumps(value), 2
|
|
|
|
def deserialize(self, key, value, flags):
|
|
if flags == 1:
|
|
return value
|
|
|
|
if flags == 2:
|
|
return json.loads(value)
|
|
|
|
raise Exception("Unknown flags for value: {1}".format(flags))
|
|
|
|
.. note::
|
|
|
|
Most write operations allow the caller to provide a ``flags`` value to
|
|
support advanced interaction with the server. This will **override** the
|
|
"flags" value returned by the serializer and should therefore only be
|
|
used when you have a complete understanding of how the value should be
|
|
serialized, stored, and deserialized.
|
|
|
|
*Error Handling*
|
|
|
|
All of the methods in this class that talk to memcached can throw one of
|
|
the following exceptions:
|
|
|
|
* :class:`pymemcache.exceptions.MemcacheUnknownCommandError`
|
|
* :class:`pymemcache.exceptions.MemcacheClientError`
|
|
* :class:`pymemcache.exceptions.MemcacheServerError`
|
|
* :class:`pymemcache.exceptions.MemcacheUnknownError`
|
|
* :class:`pymemcache.exceptions.MemcacheUnexpectedCloseError`
|
|
* :class:`pymemcache.exceptions.MemcacheIllegalInputError`
|
|
* :class:`socket.timeout`
|
|
* :class:`socket.error`
|
|
|
|
Instances of this class maintain a persistent connection to memcached
|
|
which is terminated when any of these exceptions are raised. The next
|
|
call to a method on the object will result in a new connection being made
|
|
to memcached.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
server: ServerSpec,
|
|
serde=None,
|
|
serializer=None,
|
|
deserializer=None,
|
|
connect_timeout: Optional[float] = None,
|
|
timeout: Optional[float] = None,
|
|
no_delay: bool = False,
|
|
ignore_exc: bool = False,
|
|
socket_module: ModuleType = socket,
|
|
socket_keepalive: Optional[KeepaliveOpts] = None,
|
|
key_prefix: bytes = b"",
|
|
default_noreply: bool = True,
|
|
allow_unicode_keys: bool = False,
|
|
encoding: str = "ascii",
|
|
tls_context: Optional[SSLContext] = None,
|
|
):
|
|
"""
|
|
Constructor.
|
|
|
|
Args:
|
|
server: tuple(hostname, port) or string containing a UNIX socket path.
|
|
serde: optional serializer object, see notes in the class docs.
|
|
serializer: deprecated serialization function
|
|
deserializer: deprecated deserialization function
|
|
connect_timeout: optional float, seconds to wait for a connection to
|
|
the memcached server. Defaults to "forever" (uses the underlying
|
|
default socket timeout, which can be very long).
|
|
timeout: optional float, seconds to wait for send or recv calls on
|
|
the socket connected to memcached. Defaults to "forever" (uses the
|
|
underlying default socket timeout, which can be very long).
|
|
no_delay: optional bool, set the TCP_NODELAY flag, which may help
|
|
with performance in some cases. Defaults to False.
|
|
ignore_exc: optional bool, True to cause the "get", "gets",
|
|
"get_many" and "gets_many" calls to treat any errors as cache
|
|
misses. Defaults to False.
|
|
socket_module: socket module to use, e.g. gevent.socket. Defaults to
|
|
the standard library's socket module.
|
|
socket_keepalive: Activate the socket keepalive feature by passing
|
|
a KeepaliveOpts structure in this parameter. Disabled by default
|
|
(None). This feature is only supported on Linux platforms.
|
|
key_prefix: Prefix of key. You can use this as namespace. Defaults
|
|
to b''.
|
|
default_noreply: bool, the default value for 'noreply' as passed to
|
|
store commands (except from cas, incr, and decr, which default to
|
|
False).
|
|
allow_unicode_keys: bool, support unicode (utf8) keys
|
|
encoding: optional str, controls data encoding (defaults to 'ascii').
|
|
|
|
Notes:
|
|
The constructor does not make a connection to memcached. The first
|
|
call to a method on the object will do that.
|
|
"""
|
|
self.server = normalize_server_spec(server)
|
|
self.serde = serde or LegacyWrappingSerde(serializer, deserializer)
|
|
self.connect_timeout = connect_timeout
|
|
self.timeout = timeout
|
|
self.no_delay = no_delay
|
|
self.ignore_exc = ignore_exc
|
|
self.socket_module = socket_module
|
|
self.socket_keepalive = socket_keepalive
|
|
user_system = platform.system()
|
|
if self.socket_keepalive is not None:
|
|
if user_system not in SOCKET_KEEPALIVE_SUPPORTED_SYSTEM:
|
|
raise SystemError(
|
|
"Pymemcache's socket keepalive mechanism doesn't "
|
|
"support your system ({user_system}). If "
|
|
"you see this message it mean that you tried to "
|
|
"configure your socket keepalive on an unsupported "
|
|
"system. To fix the problem pass `socket_"
|
|
"keepalive=False` or use a supported system. "
|
|
"Supported systems are: {systems}".format(
|
|
user_system=user_system,
|
|
systems=", ".join(sorted(SOCKET_KEEPALIVE_SUPPORTED_SYSTEM)),
|
|
)
|
|
)
|
|
if not isinstance(self.socket_keepalive, KeepaliveOpts):
|
|
raise ValueError(
|
|
"Unsupported keepalive options. If you see this message "
|
|
"it means that you passed an unsupported object within "
|
|
"the param `socket_keepalive`. To fix it "
|
|
"please instantiate and pass to socket_keepalive a "
|
|
"KeepaliveOpts object. That's the only supported type "
|
|
"of structure."
|
|
)
|
|
self.sock: Optional[socket.socket] = None
|
|
if isinstance(key_prefix, str):
|
|
key_prefix = key_prefix.encode("ascii")
|
|
if not isinstance(key_prefix, bytes):
|
|
raise TypeError("key_prefix should be bytes.")
|
|
self.key_prefix = key_prefix
|
|
self.default_noreply = default_noreply
|
|
self.allow_unicode_keys = allow_unicode_keys
|
|
self.encoding = encoding
|
|
self.tls_context = tls_context
|
|
|
|
def check_key(self, key: Key, key_prefix: bytes) -> bytes:
|
|
"""Checks key and add key_prefix."""
|
|
return check_key_helper(
|
|
key, allow_unicode_keys=self.allow_unicode_keys, key_prefix=key_prefix
|
|
)
|
|
|
|
def _connect(self) -> None:
|
|
self.close()
|
|
|
|
s = self.socket_module
|
|
|
|
if not isinstance(self.server, tuple):
|
|
sockaddr = self.server
|
|
sock = s.socket(s.AF_UNIX, s.SOCK_STREAM)
|
|
else:
|
|
sock = None
|
|
error = None
|
|
host, port = self.server
|
|
info = s.getaddrinfo(host, port, s.AF_UNSPEC, s.SOCK_STREAM, s.IPPROTO_TCP)
|
|
for family, socktype, proto, _, sockaddr in info:
|
|
try:
|
|
sock = s.socket(family, socktype, proto)
|
|
if self.no_delay:
|
|
sock.setsockopt(s.IPPROTO_TCP, s.TCP_NODELAY, 1)
|
|
if self.tls_context:
|
|
context = self.tls_context
|
|
sock = context.wrap_socket(sock, server_hostname=host)
|
|
except Exception as e:
|
|
error = e
|
|
if sock is not None:
|
|
sock.close()
|
|
sock = None
|
|
else:
|
|
break
|
|
|
|
if error is not None:
|
|
raise error
|
|
|
|
try:
|
|
sock.settimeout(self.connect_timeout)
|
|
if self.socket_keepalive is not None:
|
|
sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
|
|
sock.setsockopt(
|
|
socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, self.socket_keepalive.idle
|
|
)
|
|
sock.setsockopt(
|
|
socket.IPPROTO_TCP,
|
|
socket.TCP_KEEPINTVL,
|
|
self.socket_keepalive.intvl,
|
|
)
|
|
sock.setsockopt(
|
|
socket.IPPROTO_TCP, socket.TCP_KEEPCNT, self.socket_keepalive.cnt
|
|
)
|
|
sock.connect(sockaddr)
|
|
sock.settimeout(self.timeout)
|
|
except Exception:
|
|
sock.close()
|
|
raise
|
|
|
|
self.sock = sock
|
|
|
|
def close(self) -> None:
|
|
"""Close the connection to memcached, if it is open. The next call to a
|
|
method that requires a connection will re-open it."""
|
|
if self.sock is not None:
|
|
try:
|
|
self.sock.close()
|
|
except Exception:
|
|
pass
|
|
finally:
|
|
self.sock = None
|
|
|
|
disconnect_all = close
|
|
|
|
def set(
|
|
self,
|
|
key: Key,
|
|
value: Any,
|
|
expire: int = 0,
|
|
noreply: Optional[bool] = None,
|
|
flags: Optional[int] = None,
|
|
) -> Optional[bool]:
|
|
"""
|
|
The memcached "set" command.
|
|
|
|
Args:
|
|
key: str, see class docs for details.
|
|
value: str, see class docs for details.
|
|
expire: optional int, number of seconds until the item is expired
|
|
from the cache, or zero for no expiry (the default).
|
|
noreply: optional bool, True to not wait for the reply (defaults to
|
|
self.default_noreply).
|
|
flags: optional int, arbitrary bit field used for server-specific
|
|
flags
|
|
|
|
Returns:
|
|
If no exception is raised, always returns True. If an exception is
|
|
raised, the set may or may not have occurred. If noreply is True,
|
|
then a successful return does not guarantee a successful set.
|
|
"""
|
|
if noreply is None:
|
|
noreply = self.default_noreply
|
|
# Optional because _store_cmd lookup in STORE_RESULTS_VALUE can return None in some cases.
|
|
# TODO: refactor to fix
|
|
return self._store_cmd(b"set", {key: value}, expire, noreply, flags=flags)[key]
|
|
|
|
def set_many(
|
|
self,
|
|
values: Dict[Key, Any],
|
|
expire: int = 0,
|
|
noreply: Optional[bool] = None,
|
|
flags: Optional[int] = None,
|
|
) -> List[Key]:
|
|
"""
|
|
A convenience function for setting multiple values.
|
|
|
|
Args:
|
|
values: dict(str, str), a dict of keys and values, see class docs
|
|
for details.
|
|
expire: optional int, number of seconds until the item is expired
|
|
from the cache, or zero for no expiry (the default).
|
|
noreply: optional bool, True to not wait for the reply (defaults to
|
|
self.default_noreply).
|
|
flags: optional int, arbitrary bit field used for server-specific
|
|
flags
|
|
|
|
Returns:
|
|
Returns a list of keys that failed to be inserted.
|
|
If noreply is True, always returns empty list.
|
|
"""
|
|
if noreply is None:
|
|
noreply = self.default_noreply
|
|
result = self._store_cmd(b"set", values, expire, noreply, flags=flags)
|
|
return [k for k, v in result.items() if not v]
|
|
|
|
set_multi = set_many
|
|
|
|
def add(
|
|
self,
|
|
key: Key,
|
|
value: Any,
|
|
expire: int = 0,
|
|
noreply: Optional[bool] = None,
|
|
flags: Optional[int] = None,
|
|
) -> bool:
|
|
"""
|
|
The memcached "add" command.
|
|
|
|
Args:
|
|
key: str, see class docs for details.
|
|
value: str, see class docs for details.
|
|
expire: optional int, number of seconds until the item is expired
|
|
from the cache, or zero for no expiry (the default).
|
|
noreply: optional bool, True to not wait for the reply (defaults to
|
|
self.default_noreply).
|
|
flags: optional int, arbitrary bit field used for server-specific
|
|
flags
|
|
|
|
Returns:
|
|
If ``noreply`` is True (or if it is unset and ``self.default_noreply``
|
|
is True), the return value is always True. Otherwise the return value
|
|
is True if the value was stored, and False if it was not (because
|
|
the key already existed).
|
|
"""
|
|
if noreply is None:
|
|
noreply = self.default_noreply
|
|
response = self._store_cmd(b"add", {key: value}, expire, noreply, flags=flags)[
|
|
key
|
|
]
|
|
# For typing, can only be None, if cas command
|
|
assert response is not None
|
|
return response
|
|
|
|
def replace(
|
|
self,
|
|
key: Key,
|
|
value,
|
|
expire: int = 0,
|
|
noreply: Optional[bool] = None,
|
|
flags: Optional[int] = None,
|
|
) -> bool:
|
|
"""
|
|
The memcached "replace" command.
|
|
|
|
Args:
|
|
key: str, see class docs for details.
|
|
value: str, see class docs for details.
|
|
expire: optional int, number of seconds until the item is expired
|
|
from the cache, or zero for no expiry (the default).
|
|
noreply: optional bool, True to not wait for the reply (defaults to
|
|
self.default_noreply).
|
|
flags: optional int, arbitrary bit field used for server-specific
|
|
flags
|
|
|
|
Returns:
|
|
If ``noreply`` is True (or if it is unset and ``self.default_noreply``
|
|
is True), the return value is always True. Otherwise returns True if
|
|
the value was stored and False if it wasn't (because the key didn't
|
|
already exist).
|
|
"""
|
|
if noreply is None:
|
|
noreply = self.default_noreply
|
|
response = self._store_cmd(
|
|
b"replace", {key: value}, expire, noreply, flags=flags
|
|
)[key]
|
|
# for typing
|
|
assert response is not None
|
|
return response
|
|
|
|
def append(
|
|
self,
|
|
key: Key,
|
|
value,
|
|
expire: int = 0,
|
|
noreply: Optional[bool] = None,
|
|
flags: Optional[int] = None,
|
|
) -> bool:
|
|
"""
|
|
The memcached "append" command.
|
|
|
|
Args:
|
|
key: str, see class docs for details.
|
|
value: str, see class docs for details.
|
|
expire: optional int, number of seconds until the item is expired
|
|
from the cache, or zero for no expiry (the default).
|
|
noreply: optional bool, True to not wait for the reply (defaults to
|
|
self.default_noreply).
|
|
flags: optional int, arbitrary bit field used for server-specific
|
|
flags
|
|
|
|
Returns:
|
|
True.
|
|
"""
|
|
if noreply is None:
|
|
noreply = self.default_noreply
|
|
response = self._store_cmd(
|
|
b"append", {key: value}, expire, noreply, flags=flags
|
|
)[key]
|
|
# For typing
|
|
assert response is not None
|
|
return response
|
|
|
|
def prepend(
|
|
self,
|
|
key,
|
|
value,
|
|
expire: int = 0,
|
|
noreply: Optional[bool] = None,
|
|
flags: Optional[int] = None,
|
|
):
|
|
"""
|
|
The memcached "prepend" command.
|
|
|
|
Args:
|
|
key: str, see class docs for details.
|
|
value: str, see class docs for details.
|
|
expire: optional int, number of seconds until the item is expired
|
|
from the cache, or zero for no expiry (the default).
|
|
noreply: optional bool, True to not wait for the reply (defaults to
|
|
self.default_noreply).
|
|
flags: optional int, arbitrary bit field used for server-specific
|
|
flags
|
|
|
|
Returns:
|
|
True.
|
|
"""
|
|
if noreply is None:
|
|
noreply = self.default_noreply
|
|
return self._store_cmd(b"prepend", {key: value}, expire, noreply, flags=flags)[
|
|
key
|
|
]
|
|
|
|
def cas(
|
|
self,
|
|
key,
|
|
value,
|
|
cas,
|
|
expire: int = 0,
|
|
noreply=False,
|
|
flags: Optional[int] = None,
|
|
) -> Optional[bool]:
|
|
"""
|
|
The memcached "cas" command.
|
|
|
|
Args:
|
|
key: str, see class docs for details.
|
|
value: str, see class docs for details.
|
|
cas: int or str that only contains the characters '0'-'9'.
|
|
expire: optional int, number of seconds until the item is expired
|
|
from the cache, or zero for no expiry (the default).
|
|
noreply: optional bool, False to wait for the reply (the default).
|
|
flags: optional int, arbitrary bit field used for server-specific
|
|
flags
|
|
|
|
Returns:
|
|
If ``noreply`` is True (or if it is unset and ``self.default_noreply``
|
|
is True), the return value is always True. Otherwise returns None if
|
|
the key didn't exist, False if it existed but had a different cas
|
|
value and True if it existed and was changed.
|
|
"""
|
|
cas = self._check_cas(cas)
|
|
return self._store_cmd(
|
|
b"cas", {key: value}, expire, noreply, flags=flags, cas=cas
|
|
)[key]
|
|
|
|
def get(self, key: Key, default: Optional[Any] = None) -> Any:
|
|
"""
|
|
The memcached "get" command, but only for one key, as a convenience.
|
|
|
|
Args:
|
|
key: str, see class docs for details.
|
|
default: value that will be returned if the key was not found.
|
|
|
|
Returns:
|
|
The value for the key, or default if the key wasn't found.
|
|
"""
|
|
return self._fetch_cmd(b"get", [key], False, key_prefix=self.key_prefix).get(
|
|
key, default
|
|
)
|
|
|
|
def get_many(self, keys: Iterable[Key]) -> Dict[Key, Any]:
|
|
"""
|
|
The memcached "get" command.
|
|
|
|
Args:
|
|
keys: list(str), see class docs for details.
|
|
|
|
Returns:
|
|
A dict in which the keys are elements of the "keys" argument list
|
|
and the values are values from the cache. The dict may contain all,
|
|
some or none of the given keys.
|
|
"""
|
|
if not keys:
|
|
return {}
|
|
|
|
return self._fetch_cmd(b"get", keys, False, key_prefix=self.key_prefix)
|
|
|
|
get_multi = get_many
|
|
|
|
def gets(
|
|
self, key: Key, default: Any = None, cas_default: Any = None
|
|
) -> Tuple[Any, Any]:
|
|
"""
|
|
The memcached "gets" command for one key, as a convenience.
|
|
|
|
Args:
|
|
key: str, see class docs for details.
|
|
default: value that will be returned if the key was not found.
|
|
cas_default: same behaviour as default argument.
|
|
|
|
Returns:
|
|
A tuple of (value, cas)
|
|
or (default, cas_defaults) if the key was not found.
|
|
"""
|
|
defaults = (default, cas_default)
|
|
return self._fetch_cmd(b"gets", [key], True, key_prefix=self.key_prefix).get(
|
|
key, defaults
|
|
)
|
|
|
|
def gets_many(self, keys: Iterable[Key]) -> Dict[Key, Tuple[Any, Any]]:
|
|
"""
|
|
The memcached "gets" command.
|
|
|
|
Args:
|
|
keys: list(str), see class docs for details.
|
|
|
|
Returns:
|
|
A dict in which the keys are elements of the "keys" argument list and
|
|
the values are tuples of (value, cas) from the cache. The dict may
|
|
contain all, some or none of the given keys.
|
|
"""
|
|
if not keys:
|
|
return {}
|
|
|
|
return self._fetch_cmd(b"gets", keys, True, key_prefix=self.key_prefix)
|
|
|
|
def delete(self, key: Key, noreply: Optional[bool] = None) -> bool:
|
|
"""
|
|
The memcached "delete" command.
|
|
|
|
Args:
|
|
key: str, see class docs for details.
|
|
noreply: optional bool, True to not wait for the reply (defaults to
|
|
self.default_noreply).
|
|
|
|
Returns:
|
|
If ``noreply`` is True (or if it is unset and ``self.default_noreply``
|
|
is True), the return value is always True. Otherwise returns True if
|
|
the key was deleted, and False if it wasn't found.
|
|
"""
|
|
if noreply is None:
|
|
noreply = self.default_noreply
|
|
cmd = b"delete " + self.check_key(key, self.key_prefix)
|
|
if noreply:
|
|
cmd += b" noreply"
|
|
cmd += b"\r\n"
|
|
results = self._misc_cmd([cmd], b"delete", noreply)
|
|
if noreply:
|
|
return True
|
|
return results[0] == b"DELETED"
|
|
|
|
def delete_many(self, keys: Iterable[Key], noreply: Optional[bool] = None) -> bool:
|
|
"""
|
|
A convenience function to delete multiple keys.
|
|
|
|
Args:
|
|
keys: list(str), the list of keys to delete.
|
|
noreply: optional bool, True to not wait for the reply (defaults to
|
|
self.default_noreply).
|
|
|
|
Returns:
|
|
True. If an exception is raised then all, some or none of the keys
|
|
may have been deleted. Otherwise all the keys have been sent to
|
|
memcache for deletion and if noreply is False, they have been
|
|
acknowledged by memcache.
|
|
"""
|
|
if not keys:
|
|
return True
|
|
|
|
if noreply is None:
|
|
noreply = self.default_noreply
|
|
|
|
cmds = []
|
|
for key in keys:
|
|
cmds.append(
|
|
b"delete "
|
|
+ self.check_key(key, self.key_prefix)
|
|
+ (b" noreply" if noreply else b"")
|
|
+ b"\r\n"
|
|
)
|
|
self._misc_cmd(cmds, b"delete", noreply)
|
|
return True
|
|
|
|
delete_multi = delete_many
|
|
|
|
def incr(
|
|
self, key: Key, value: int, noreply: Optional[bool] = False
|
|
) -> Optional[int]:
|
|
"""
|
|
The memcached "incr" command.
|
|
|
|
Args:
|
|
key: str, see class docs for details.
|
|
value: int, the amount by which to increment the value.
|
|
noreply: optional bool, False to wait for the reply (the default).
|
|
|
|
Returns:
|
|
If noreply is True, always returns None. Otherwise returns the new
|
|
value of the key, or None if the key wasn't found.
|
|
"""
|
|
key = self.check_key(key, self.key_prefix)
|
|
val = self._check_integer(value, "value")
|
|
cmd = b"incr " + key + b" " + val
|
|
if noreply:
|
|
cmd += b" noreply"
|
|
cmd += b"\r\n"
|
|
results = self._misc_cmd([cmd], b"incr", noreply)
|
|
if noreply:
|
|
return None
|
|
if results[0] == b"NOT_FOUND":
|
|
return None
|
|
return int(results[0])
|
|
|
|
def decr(
|
|
self, key: Key, value: int, noreply: Optional[bool] = False
|
|
) -> Optional[int]:
|
|
"""
|
|
The memcached "decr" command.
|
|
|
|
Args:
|
|
key: str, see class docs for details.
|
|
value: int, the amount by which to decrement the value.
|
|
noreply: optional bool, False to wait for the reply (the default).
|
|
|
|
Returns:
|
|
If noreply is True, always returns None. Otherwise returns the new
|
|
value of the key, or None if the key wasn't found.
|
|
"""
|
|
key = self.check_key(key, self.key_prefix)
|
|
val = self._check_integer(value, "value")
|
|
cmd = b"decr " + key + b" " + val
|
|
if noreply:
|
|
cmd += b" noreply"
|
|
cmd += b"\r\n"
|
|
results = self._misc_cmd([cmd], b"decr", noreply)
|
|
if noreply:
|
|
return None
|
|
if results[0] == b"NOT_FOUND":
|
|
return None
|
|
return int(results[0])
|
|
|
|
def touch(self, key: Key, expire: int = 0, noreply: Optional[bool] = None) -> bool:
|
|
"""
|
|
The memcached "touch" command.
|
|
|
|
Args:
|
|
key: str, see class docs for details.
|
|
expire: optional int, number of seconds until the item is expired
|
|
from the cache, or zero for no expiry (the default).
|
|
noreply: optional bool, True to not wait for the reply (defaults to
|
|
self.default_noreply).
|
|
|
|
Returns:
|
|
True if the expiration time was updated, False if the key wasn't
|
|
found.
|
|
"""
|
|
if noreply is None:
|
|
noreply = self.default_noreply
|
|
key = self.check_key(key, self.key_prefix)
|
|
expire_bytes = self._check_integer(expire, "expire")
|
|
cmd = b"touch " + key + b" " + expire_bytes
|
|
if noreply:
|
|
cmd += b" noreply"
|
|
cmd += b"\r\n"
|
|
results = self._misc_cmd([cmd], b"touch", noreply)
|
|
if noreply:
|
|
return True
|
|
return results[0] == b"TOUCHED"
|
|
|
|
def stats(self, *args):
|
|
"""
|
|
The memcached "stats" command.
|
|
|
|
The returned keys depend on what the "stats" command returns.
|
|
A best effort is made to convert values to appropriate Python
|
|
types, defaulting to strings when a conversion cannot be made.
|
|
|
|
Args:
|
|
*arg: extra string arguments to the "stats" command. See the
|
|
memcached protocol documentation for more information.
|
|
|
|
Returns:
|
|
A dict of the returned stats.
|
|
"""
|
|
result = self._fetch_cmd(b"stats", args, False)
|
|
|
|
for key, value in result.items():
|
|
converter = STAT_TYPES.get(key, int)
|
|
try:
|
|
result[key] = converter(value)
|
|
except Exception:
|
|
pass
|
|
|
|
return result
|
|
|
|
def cache_memlimit(self, memlimit) -> bool:
|
|
"""
|
|
The memcached "cache_memlimit" command.
|
|
|
|
Args:
|
|
memlimit: int, the number of megabytes to set as the new cache memory
|
|
limit.
|
|
|
|
Returns:
|
|
If no exception is raised, always returns True.
|
|
"""
|
|
memlimit = self._check_integer(memlimit, "memlimit")
|
|
self._fetch_cmd(b"cache_memlimit", [memlimit], False)
|
|
return True
|
|
|
|
def version(self) -> bytes:
|
|
"""
|
|
The memcached "version" command.
|
|
|
|
Returns:
|
|
A string of the memcached version.
|
|
"""
|
|
cmd = b"version\r\n"
|
|
results = self._misc_cmd([cmd], b"version", False)
|
|
before, _, after = results[0].partition(b" ")
|
|
|
|
if before != b"VERSION":
|
|
raise MemcacheUnknownError(f"Received unexpected response: {results[0]!r}")
|
|
return after
|
|
|
|
def raw_command(
|
|
self, command: Union[str, bytes], end_tokens: Union[str, bytes] = "\r\n"
|
|
) -> bytes:
|
|
"""
|
|
Sends an arbitrary command to the server and parses the response until a
|
|
specified token is encountered.
|
|
|
|
Args:
|
|
command: str|bytes: The command to send.
|
|
end_tokens: str|bytes: The token expected at the end of the
|
|
response. If the `end_token` is not found, the client will wait
|
|
until the timeout specified in the constructor.
|
|
|
|
Returns:
|
|
The response from the server, with the `end_token` removed.
|
|
"""
|
|
encoding = "utf8" if self.allow_unicode_keys else "ascii"
|
|
command = command.encode(encoding) if isinstance(command, str) else command
|
|
end_tokens = (
|
|
end_tokens.encode(encoding) if isinstance(end_tokens, str) else end_tokens
|
|
)
|
|
return self._misc_cmd([b"" + command + b"\r\n"], command, False, end_tokens)[0]
|
|
|
|
def flush_all(self, delay: int = 0, noreply: Optional[bool] = None) -> bool:
|
|
"""
|
|
The memcached "flush_all" command.
|
|
|
|
Args:
|
|
delay: optional int, the number of seconds to wait before flushing,
|
|
or zero to flush immediately (the default).
|
|
noreply: optional bool, True to not wait for the reply (defaults to
|
|
self.default_noreply).
|
|
|
|
Returns:
|
|
True.
|
|
"""
|
|
if noreply is None:
|
|
noreply = self.default_noreply
|
|
delay_bytes = self._check_integer(delay, "delay")
|
|
cmd = b"flush_all " + delay_bytes
|
|
if noreply:
|
|
cmd += b" noreply"
|
|
cmd += b"\r\n"
|
|
results = self._misc_cmd([cmd], b"flush_all", noreply)
|
|
if noreply:
|
|
return True
|
|
return results[0] == b"OK"
|
|
|
|
def quit(self) -> None:
|
|
"""
|
|
The memcached "quit" command.
|
|
|
|
This will close the connection with memcached. Calling any other
|
|
method on this object will re-open the connection, so this object can
|
|
be re-used after quit.
|
|
"""
|
|
cmd = b"quit\r\n"
|
|
self._misc_cmd([cmd], b"quit", True)
|
|
self.close()
|
|
|
|
def shutdown(self, graceful: bool = False) -> None:
|
|
"""
|
|
The memcached "shutdown" command.
|
|
|
|
This will request shutdown and eventual termination of the server,
|
|
optionally preceded by a graceful stop of memcached's internal state
|
|
machine. Note that the server needs to have been started with the
|
|
shutdown protocol command enabled with the --enable-shutdown flag.
|
|
|
|
Args:
|
|
graceful: optional bool, True to request a graceful shutdown with
|
|
SIGUSR1 (defaults to False, i.e. SIGINT shutdown).
|
|
"""
|
|
cmd = b"shutdown"
|
|
if graceful:
|
|
cmd += b" graceful"
|
|
cmd += b"\r\n"
|
|
|
|
# The shutdown command raises a server-side error if the shutdown
|
|
# protocol command is not enabled. Otherwise, a successful shutdown
|
|
# is expected to close the remote end of the transport.
|
|
try:
|
|
self._misc_cmd([cmd], b"shutdown", False)
|
|
except MemcacheUnexpectedCloseError:
|
|
pass
|
|
|
|
def _raise_errors(self, line: bytes, name: bytes) -> None:
|
|
if line.startswith(b"ERROR"):
|
|
raise MemcacheUnknownCommandError(name)
|
|
|
|
if line.startswith(b"CLIENT_ERROR"):
|
|
error = line[line.find(b" ") + 1 :]
|
|
raise MemcacheClientError(error)
|
|
|
|
if line.startswith(b"SERVER_ERROR"):
|
|
error = line[line.find(b" ") + 1 :]
|
|
raise MemcacheServerError(error)
|
|
|
|
def _check_integer(self, value: int, name: str) -> bytes:
|
|
"""Check that a value is an integer and encode it as a binary string"""
|
|
if not isinstance(value, int):
|
|
raise MemcacheIllegalInputError(
|
|
f"{name} must be integer, got bad value: {value!r}"
|
|
)
|
|
|
|
return str(value).encode(self.encoding)
|
|
|
|
def _check_cas(self, cas: Union[int, str, bytes]) -> bytes:
|
|
"""Check that a value is a valid input for 'cas' -- either an int or a
|
|
string containing only 0-9
|
|
|
|
The value will be (re)encoded so that we can accept strings or bytes.
|
|
"""
|
|
# convert non-binary values to binary
|
|
if isinstance(cas, (int, str)):
|
|
try:
|
|
cas = str(cas).encode(self.encoding)
|
|
except UnicodeEncodeError:
|
|
raise MemcacheIllegalInputError("non-ASCII cas value: %r" % cas)
|
|
elif not isinstance(cas, bytes):
|
|
raise MemcacheIllegalInputError(
|
|
"cas must be integer, string, or bytes, got bad value: %r" % cas
|
|
)
|
|
|
|
if not cas.isdigit():
|
|
raise MemcacheIllegalInputError(
|
|
"cas must only contain values in 0-9, got bad value: %r" % cas
|
|
)
|
|
|
|
return cas
|
|
|
|
def _extract_value(
|
|
self,
|
|
expect_cas: bool,
|
|
line: bytes,
|
|
buf: bytes,
|
|
remapped_keys: Dict[bytes, Key],
|
|
prefixed_keys: List[bytes],
|
|
) -> Tuple[Key, Union[Any, Tuple[Any, bytes]], bytes]:
|
|
"""
|
|
This function is abstracted from _fetch_cmd to support different ways
|
|
of value extraction. In order to use this feature, _extract_value needs
|
|
to be overridden in the subclass.
|
|
"""
|
|
if expect_cas:
|
|
_, key, flags, size, cas = line.split()
|
|
else:
|
|
try:
|
|
_, key, flags, size = line.split()
|
|
except Exception as e:
|
|
raise ValueError(f"Unable to parse line {line!r}: {e}")
|
|
|
|
value = None
|
|
try:
|
|
# For typing
|
|
assert self.sock is not None
|
|
|
|
buf, value = _readvalue(self.sock, buf, int(size))
|
|
except MemcacheUnexpectedCloseError:
|
|
self.close()
|
|
raise
|
|
original_key = remapped_keys[key]
|
|
value = self.serde.deserialize(original_key, value, int(flags))
|
|
|
|
if expect_cas:
|
|
return original_key, (value, cas), buf
|
|
else:
|
|
return original_key, value, buf
|
|
|
|
def _fetch_cmd(
|
|
self,
|
|
name: bytes,
|
|
keys: Iterable[Key],
|
|
expect_cas: bool,
|
|
key_prefix: bytes = b"",
|
|
) -> Dict[Key, Any]:
|
|
prefixed_keys = [self.check_key(k, key_prefix=key_prefix) for k in keys]
|
|
remapped_keys = dict(zip(prefixed_keys, keys))
|
|
|
|
# It is important for all keys to be listed in their original order.
|
|
cmd = name
|
|
if prefixed_keys:
|
|
cmd += b" " + b" ".join(prefixed_keys)
|
|
cmd += b"\r\n"
|
|
|
|
try:
|
|
if self.sock is None:
|
|
self._connect()
|
|
|
|
# For typing
|
|
assert self.sock is not None
|
|
|
|
self.sock.sendall(cmd)
|
|
|
|
buf = b""
|
|
line = None
|
|
result: Dict[Key, Any] = {}
|
|
while True:
|
|
try:
|
|
buf, line = _readline(self.sock, buf)
|
|
except MemcacheUnexpectedCloseError:
|
|
self.close()
|
|
raise
|
|
self._raise_errors(line, name)
|
|
if line == b"END" or line == b"OK":
|
|
return result
|
|
elif line.startswith(b"VALUE"):
|
|
key, value, buf = self._extract_value(
|
|
expect_cas, line, buf, remapped_keys, prefixed_keys
|
|
)
|
|
result[key] = value
|
|
elif name == b"stats" and line.startswith(b"STAT"):
|
|
key_value = line.split()
|
|
result[key_value[1]] = key_value[2] if len(key_value) > 2 else b""
|
|
elif name == b"stats" and line.startswith(b"ITEM"):
|
|
# For 'stats cachedump' commands
|
|
key_value = line.split()
|
|
result[key_value[1]] = b" ".join(key_value[2:])
|
|
else:
|
|
raise MemcacheUnknownError(line[:32])
|
|
except Exception:
|
|
self.close()
|
|
if self.ignore_exc:
|
|
return {}
|
|
raise
|
|
|
|
def _store_cmd(
|
|
self,
|
|
name: bytes,
|
|
values: Dict[Key, Any],
|
|
expire: int,
|
|
noreply: bool,
|
|
flags: Optional[int] = None,
|
|
cas: Optional[bytes] = None,
|
|
) -> Dict[Key, Optional[bool]]:
|
|
cmds = []
|
|
keys = []
|
|
|
|
extra = b""
|
|
if cas is not None:
|
|
extra += b" " + cas
|
|
if noreply:
|
|
extra += b" noreply"
|
|
expire_bytes = self._check_integer(expire, "expire")
|
|
|
|
for key, data in values.items():
|
|
# must be able to reliably map responses back to the original order
|
|
keys.append(key)
|
|
|
|
key = self.check_key(key, self.key_prefix)
|
|
data, data_flags = self.serde.serialize(key, data)
|
|
|
|
# If 'flags' was explicitly provided, it overrides the value
|
|
# returned by the serializer.
|
|
if flags is not None:
|
|
data_flags = flags
|
|
|
|
if not isinstance(data, bytes):
|
|
try:
|
|
data = str(data).encode(self.encoding)
|
|
except UnicodeEncodeError as e:
|
|
raise MemcacheIllegalInputError(
|
|
"Data values must be binary-safe: %s" % e
|
|
)
|
|
|
|
cmds.append(
|
|
name
|
|
+ b" "
|
|
+ key
|
|
+ b" "
|
|
+ str(data_flags).encode(self.encoding)
|
|
+ b" "
|
|
+ expire_bytes
|
|
+ b" "
|
|
+ str(len(data)).encode(self.encoding)
|
|
+ extra
|
|
+ b"\r\n"
|
|
+ data
|
|
+ b"\r\n"
|
|
)
|
|
|
|
if self.sock is None:
|
|
self._connect()
|
|
|
|
# For typing
|
|
assert self.sock is not None
|
|
|
|
try:
|
|
self.sock.sendall(b"".join(cmds))
|
|
if noreply:
|
|
return {k: True for k in keys}
|
|
|
|
results = {}
|
|
buf = b""
|
|
line = None
|
|
for key in keys:
|
|
try:
|
|
buf, line = _readline(self.sock, buf)
|
|
except MemcacheUnexpectedCloseError:
|
|
self.close()
|
|
raise
|
|
self._raise_errors(line, name)
|
|
|
|
if line in VALID_STORE_RESULTS[name]:
|
|
results[key] = STORE_RESULTS_VALUE[line]
|
|
else:
|
|
raise MemcacheUnknownError(line[:32])
|
|
return results
|
|
except Exception:
|
|
self.close()
|
|
raise
|
|
|
|
def _misc_cmd(
|
|
self,
|
|
cmds: Iterable[bytes],
|
|
cmd_name: bytes,
|
|
noreply: Optional[bool],
|
|
end_tokens=None,
|
|
) -> List[bytes]:
|
|
|
|
# If no end_tokens have been given, just assume standard memcached
|
|
# operations, which end in "\r\n", use regular code for that.
|
|
_reader: Callable[[socket.socket, bytes], Tuple[bytes, bytes]]
|
|
if end_tokens:
|
|
_reader = partial(_readsegment, end_tokens=end_tokens)
|
|
else:
|
|
_reader = _readline
|
|
|
|
if self.sock is None:
|
|
self._connect()
|
|
|
|
# For typing
|
|
assert self.sock is not None
|
|
|
|
try:
|
|
self.sock.sendall(b"".join(cmds))
|
|
|
|
if noreply:
|
|
return []
|
|
|
|
results = []
|
|
buf = b""
|
|
line = None
|
|
for cmd in cmds:
|
|
try:
|
|
buf, line = _reader(self.sock, buf)
|
|
except MemcacheUnexpectedCloseError:
|
|
self.close()
|
|
raise
|
|
self._raise_errors(line, cmd_name)
|
|
results.append(line)
|
|
return results
|
|
|
|
except Exception:
|
|
self.close()
|
|
raise
|
|
|
|
def __setitem__(self, key: Key, value):
|
|
self.set(key, value, noreply=True)
|
|
|
|
def __getitem__(self, key):
|
|
value = self.get(key)
|
|
if value is None:
|
|
raise KeyError
|
|
return value
|
|
|
|
def __delitem__(self, key):
|
|
self.delete(key, noreply=True)
|
|
|
|
|
|
class PooledClient:
|
|
"""A thread-safe pool of clients (with the same client api).
|
|
|
|
Args:
|
|
max_pool_size: maximum pool size to use (going above this amount
|
|
triggers a runtime error), by default this is 2147483648L
|
|
when not provided (or none).
|
|
pool_idle_timeout: pooled connections are discarded if they have been
|
|
unused for this many seconds. A value of 0 indicates
|
|
that pooled connections are never discarded.
|
|
lock_generator: a callback/type that takes no arguments that will
|
|
be called to create a lock or semaphore that can
|
|
protect the pool from concurrent access (for example a
|
|
eventlet lock or semaphore could be used instead)
|
|
|
|
Further arguments are interpreted as for :py:class:`.Client` constructor.
|
|
|
|
Note: if `serde` is given, the same object will be used for *all* clients
|
|
in the pool. Your serde object must therefore be thread-safe.
|
|
"""
|
|
|
|
#: :class:`Client` class used to create new clients
|
|
client_class = Client
|
|
|
|
def __init__(
|
|
self,
|
|
server: ServerSpec,
|
|
serde=None,
|
|
serializer=None,
|
|
deserializer=None,
|
|
connect_timeout=None,
|
|
timeout=None,
|
|
no_delay=False,
|
|
ignore_exc=False,
|
|
socket_module=socket,
|
|
socket_keepalive=None,
|
|
key_prefix=b"",
|
|
max_pool_size=None,
|
|
pool_idle_timeout=0,
|
|
lock_generator=None,
|
|
default_noreply: bool = True,
|
|
allow_unicode_keys=False,
|
|
encoding="ascii",
|
|
tls_context=None,
|
|
):
|
|
self.server = normalize_server_spec(server)
|
|
self.serde = serde or LegacyWrappingSerde(serializer, deserializer)
|
|
self.connect_timeout = connect_timeout
|
|
self.timeout = timeout
|
|
self.no_delay = no_delay
|
|
self.ignore_exc = ignore_exc
|
|
self.socket_module = socket_module
|
|
self.socket_keepalive = socket_keepalive
|
|
self.default_noreply = default_noreply
|
|
self.allow_unicode_keys = allow_unicode_keys
|
|
if isinstance(key_prefix, str):
|
|
key_prefix = key_prefix.encode("ascii")
|
|
if not isinstance(key_prefix, bytes):
|
|
raise TypeError("key_prefix should be bytes.")
|
|
self.key_prefix = key_prefix
|
|
self.client_pool = pool.ObjectPool(
|
|
self._create_client,
|
|
after_remove=lambda client: client.close(),
|
|
max_size=max_pool_size,
|
|
idle_timeout=pool_idle_timeout,
|
|
lock_generator=lock_generator,
|
|
)
|
|
self.encoding = encoding
|
|
self.tls_context = tls_context
|
|
|
|
def check_key(self, key: Key) -> bytes:
|
|
"""Checks key and add key_prefix."""
|
|
return check_key_helper(
|
|
key, allow_unicode_keys=self.allow_unicode_keys, key_prefix=self.key_prefix
|
|
)
|
|
|
|
def _create_client(self) -> Client:
|
|
return self.client_class(
|
|
self.server,
|
|
serde=self.serde,
|
|
connect_timeout=self.connect_timeout,
|
|
timeout=self.timeout,
|
|
no_delay=self.no_delay,
|
|
# We need to know when it fails *always* so that we
|
|
# can remove/destroy it from the pool...
|
|
ignore_exc=False,
|
|
socket_module=self.socket_module,
|
|
socket_keepalive=self.socket_keepalive,
|
|
key_prefix=self.key_prefix,
|
|
default_noreply=self.default_noreply,
|
|
allow_unicode_keys=self.allow_unicode_keys,
|
|
tls_context=self.tls_context,
|
|
)
|
|
|
|
def close(self) -> None:
|
|
self.client_pool.clear()
|
|
|
|
disconnect_all = close
|
|
|
|
def set(
|
|
self,
|
|
key,
|
|
value,
|
|
expire: int = 0,
|
|
noreply: Optional[bool] = None,
|
|
flags: Optional[int] = None,
|
|
):
|
|
with self.client_pool.get_and_release(destroy_on_fail=True) as client:
|
|
return client.set(key, value, expire=expire, noreply=noreply, flags=flags)
|
|
|
|
def set_many(
|
|
self,
|
|
values,
|
|
expire: int = 0,
|
|
noreply: Optional[bool] = None,
|
|
flags: Optional[int] = None,
|
|
):
|
|
with self.client_pool.get_and_release(destroy_on_fail=True) as client:
|
|
return client.set_many(values, expire=expire, noreply=noreply, flags=flags)
|
|
|
|
set_multi = set_many
|
|
|
|
def replace(
|
|
self,
|
|
key,
|
|
value,
|
|
expire: int = 0,
|
|
noreply: Optional[bool] = None,
|
|
flags: Optional[int] = None,
|
|
):
|
|
with self.client_pool.get_and_release(destroy_on_fail=True) as client:
|
|
return client.replace(
|
|
key, value, expire=expire, noreply=noreply, flags=flags
|
|
)
|
|
|
|
def append(
|
|
self,
|
|
key,
|
|
value,
|
|
expire: int = 0,
|
|
noreply: Optional[bool] = None,
|
|
flags: Optional[int] = None,
|
|
):
|
|
with self.client_pool.get_and_release(destroy_on_fail=True) as client:
|
|
return client.append(
|
|
key, value, expire=expire, noreply=noreply, flags=flags
|
|
)
|
|
|
|
def prepend(
|
|
self,
|
|
key,
|
|
value,
|
|
expire: int = 0,
|
|
noreply: Optional[bool] = None,
|
|
flags: Optional[int] = None,
|
|
):
|
|
with self.client_pool.get_and_release(destroy_on_fail=True) as client:
|
|
return client.prepend(
|
|
key, value, expire=expire, noreply=noreply, flags=flags
|
|
)
|
|
|
|
def cas(
|
|
self,
|
|
key,
|
|
value,
|
|
cas,
|
|
expire: int = 0,
|
|
noreply=False,
|
|
flags: Optional[int] = None,
|
|
):
|
|
with self.client_pool.get_and_release(destroy_on_fail=True) as client:
|
|
return client.cas(
|
|
key, value, cas, expire=expire, noreply=noreply, flags=flags
|
|
)
|
|
|
|
def get(self, key: Key, default: Any = None) -> Any:
|
|
with self.client_pool.get_and_release(destroy_on_fail=True) as client:
|
|
try:
|
|
return client.get(key, default)
|
|
except Exception:
|
|
if self.ignore_exc:
|
|
return default
|
|
else:
|
|
raise
|
|
|
|
def get_many(self, keys: Iterable[Key]) -> Dict[Key, Any]:
|
|
with self.client_pool.get_and_release(destroy_on_fail=True) as client:
|
|
try:
|
|
return client.get_many(keys)
|
|
except Exception:
|
|
if self.ignore_exc:
|
|
return {}
|
|
else:
|
|
raise
|
|
|
|
get_multi = get_many
|
|
|
|
def gets(self, key: Key) -> Tuple[Any, Any]:
|
|
with self.client_pool.get_and_release(destroy_on_fail=True) as client:
|
|
try:
|
|
return client.gets(key)
|
|
except Exception:
|
|
if self.ignore_exc:
|
|
return (None, None)
|
|
else:
|
|
raise
|
|
|
|
def gets_many(self, keys: Iterable[Key]) -> Dict[Key, Tuple[Any, Any]]:
|
|
with self.client_pool.get_and_release(destroy_on_fail=True) as client:
|
|
try:
|
|
return client.gets_many(keys)
|
|
except Exception:
|
|
if self.ignore_exc:
|
|
return {}
|
|
else:
|
|
raise
|
|
|
|
def delete(self, key: Key, noreply: Optional[bool] = None) -> bool:
|
|
with self.client_pool.get_and_release(destroy_on_fail=True) as client:
|
|
return client.delete(key, noreply=noreply)
|
|
|
|
def delete_many(self, keys: Iterable[Key], noreply: Optional[bool] = None) -> bool:
|
|
with self.client_pool.get_and_release(destroy_on_fail=True) as client:
|
|
return client.delete_many(keys, noreply=noreply)
|
|
|
|
delete_multi = delete_many
|
|
|
|
def add(
|
|
self,
|
|
key: Key,
|
|
value,
|
|
expire: int = 0,
|
|
noreply: Optional[bool] = None,
|
|
flags: Optional[int] = None,
|
|
):
|
|
with self.client_pool.get_and_release(destroy_on_fail=True) as client:
|
|
return client.add(key, value, expire=expire, noreply=noreply, flags=flags)
|
|
|
|
def incr(self, key: Key, value, noreply=False):
|
|
with self.client_pool.get_and_release(destroy_on_fail=True) as client:
|
|
return client.incr(key, value, noreply=noreply)
|
|
|
|
def decr(self, key: Key, value, noreply=False):
|
|
with self.client_pool.get_and_release(destroy_on_fail=True) as client:
|
|
return client.decr(key, value, noreply=noreply)
|
|
|
|
def touch(self, key: Key, expire: int = 0, noreply=None):
|
|
with self.client_pool.get_and_release(destroy_on_fail=True) as client:
|
|
return client.touch(key, expire=expire, noreply=noreply)
|
|
|
|
def stats(self, *args):
|
|
with self.client_pool.get_and_release(destroy_on_fail=True) as client:
|
|
try:
|
|
return client.stats(*args)
|
|
except Exception:
|
|
if self.ignore_exc:
|
|
return {}
|
|
else:
|
|
raise
|
|
|
|
def version(self) -> bytes:
|
|
with self.client_pool.get_and_release(destroy_on_fail=True) as client:
|
|
return client.version()
|
|
|
|
def flush_all(self, delay=0, noreply=None) -> bool:
|
|
with self.client_pool.get_and_release(destroy_on_fail=True) as client:
|
|
return client.flush_all(delay=delay, noreply=noreply)
|
|
|
|
def quit(self) -> None:
|
|
with self.client_pool.get_and_release(destroy_on_fail=True) as client:
|
|
try:
|
|
client.quit()
|
|
finally:
|
|
self.client_pool.destroy(client)
|
|
|
|
def shutdown(self, graceful: bool = False) -> None:
|
|
with self.client_pool.get_and_release(destroy_on_fail=True) as client:
|
|
client.shutdown(graceful)
|
|
|
|
def raw_command(self, command, end_tokens=b"\r\n"):
|
|
with self.client_pool.get_and_release(destroy_on_fail=True) as client:
|
|
return client.raw_command(command, end_tokens)
|
|
|
|
def __setitem__(self, key: Key, value):
|
|
self.set(key, value, noreply=True)
|
|
|
|
def __getitem__(self, key):
|
|
value = self.get(key)
|
|
if value is None:
|
|
raise KeyError
|
|
return value
|
|
|
|
def __delitem__(self, key):
|
|
self.delete(key, noreply=True)
|
|
|
|
|
|
def _readline(sock: socket.socket, buf: bytes) -> Tuple[bytes, bytes]:
|
|
"""Read line of text from the socket.
|
|
|
|
Read a line of text (delimited by "\r\n") from the socket, and
|
|
return that line along with any trailing characters read from the
|
|
socket.
|
|
|
|
Args:
|
|
sock: Socket object, should be connected.
|
|
buf: Bytes, zero or more characters, returned from an earlier
|
|
call to _readline or _readvalue (pass an empty byte string on the
|
|
first call).
|
|
|
|
Returns:
|
|
A tuple of (buf, line) where line is the full line read from the
|
|
socket (minus the "\r\n" characters) and buf is any trailing
|
|
characters read after the "\r\n" was found (which may be an empty
|
|
byte string).
|
|
|
|
"""
|
|
chunks: List[bytes] = []
|
|
last_char = b""
|
|
|
|
while True:
|
|
# We're reading in chunks, so "\r\n" could appear in one chunk,
|
|
# or across the boundary of two chunks, so we check for both
|
|
# cases.
|
|
|
|
# This case must appear first, since the buffer could have
|
|
# later \r\n characters in it and we want to get the first \r\n.
|
|
if last_char == b"\r" and buf[0:1] == b"\n":
|
|
# Strip the last character from the last chunk.
|
|
chunks[-1] = chunks[-1][:-1]
|
|
return buf[1:], b"".join(chunks)
|
|
else:
|
|
token_pos = buf.find(b"\r\n")
|
|
if token_pos != -1:
|
|
# Note: 2 == len(b"\r\n")
|
|
before, after = buf[:token_pos], buf[token_pos + 2 :]
|
|
chunks.append(before)
|
|
return after, b"".join(chunks)
|
|
|
|
if buf:
|
|
chunks.append(buf)
|
|
last_char = buf[-1:]
|
|
|
|
buf = _recv(sock, RECV_SIZE)
|
|
if not buf:
|
|
raise MemcacheUnexpectedCloseError()
|
|
|
|
|
|
def _readvalue(sock: socket.socket, buf: bytes, size: int):
|
|
"""Read specified amount of bytes from the socket.
|
|
|
|
Read size bytes, followed by the "\r\n" characters, from the socket,
|
|
and return those bytes and any trailing bytes read after the "\r\n".
|
|
|
|
Args:
|
|
sock: Socket object, should be connected.
|
|
buf: String, zero or more characters, returned from an earlier
|
|
call to _readline or _readvalue (pass an empty string on the
|
|
first call).
|
|
size: Integer, number of bytes to read from the socket.
|
|
|
|
Returns:
|
|
A tuple of (buf, value) where value is the bytes read from the
|
|
socket (there will be exactly size bytes) and buf is trailing
|
|
characters read after the "\r\n" following the bytes (but not
|
|
including the \r\n).
|
|
|
|
"""
|
|
chunks = []
|
|
rlen = size + 2
|
|
while rlen - len(buf) > 0:
|
|
if buf:
|
|
rlen -= len(buf)
|
|
chunks.append(buf)
|
|
buf = _recv(sock, RECV_SIZE)
|
|
if not buf:
|
|
raise MemcacheUnexpectedCloseError()
|
|
|
|
# Now we need to remove the \r\n from the end. There are two cases we care
|
|
# about: the \r\n is all in the last buffer, or only the \n is in the last
|
|
# buffer, and we need to remove the \r from the penultimate buffer.
|
|
|
|
if rlen == 1:
|
|
# replace the last chunk with the same string minus the last character,
|
|
# which is always '\r' in this case.
|
|
chunks[-1] = chunks[-1][:-1]
|
|
else:
|
|
# Just remove the "\r\n" from the latest chunk
|
|
chunks.append(buf[: rlen - 2])
|
|
|
|
return buf[rlen:], b"".join(chunks)
|
|
|
|
|
|
def _readsegment(
|
|
sock: socket.socket, buf: bytes, end_tokens: bytes
|
|
) -> Tuple[bytes, bytes]:
|
|
"""Read a segment from the socket.
|
|
|
|
Read a segment from the socket, up to the first end_token sub-string/bytes,
|
|
and return that segment.
|
|
|
|
Args:
|
|
sock: Socket object, should be connected.
|
|
buf: bytes, zero or more bytes, returned from an earlier
|
|
call to _readline, _readsegment or _readvalue (pass an empty
|
|
byte-string on the first call).
|
|
end_tokens: bytes, indicates the end of the segment, generally this is
|
|
b"\\r\\n" for memcached.
|
|
|
|
Returns:
|
|
A tuple of (buf, line) where line is the full line read from the
|
|
socket (minus the end_tokens bytes) and buf is any trailing
|
|
characters read after the end_tokens was found (which may be an empty
|
|
bytes object).
|
|
|
|
"""
|
|
result = bytes()
|
|
|
|
while True:
|
|
|
|
tokens_pos = buf.find(end_tokens)
|
|
if tokens_pos != -1:
|
|
before, after = buf[:tokens_pos], buf[tokens_pos + len(end_tokens) :]
|
|
result += before
|
|
return after, result
|
|
|
|
buf = _recv(sock, RECV_SIZE)
|
|
if not buf:
|
|
raise MemcacheUnexpectedCloseError()
|
|
|
|
|
|
def _recv(sock: socket.socket, size: int) -> bytes:
|
|
"""sock.recv() with retry on EINTR"""
|
|
while True:
|
|
try:
|
|
return sock.recv(size)
|
|
except OSError as e:
|
|
if e.errno != errno.EINTR:
|
|
raise
|