unix_socket_path in RedisSettings (#336)

* from rebase

* update: testing requirements

* Update tests/conftest.py

Co-authored-by: Samuel Colvin <samcolvin@gmail.com>

* update: testing requirements

Co-authored-by: Samuel Colvin <samcolvin@gmail.com>

Author: lazToum

arq/connections.py

@@ -5,7 +5,7 @@
 from datetime import datetime, timedelta
 from operator import attrgetter
 from typing import TYPE_CHECKING, Any, Callable, List, Optional, Tuple, Union
-from urllib.parse import urlparse
+from urllib.parse import parse_qs, urlparse
 from uuid import uuid4
 
 from redis.asyncio import ConnectionPool, Redis
@@ -29,6 +29,7 @@ class RedisSettings:
 
     host: Union[str, List[Tuple[str, int]]] = 'localhost'
     port: int = 6379
+    unix_socket_path: Optional[str] = None
     database: int = 0
     username: Optional[str] = None
     password: Optional[str] = None
@@ -49,14 +50,21 @@ class RedisSettings:
     @classmethod
     def from_dsn(cls, dsn: str) -> 'RedisSettings':
         conf = urlparse(dsn)
-        assert conf.scheme in {'redis', 'rediss'}, 'invalid DSN scheme'
+        assert conf.scheme in {'redis', 'rediss', 'unix'}, 'invalid DSN scheme'
+        query_db = parse_qs(conf.query).get('db')
+        if query_db:
+            # e.g. redis://localhost:6379?db=1
+            database = int(query_db[0])
+        else:
+            database = int(conf.path.lstrip('/')) if conf.path else 0
         return RedisSettings(
             host=conf.hostname or 'localhost',
             port=conf.port or 6379,
             ssl=conf.scheme == 'rediss',
             username=conf.username,
             password=conf.password,
-            database=int((conf.path or '0').strip('/')),
+            database=database,
+            unix_socket_path=conf.path if conf.scheme == 'unix' else None,
         )
 
     def __repr__(self) -> str:
@@ -230,6 +238,7 @@ def pool_factory(*args: Any, **kwargs: Any) -> ArqRedis:
             ArqRedis,
             host=settings.host,
             port=settings.port,
+            unix_socket_path=settings.unix_socket_path,
             socket_connect_timeout=settings.conn_timeout,
             ssl=settings.ssl,
             ssl_keyfile=settings.ssl_keyfile,

requirements/testing.in

@@ -8,3 +8,4 @@ pytest-mock>=3,<4
 pytest-sugar>=0.9,<1
 pytest-timeout>=2,<3
 pytz
+redislite

requirements/testing.txt

@@ -4,10 +4,14 @@
 #
 #    pip-compile --output-file=requirements/testing.txt requirements/testing.in
 #
+async-timeout==4.0.2
+    # via redis
 attrs==22.1.0
     # via pytest
 coverage[toml]==6.4.4
     # via -r requirements/testing.in
+deprecated==1.2.13
+    # via redis
 dirty-equals==0.4
     # via -r requirements/testing.in
 iniconfig==1.1.1
@@ -18,8 +22,11 @@ packaging==21.3
     # via
     #   pytest
     #   pytest-sugar
+    #   redis
 pluggy==1.0.0
     # via pytest
+psutil==5.9.1
+    # via redislite
 py==1.11.0
     # via pytest
 pydantic==1.9.2
@@ -45,6 +52,10 @@ pytz==2022.2.1
     # via
     #   -r requirements/testing.in
     #   dirty-equals
+redis==4.2.2
+    # via redislite
+redislite==6.2.805324
+    # via -r requirements/testing.in
 termcolor==1.1.0
     # via pytest-sugar
 tomli==2.0.1
@@ -53,3 +64,8 @@ tomli==2.0.1
     #   pytest
 typing-extensions==4.3.0
     # via pydantic
+wrapt==1.14.1
+    # via deprecated
+
+# The following packages are considered to be unsafe in a requirements file:
+# setuptools

tests/conftest.py

@@ -5,6 +5,7 @@
 
 import msgpack
 import pytest
+from redislite import Redis
 
 from arq.connections import ArqRedis, create_pool
 from arq.worker import Worker
@@ -30,6 +31,13 @@ async def arq_redis(loop):
     await redis_.close(close_connection_pool=True)
 
 
+@pytest.fixture
+async def unix_socket_path(loop, tmp_path):
+    rdb = Redis(str(tmp_path / 'redis_test.db'))
+    yield rdb.socket_file
+    rdb.close()
+
+
 @pytest.fixture
 async def arq_redis_msgpack(loop):
     redis_ = ArqRedis(

tests/test_jobs.py

@@ -76,6 +76,17 @@ async def foobar(ctx, *args, **kwargs):
     ]
 
 
+async def test_enqueue_job_with_unix_socket(worker, unix_socket_path):
+    """Test initializing arq_redis using a unix socket connection, and the worker using it."""
+    settings = RedisSettings(unix_socket_path=unix_socket_path)
+    arq_redis = await create_pool(settings, default_queue_name='socket_queue')
+    await test_enqueue_job(
+        arq_redis,
+        lambda functions, **_: worker(functions=functions, arq_redis=arq_redis, queue_name=None),
+        queue_name=None,
+    )
+
+
 async def test_enqueue_job_alt_queue(arq_redis: ArqRedis, worker):
     await test_enqueue_job(arq_redis, worker, queue_name='custom_queue')
 

tests/test_utils.py

@@ -18,8 +18,8 @@ def test_settings_changed():
     settings = RedisSettings(port=123)
     assert settings.port == 123
     assert (
-        "RedisSettings(host='localhost', port=123, database=0, username=None, password=None, ssl=False, "
-        "ssl_keyfile=None, ssl_certfile=None, ssl_cert_reqs='required', ssl_ca_certs=None, "
+        "RedisSettings(host='localhost', port=123, unix_socket_path=None, database=0, username=None, password=None, "
+        "ssl=False, ssl_keyfile=None, ssl_certfile=None, ssl_cert_reqs='required', ssl_ca_certs=None, "
         'ssl_ca_data=None, ssl_check_hostname=False, conn_timeout=1, conn_retries=5, conn_retry_delay=1, '
         "sentinel=False, sentinel_master='mymaster')"
     ) == str(settings)
@@ -141,6 +141,14 @@ def parse_redis_settings(cls, v):
     assert s4.redis_settings.username == 'user'
     assert s4.redis_settings.password == 'pass'
 
+    s5 = Settings(redis_settings={'unix_socket_path': '/tmp/redis.sock'})
+    assert s5.redis_settings.unix_socket_path == '/tmp/redis.sock'
+    assert s5.redis_settings.database == 0
+
+    s6 = Settings(redis_settings='unix:///tmp/redis.socket?db=6')
+    assert s6.redis_settings.unix_socket_path == '/tmp/redis.socket'
+    assert s6.redis_settings.database == 6
+
 
 def test_ms_to_datetime_tz(env: SetEnv):
     arq.utils.get_tz.cache_clear()