diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 715d9468..75256c16 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -68,8 +68,12 @@ jobs: pip install -U setuptools pip install -r requirements.txt pip install codecov - - name: Run memcached service - uses: jkeys089/actions-memcached@master + - name: Run memcached + run: | + sudo apt-get install -y memcached + sudo systemctl disable --now memcached + memcached -d -p 11211 -l 127.0.0.1 + memcached -d -s /tmp/memcached.sock - name: Run tests run: pytest - run: python -m coverage xml diff --git a/aiomcache/client.py b/aiomcache/client.py index 78896f48..27a92fa8 100644 --- a/aiomcache/client.py +++ b/aiomcache/client.py @@ -45,7 +45,7 @@ async def wrapper(self: _Client, *args: _P.args, # type: ignore[misc] class FlagClient(Generic[_T]): - def __init__(self, host: str, port: int = 11211, *, + def __init__(self, host: str, port: int, *, pool_size: int = 2, pool_minsize: Optional[int] = None, conn_args: Optional[Mapping[str, Any]] = None, get_flag_handler: Optional[_GetFlagHandler[_T]] = None, @@ -54,7 +54,7 @@ def __init__(self, host: str, port: int = 11211, *, Creates new Client instance. :param host: memcached host - :param port: memcached port + :param port: memcached port (-1 implies local unix socket) :param pool_size: max connection pool size :param pool_minsize: min connection pool size :param conn_args: extra arguments passed to @@ -492,9 +492,28 @@ async def flush_all(self, conn: Connection) -> None: class Client(FlagClient[bytes]): - def __init__(self, host: str, port: int = 11211, *, + @overload + def __init__(self, host: str, port: int, *, + pool_size: int = 2, pool_minsize: Optional[int] = None, + conn_args: Optional[Mapping[str, Any]] = None): + ... + + @overload + def __init__(self, *, path: str, pool_size: int = 2, pool_minsize: Optional[int] = None, conn_args: Optional[Mapping[str, Any]] = None): + ... + + def __init__(self, host: str = '127.0.0.1', port: int = 11211, *, + path: str = '', + pool_size: int = 2, pool_minsize: Optional[int] = None, + conn_args: Optional[Mapping[str, Any]] = None): + + # unlikely to provide host/port with the overloads, but still need to deal with it + if path: + host = path + port = -1 + super().__init__(host, port, pool_size=pool_size, pool_minsize=pool_minsize, conn_args=conn_args, get_flag_handler=None, set_flag_handler=None) diff --git a/aiomcache/pool.py b/aiomcache/pool.py index edee644f..55a1e266 100644 --- a/aiomcache/pool.py +++ b/aiomcache/pool.py @@ -12,7 +12,7 @@ class Connection(NamedTuple): class MemcachePool: def __init__(self, host: str, port: int, *, minsize: int, maxsize: int, conn_args: Optional[Mapping[str, Any]] = None): - self._host = host + self._target = host self._port = port self._minsize = minsize self._maxsize = maxsize @@ -67,8 +67,12 @@ def release(self, conn: Connection) -> None: async def _create_new_conn(self) -> Optional[Connection]: if self.size() < self._maxsize: - reader, writer = await asyncio.open_connection( - self._host, self._port, **self.conn_args) + if self._port == -1: # -1 implies unix socket + reader, writer = await asyncio.open_unix_connection(self._target, **self.conn_args) + else: + reader, writer = await asyncio.open_connection(self._target, self._port, + **self.conn_args) + if self.size() < self._maxsize: return Connection(reader, writer) else: diff --git a/examples/simple.py b/examples/simple.py index e182d420..39ae898e 100644 --- a/examples/simple.py +++ b/examples/simple.py @@ -5,6 +5,7 @@ async def hello_aiomcache() -> None: mc = aiomcache.Client("127.0.0.1", 11211) + # Can also use Client(path="/var/run/memcached/mc.sock") for unix sockets await mc.set(b"some_key", b"Some value") value = await mc.get(b"some_key") print(value) diff --git a/tests/commands_test.py b/tests/commands_test.py index ae5768e6..cc0f9553 100644 --- a/tests/commands_test.py +++ b/tests/commands_test.py @@ -8,6 +8,7 @@ from aiomcache import Client, FlagClient from aiomcache.exceptions import ClientException, ValidationException +from .conftest import McacheUnixParams from .flag_helper import FlagHelperDemo @@ -427,3 +428,15 @@ async def test_flag_handler_invoked_only_when_expected( assert orig_get_count + 1 == demo_flag_helper.get_invocation_count assert orig_set_count + 1 == demo_flag_helper.set_invocation_count + + +async def test_client_unix_socket_set_get( + mcache_unix_params: McacheUnixParams, +) -> None: + client = Client(path=mcache_unix_params["path"]) + try: + await client.set(b"key", b"value") + v = await client.get(b"key") + assert v == b"value" + finally: + await client.close() diff --git a/tests/conftest.py b/tests/conftest.py index a5832541..aa40d7e2 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,12 +1,9 @@ -import contextlib import socket import sys -import time import uuid -from typing import Any, AsyncIterator, Callable, Iterator, TypedDict +from typing import Any, AsyncIterator, Callable, TypedDict import docker as docker_mod -import memcache import pytest import aiomcache @@ -23,6 +20,10 @@ class McacheParams(TypedDict): port: int +class McacheUnixParams(TypedDict): + path: str + + class ServerParams(TypedDict): Id: NotRequired[str] host: str @@ -30,6 +31,12 @@ class ServerParams(TypedDict): mcache_params: McacheParams +class UnixServerParams(TypedDict): + Id: NotRequired[str] + path: str + mcache_unix_params: McacheUnixParams + + mcache_server_option = "localhost" @@ -72,58 +79,35 @@ def mcache_server_actual(host: str, port: int = 11211) -> ServerParams: } -@contextlib.contextmanager -def mcache_server_docker( # type: ignore[no-any-unimported] - unused_port: Callable[[], int], docker: docker_mod.Client, session_id: str -) -> Iterator[ServerParams]: - docker.images.pull("memcached:alpine") - container = docker.containers.run( - image='memcached:alpine', - name='memcached-test-server-{}'.format(session_id), - ports={"11211/tcp": None}, - detach=True, - ) - try: - container.start() - container.reload() - net_settings = container.attrs["NetworkSettings"] - host = net_settings["IPAddress"] - port = int(net_settings["Ports"]["11211/tcp"][0]["HostPort"]) - mcache_params: McacheParams = {"host": host, "port": port} - delay = 0.001 - for _i in range(10): - try: - conn = memcache.Client(["{host}:{port}".format_map(mcache_params)]) - conn.get_stats() - break - except Exception: - time.sleep(delay) - delay *= 2 - else: - pytest.fail("Cannot start memcached") - ret: ServerParams = { - "Id": container.id, - "host": host, - "port": port, - "mcache_params": mcache_params - } - time.sleep(0.1) - yield ret - finally: - container.kill() - container.remove() - - @pytest.fixture(scope='session') def mcache_server() -> ServerParams: return mcache_server_actual("localhost") +def mcache_unix_server_actual(path: str) -> UnixServerParams: + return { + "path": path, + "mcache_unix_params": {"path": path} + } + + +@pytest.fixture(scope='session') +def mcache_unix_server(session_id: str) -> UnixServerParams: + # if starting memcached via systemd, ensure privatetmp is not on + sock_path = '/tmp/memcached.sock' # noqa: S108 + return mcache_unix_server_actual(sock_path) + + @pytest.fixture def mcache_params(mcache_server: ServerParams) -> McacheParams: return mcache_server["mcache_params"] +@pytest.fixture +def mcache_unix_params(mcache_unix_server: UnixServerParams) -> McacheUnixParams: + return mcache_unix_server["mcache_unix_params"] + + @pytest.fixture async def mcache(mcache_params: McacheParams) -> AsyncIterator[aiomcache.Client]: client = aiomcache.Client(**mcache_params) @@ -131,6 +115,13 @@ async def mcache(mcache_params: McacheParams) -> AsyncIterator[aiomcache.Client] await client.close() +@pytest.fixture +async def mcache_unix(mcache_unix_params: McacheUnixParams) -> AsyncIterator[aiomcache.Client]: + client = aiomcache.Client(path=mcache_unix_params["path"]) + yield client + await client.close() + + test_only_demo_flag_helper = FlagHelperDemo() diff --git a/tests/pool_test.py b/tests/pool_test.py index bc571067..1b5506c3 100644 --- a/tests/pool_test.py +++ b/tests/pool_test.py @@ -6,7 +6,7 @@ from aiomcache.client import Client, acquire from aiomcache.pool import Connection, MemcachePool -from .conftest import McacheParams +from .conftest import McacheParams, McacheUnixParams async def test_pool_creation(mcache_params: McacheParams) -> None: @@ -139,9 +139,8 @@ async def test_0_minsize(mcache_params: McacheParams) -> None: await pool.clear() -async def test_bad_connection(mcache_params: McacheParams) -> None: - pool = MemcachePool(minsize=5, maxsize=1, **mcache_params) - pool._host = "INVALID_HOST" +async def test_bad_connection() -> None: + pool = MemcachePool("INVALID_HOST", 11211, minsize=5, maxsize=1) assert pool.size() == 0 with pytest.raises(socket.gaierror): conn = await pool.acquire() @@ -149,3 +148,21 @@ async def test_bad_connection(mcache_params: McacheParams) -> None: assert isinstance(conn.writer, asyncio.StreamWriter) pool.release(conn) assert pool.size() == 0 + + +async def test_pool_unix_socket_acquire_release( + mcache_unix_params: McacheUnixParams, +) -> None: + pool = MemcachePool(mcache_unix_params["path"], -1, minsize=1, maxsize=5) + conn = await pool.acquire() + assert isinstance(conn.reader, asyncio.StreamReader) + assert isinstance(conn.writer, asyncio.StreamWriter) + pool.release(conn) + await pool.clear() + + +async def test_pool_unix_socket_bad_path() -> None: + pool = MemcachePool("/tmp/nonexistent-mc.sock", -1, minsize=1, maxsize=1) # noqa: S108 + assert pool.size() == 0 + with pytest.raises(FileNotFoundError): + await pool.acquire()