Migrate from aioredis to redis with asyncio support

Author: moznuy

aredis_om/__init__.py

@@ -1,3 +1,4 @@
+from ._util import redis
 from .checks import has_redis_json, has_redisearch
 from .connections import get_redis_connection
 from .model.migrations.migrator import MigrationError, Migrator

aredis_om/_util.py

@@ -0,0 +1,17 @@
+import inspect
+
+
+def is_async_mode():
+    async def f():
+        """Unasync transforms async functions in sync functions"""
+        return None
+
+    return inspect.iscoroutinefunction(f)
+
+
+ASYNC_MODE = is_async_mode()
+
+if ASYNC_MODE:
+    from redis import asyncio as redis
+else:
+    import redis  # type: ignore

aredis_om/connections.py

@@ -1,19 +1,19 @@
 import os
 
-import aioredis
+from . import redis
 
 
 URL = os.environ.get("REDIS_OM_URL", None)
 
 
-def get_redis_connection(**kwargs) -> aioredis.Redis:
+def get_redis_connection(**kwargs) -> redis.Redis:
     # If someone passed in a 'url' parameter, or specified a REDIS_OM_URL
     # environment variable, we'll create the Redis client from the URL.
     url = kwargs.pop("url", URL)
     if url:
-        return aioredis.Redis.from_url(url, **kwargs)
+        return redis.Redis.from_url(url, **kwargs)
 
     # Decode from UTF-8 by default
     if "decode_responses" not in kwargs:
         kwargs["decode_responses"] = True
-    return aioredis.Redis(**kwargs)
+    return redis.Redis(**kwargs)

aredis_om/model/migrations/migrator.py

@@ -4,7 +4,7 @@
 from enum import Enum
 from typing import List, Optional
 
-from aioredis import Redis, ResponseError
+from ... import redis as redis_module
 
 
 log = logging.getLogger(__name__)
@@ -39,7 +39,7 @@ def schema_hash_key(index_name):
     return f"{index_name}:hash"
 
 
-async def create_index(redis: Redis, index_name, schema, current_hash):
+async def create_index(redis: redis_module.Redis, index_name, schema, current_hash):
     db_number = redis.connection_pool.connection_kwargs.get("db")
     if db_number and db_number > 0:
         raise MigrationError(
@@ -48,7 +48,7 @@ async def create_index(redis: Redis, index_name, schema, current_hash):
         )
     try:
         await redis.execute_command(f"ft.info {index_name}")
-    except ResponseError:
+    except redis_module.ResponseError:
         await redis.execute_command(f"ft.create {index_name} {schema}")
         await redis.set(schema_hash_key(index_name), current_hash)
     else:
@@ -67,7 +67,7 @@ class IndexMigration:
     schema: str
     hash: str
     action: MigrationAction
-    redis: Redis
+    redis: redis_module.Redis
     previous_hash: Optional[str] = None
 
     async def run(self):
@@ -79,13 +79,13 @@ async def run(self):
     async def create(self):
         try:
             await create_index(self.redis, self.index_name, self.schema, self.hash)
-        except ResponseError:
+        except redis_module.ResponseError:
             log.info("Index already exists: %s", self.index_name)
 
     async def drop(self):
         try:
             await self.redis.execute_command(f"FT.DROPINDEX {self.index_name}")
-        except ResponseError:
+        except redis_module.ResponseError:
             log.info("Index does not exist: %s", self.index_name)
 
 
@@ -115,7 +115,7 @@ async def detect_migrations(self):
 
             try:
                 await redis.execute_command("ft.info", cls.Meta.index_name)
-            except ResponseError:
+            except redis_module.ResponseError:
                 self.migrations.append(
                     IndexMigration(
                         name,

aredis_om/model/model.py

@@ -24,8 +24,6 @@
     no_type_check,
 )
 
-import aioredis
-from aioredis.client import Pipeline
 from pydantic import BaseModel, validator
 from pydantic.fields import FieldInfo as PydanticFieldInfo
 from pydantic.fields import ModelField, Undefined, UndefinedType
@@ -35,9 +33,10 @@
 from typing_extensions import Protocol, get_args, get_origin
 from ulid import ULID
 
+from .. import redis
+from .._util import ASYNC_MODE
 from ..checks import has_redis_json, has_redisearch
 from ..connections import get_redis_connection
-from ..unasync_util import ASYNC_MODE
 from .encoders import jsonable_encoder
 from .render_tree import render_tree
 from .token_escaper import TokenEscaper
@@ -975,7 +974,7 @@ class BaseMeta(Protocol):
     global_key_prefix: str
     model_key_prefix: str
     primary_key_pattern: str
-    database: aioredis.Redis
+    database: redis.Redis
     primary_key: PrimaryKey
     primary_key_creator_cls: Type[PrimaryKeyCreator]
     index_name: str
@@ -994,7 +993,7 @@ class DefaultMeta:
     global_key_prefix: Optional[str] = None
     model_key_prefix: Optional[str] = None
     primary_key_pattern: Optional[str] = None
-    database: Optional[aioredis.Redis] = None
+    database: Optional[redis.Redis] = None
     primary_key: Optional[PrimaryKey] = None
     primary_key_creator_cls: Optional[Type[PrimaryKeyCreator]] = None
     index_name: Optional[str] = None
@@ -1127,10 +1126,14 @@ async def update(self, **field_values):
         """Update this model instance with the specified key-value pairs."""
         raise NotImplementedError
 
-    async def save(self, pipeline: Optional[Pipeline] = None) -> "RedisModel":
+    async def save(
+        self, pipeline: Optional[redis.client.Pipeline] = None
+    ) -> "RedisModel":
         raise NotImplementedError
 
-    async def expire(self, num_seconds: int, pipeline: Optional[Pipeline] = None):
+    async def expire(
+        self, num_seconds: int, pipeline: Optional[redis.client.Pipeline] = None
+    ):
         if pipeline is None:
             db = self.db()
         else:
@@ -1241,7 +1244,7 @@ def get_annotations(cls):
     async def add(
         cls,
         models: Sequence["RedisModel"],
-        pipeline: Optional[Pipeline] = None,
+        pipeline: Optional[redis.client.Pipeline] = None,
         pipeline_verifier: Callable[..., Any] = verify_pipeline_response,
     ) -> Sequence["RedisModel"]:
         if pipeline is None:
@@ -1301,7 +1304,9 @@ def __init_subclass__(cls, **kwargs):
                     f"HashModels cannot index dataclass fields. Field: {name}"
                 )
 
-    async def save(self, pipeline: Optional[Pipeline] = None) -> "HashModel":
+    async def save(
+        self, pipeline: Optional[redis.client.Pipeline] = None
+    ) -> "HashModel":
         self.check()
         if pipeline is None:
             db = self.db()
@@ -1473,7 +1478,9 @@ def __init__(self, *args, **kwargs):
             )
         super().__init__(*args, **kwargs)
 
-    async def save(self, pipeline: Optional[Pipeline] = None) -> "JsonModel":
+    async def save(
+        self, pipeline: Optional[redis.client.Pipeline] = None
+    ) -> "JsonModel":
         self.check()
         if pipeline is None:
             db = self.db()

aredis_om/unasync_util.py

@@ -5,7 +5,6 @@
 
 ADDITIONAL_REPLACEMENTS = {
     "aredis_om": "redis_om",
-    "aioredis": "redis",
     ":tests.": ":tests_sync.",
     "pytest_asyncio": "pytest",
     "py_test_mark_asyncio": "py_test_mark_sync",

make_sync.py

@@ -23,8 +23,7 @@ include=[
 
 [tool.poetry.dependencies]
 python = "^3.7"
-redis = ">=3.5.3,<5.0.0"
-aioredis = "^2.0.0"
+redis = ">=4.2.0,<5.0.0"
 pydantic = "^1.8.2"
 click = "^8.0.1"
 six = "^1.16.0"

pyproject.toml

@@ -25,6 +25,7 @@
 from redis_om import has_redisearch
 from tests.conftest import py_test_mark_asyncio
 
+
 if not has_redisearch():
     pytestmark = pytest.mark.skip
 

tests/test_hash_model.py

@@ -27,6 +27,7 @@
 from redis_om import has_redis_json
 from tests.conftest import py_test_mark_asyncio
 
+
 if not has_redis_json():
     pytestmark = pytest.mark.skip