Support both sync and asyncio uses

aredis_om/model/__init__.py

@@ -1,2 +1,2 @@
 from .migrations.migrator import MigrationError, Migrator
-from .model import EmbeddedJsonModel, Field, HashModel, JsonModel, RedisModel
+from .model import EmbeddedJsonModel, Field, HashModel, JsonModel, RedisModel, NotFoundError

aredis_om/model/migrations/migrator.py

@@ -6,8 +6,6 @@ from typing import List, Optional
 
 from aioredis import Redis, ResponseError
 
-from aredis_om.model.model import model_registry
-
 
 log = logging.getLogger(__name__)
 
@@ -96,6 +94,10 @@ class Migrator:
         if self.module:
             import_submodules(self.module)
 
+        # Import this at run-time to avoid triggering import-time side effects,
+        # e.g. checks for RedisJSON, etc.
+        from aredis_om.model.model import model_registry
+
         for name, cls in model_registry.items():
             hash_key = schema_hash_key(cls.Meta.index_name)
             try:

aredis_om/model/model.py

@@ -10,7 +10,6 @@ from functools import reduce
 from typing import (
     AbstractSet,
     Any,
-    AsyncGenerator,
     Callable,
     Dict,
     List,
@@ -1295,7 +1294,7 @@ class HashModel(RedisModel, abc.ABC):
         return self
 
     @classmethod
-    async def all_pks(cls) -> AsyncGenerator[str, None]:  # type: ignore
+    async def all_pks(cls):  # type: ignore
         key_prefix = cls.make_key(cls._meta.primary_key_pattern.format(pk=""))
         # TODO: We assume the key ends with the default separator, ":" -- when
         #  we make the separator configurable, we need to update this as well.
@@ -1437,13 +1436,16 @@ class HashModel(RedisModel, abc.ABC):
 
 class JsonModel(RedisModel, abc.ABC):
     def __init_subclass__(cls, **kwargs):
-        if not has_redis_json(cls.db()):
+        # Generate the RediSearch schema once to validate fields.
+        cls.redisearch_schema()
+
+    def __init__(self, *args, **kwargs):
+        if not has_redis_json(self.db()):
             log.error(
                 "Your Redis instance does not have the RedisJson module "
                 "loaded. JsonModel depends on RedisJson."
             )
-        # Generate the RediSearch schema once to validate fields.
-        cls.redisearch_schema()
+        super().__init__(*args, **kwargs)
 
     async def save(self, pipeline: Optional[Pipeline] = None) -> "JsonModel":
         self.check()