From ea38eb01b21844aadc99956b54cd3b65ffd458a1 Mon Sep 17 00:00:00 2001
From: Trenton H <797416+stumpylog@users.noreply.github.com>
Date: Fri, 2 Dec 2022 13:54:15 -0800
Subject: [PATCH] Adds support for database number specification

---
 src/paperless/settings.py            | 13 +++++++++++--
 src/paperless/tests/test_settings.py | 19 +++++++++++++++++++
 2 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/src/paperless/settings.py b/src/paperless/settings.py
index eef7344da..3271ae7e6 100644
--- a/src/paperless/settings.py
+++ b/src/paperless/settings.py
@@ -83,12 +83,21 @@ def _parse_redis_url(env_redis: Optional[str]) -> Tuple[str]:
     if "unix" in env_redis.lower():
         # channels_redis socket format, looks like:
         # "unix:///path/to/redis.sock"
-        return (f"redis+socket:{path}", env_redis)
+        if "?db=" in env_redis:
+            path, number = path.split("?db=")
+            return (f"redis+socket:{path}?virtual_host={number}", env_redis)
+        else:
+            return (f"redis+socket:{path}", env_redis)
 
     elif "+socket" in env_redis.lower():
         # celery socket style, looks like:
         # "redis+socket:///path/to/redis.sock"
-        return (env_redis, f"unix:{path}")
+        if "?virtual_host=" in env_redis:
+            # Virtual host (aka db number)
+            path, number = path.split("?virtual_host=")
+            return (env_redis, f"unix:{path}?db={number}")
+        else:
+            return (env_redis, f"unix:{path}")
 
     # Not a socket
     return (env_redis, env_redis)
diff --git a/src/paperless/tests/test_settings.py b/src/paperless/tests/test_settings.py
index fa839299f..346f2d098 100644
--- a/src/paperless/tests/test_settings.py
+++ b/src/paperless/tests/test_settings.py
@@ -97,7 +97,9 @@ class TestIgnoreDateParsing(TestCase):
         """
 
         for input, expected in [
+            # Nothing is set
             (None, ("redis://localhost:6379", "redis://localhost:6379")),
+            # celery style
             (
                 "redis+socket:///run/redis/redis.sock",
                 (
@@ -105,6 +107,7 @@ class TestIgnoreDateParsing(TestCase):
                     "unix:///run/redis/redis.sock",
                 ),
             ),
+            # redis-py / channels-redis style
             (
                 "unix:///run/redis/redis.sock",
                 (
@@ -112,6 +115,22 @@ class TestIgnoreDateParsing(TestCase):
                     "unix:///run/redis/redis.sock",
                 ),
             ),
+            # celery style with db
+            (
+                "redis+socket:///run/redis/redis.sock?virtual_host=5",
+                (
+                    "redis+socket:///run/redis/redis.sock?virtual_host=5",
+                    "unix:///run/redis/redis.sock?db=5",
+                ),
+            ),
+            # redis-py / channels-redis style with db
+            (
+                "unix:///run/redis/redis.sock?db=10",
+                (
+                    "redis+socket:///run/redis/redis.sock?virtual_host=10",
+                    "unix:///run/redis/redis.sock?db=10",
+                ),
+            ),
         ]:
             result = _parse_redis_url(input)
             self.assertTupleEqual(expected, result)