Add async pydantic redis

Author: Tinitto

pydantic_redis/__init__.py

@@ -1,9 +1,8 @@
 """Entry point for redisy"""
 
-from pydantic_redis.shared.config import RedisConfig
-from pydantic_redis.syncio.model import Model
-from pydantic_redis.syncio.store import Store
+from pydantic_redis.syncio import Store, Model, RedisConfig
+import pydantic_redis.asyncio
 
-__all__ = [Store, RedisConfig, Model]
+__all__ = [Store, RedisConfig, Model, asyncio]
 
 __version__ = "0.2.0"

pydantic_redis/asyncio/__init__.py

@@ -0,0 +1,7 @@
+"""Package containing the async version of pydantic_redis"""
+
+from .model import Model
+from .store import Store
+from ..shared.config import RedisConfig
+
+__all__ = [Model, Store, RedisConfig]

pydantic_redis/asyncio/model.py

@@ -0,0 +1,143 @@
+"""Module containing the model classes"""
+from typing import Optional, List, Any, Union, Dict, Tuple, Type
+
+import redis.asyncio
+
+from pydantic_redis.shared.model import AbstractModel
+from pydantic_redis.shared.model.insert_utils import insert_on_pipeline
+from pydantic_redis.shared.model.prop_utils import get_primary_key, get_table_index_key
+from pydantic_redis.shared.model.select_utils import (
+    select_all_fields_all_ids,
+    select_all_fields_some_ids,
+    select_some_fields_all_ids,
+    select_some_fields_some_ids,
+    parse_select_response,
+)
+
+from .store import Store
+
+
+class Model(AbstractModel):
+    """
+    The section in the store that saves rows of the same kind
+    """
+
+    _store: Store
+
+    @classmethod
+    async def insert(
+        cls,
+        data: Union[List[AbstractModel], AbstractModel],
+        life_span_seconds: Optional[float] = None,
+    ):
+        """
+        Inserts a given row or sets of rows into the table
+        """
+        store = cls.get_store()
+        life_span = (
+            life_span_seconds
+            if life_span_seconds is not None
+            else store.life_span_in_seconds
+        )
+
+        async with store.redis_store.pipeline(transaction=True) as pipeline:
+            data_list = []
+
+            if isinstance(data, list):
+                data_list = data
+            elif isinstance(data, AbstractModel):
+                data_list = [data]
+
+            for record in data_list:
+                insert_on_pipeline(
+                    model=cls,
+                    _id=None,
+                    pipeline=pipeline,
+                    record=record,
+                    life_span=life_span,
+                )
+
+            return await pipeline.execute()
+
+    @classmethod
+    async def update(
+        cls, _id: Any, data: Dict[str, Any], life_span_seconds: Optional[float] = None
+    ):
+        """
+        Updates a given row or sets of rows in the table
+        """
+        store = cls.get_store()
+        life_span = (
+            life_span_seconds
+            if life_span_seconds is not None
+            else store.life_span_in_seconds
+        )
+        async with store.redis_store.pipeline(transaction=True) as pipeline:
+            if isinstance(data, dict):
+                insert_on_pipeline(
+                    model=cls,
+                    _id=_id,
+                    pipeline=pipeline,
+                    record=data,
+                    life_span=life_span,
+                )
+
+            return await pipeline.execute()
+
+    @classmethod
+    async def delete(cls, ids: Union[Any, List[Any]]):
+        """
+        deletes a given row or sets of rows in the table
+        """
+        store = cls.get_store()
+
+        async with store.redis_store.pipeline() as pipeline:
+            primary_keys = []
+
+            if isinstance(ids, list):
+                primary_keys = ids
+            elif ids is not None:
+                primary_keys = [ids]
+
+            names = [
+                get_primary_key(model=cls, primary_key_value=primary_key_value)
+                for primary_key_value in primary_keys
+            ]
+            pipeline.delete(*names)
+            # remove the primary keys from the index
+            table_index_key = get_table_index_key(model=cls)
+            pipeline.srem(table_index_key, *names)
+            return await pipeline.execute()
+
+    @classmethod
+    async def select(
+        cls,
+        columns: Optional[List[str]] = None,
+        ids: Optional[List[Any]] = None,
+        **kwargs,
+    ):
+        """
+        Selects given rows or sets of rows in the table
+        """
+        if columns is None and ids is None:
+            response = await select_all_fields_all_ids(model=cls)
+
+        elif columns is None and isinstance(ids, list):
+            response = await select_all_fields_some_ids(model=cls, ids=ids)
+
+        elif isinstance(columns, list) and ids is None:
+            response = await select_some_fields_all_ids(model=cls, fields=columns)
+
+        elif isinstance(columns, list) and isinstance(ids, list):
+            response = await select_some_fields_some_ids(
+                model=cls, fields=columns, ids=ids
+            )
+
+        else:
+            raise ValueError(
+                f"columns {columns}, ids: {ids} should be either None or lists"
+            )
+
+        return parse_select_response(
+            model=cls, response=response, as_models=(columns is None)
+        )

pydantic_redis/asyncio/store.py

@@ -0,0 +1,25 @@
+"""Module containing the store class for async io"""
+from typing import Dict, Type, TYPE_CHECKING
+
+from redis import asyncio as redis
+
+from ..shared.store import AbstractStore
+
+if TYPE_CHECKING:
+    from .model import Model
+
+
+class Store(AbstractStore):
+    """
+    A store that allows a declarative way of querying for data in redis
+    """
+
+    models: Dict[str, Type["Model"]] = {}
+
+    def _connect_to_redis(self) -> redis.Redis:
+        """Connects the store to redis, returning a proper connection"""
+        return redis.from_url(
+            self.redis_config.redis_url,
+            encoding=self.redis_config.encoding,
+            decode_responses=True,
+        )

pydantic_redis/syncio/__init__.py

@@ -1 +1,6 @@
 """Package containing the synchronous and thus default version of pydantic_redis"""
+from .model import Model
+from .store import Store
+from ..shared.config import RedisConfig
+
+__all__ = [Model, Store, RedisConfig]

pydantic_redis/syncio/model.py

@@ -87,7 +87,8 @@ def delete(cls, ids: Union[Any, List[Any]]):
         """
         deletes a given row or sets of rows in the table
         """
-        with cls._store.redis_store.pipeline() as pipeline:
+        store = cls.get_store()
+        with store.redis_store.pipeline() as pipeline:
             primary_keys = []
 
             if isinstance(ids, list):

requirements.txt

@@ -10,3 +10,4 @@ twine==3.8.0
 black==22.8.0
 pre-commit
 build
+pytest-asyncio

test/conftest.py

@@ -3,19 +3,26 @@
 from typing import Tuple, List, Optional
 
 import pytest
+import pytest_asyncio
 import redislite
 from pytest_lazyfixture import lazy_fixture
 
-from pydantic_redis import Store, RedisConfig, Model
+from pydantic_redis import syncio as syn, asyncio as asy
 
 
-class Author(Model):
+class Author(syn.Model):
     _primary_key_field: str = "name"
     name: str
     active_years: Tuple[int, int]
 
 
-class Book(Model):
+class AsyncAuthor(asy.Model):
+    _primary_key_field: str = "name"
+    name: str
+    active_years: Tuple[int, int]
+
+
+class Book(syn.Model):
     _primary_key_field: str = "title"
     title: str
     author: Author
@@ -25,7 +32,17 @@ class Book(Model):
     in_stock: bool = True
 
 
-class Library(Model):
+class AsyncBook(asy.Model):
+    _primary_key_field: str = "title"
+    title: str
+    author: AsyncAuthor
+    rating: float
+    published_on: date
+    tags: List[str] = []
+    in_stock: bool = True
+
+
+class Library(syn.Model):
     # the _primary_key_field is mandatory
     _primary_key_field: str = "name"
     name: str
@@ -36,11 +53,27 @@ class Library(Model):
     new: Tuple[Book, Author, Book, int] = None
 
 
+class AsyncLibrary(asy.Model):
+    # the _primary_key_field is mandatory
+    _primary_key_field: str = "name"
+    name: str
+    address: str
+    books: List[AsyncBook] = None
+    lost: Optional[List[AsyncBook]] = None
+    popular: Optional[Tuple[AsyncBook, AsyncBook]] = None
+    new: Tuple[AsyncBook, AsyncAuthor, AsyncBook, int] = None
+
+
 authors = {
     "charles": Author(name="Charles Dickens", active_years=(1220, 1280)),
     "jane": Author(name="Jane Austen", active_years=(1580, 1640)),
 }
 
+async_authors = {
+    "charles": AsyncAuthor(name="Charles Dickens", active_years=(1220, 1280)),
+    "jane": AsyncAuthor(name="Jane Austen", active_years=(1580, 1640)),
+}
+
 books = [
     Book(
         title="Oliver Twist",
@@ -74,6 +107,40 @@ class Library(Model):
     ),
 ]
 
+async_books = [
+    AsyncBook(
+        title="Oliver Twist",
+        author=authors["charles"],
+        published_on=date(year=1215, month=4, day=4),
+        in_stock=False,
+        rating=2,
+        tags=["Classic"],
+    ),
+    AsyncBook(
+        title="Great Expectations",
+        author=authors["charles"],
+        published_on=date(year=1220, month=4, day=4),
+        rating=5,
+        tags=["Classic"],
+    ),
+    AsyncBook(
+        title="Jane Eyre",
+        author=authors["charles"],
+        published_on=date(year=1225, month=6, day=4),
+        in_stock=False,
+        rating=3.4,
+        tags=["Classic", "Romance"],
+    ),
+    AsyncBook(
+        title="Wuthering Heights",
+        author=authors["jane"],
+        published_on=date(year=1600, month=4, day=4),
+        rating=4.0,
+        tags=["Classic", "Romance"],
+    ),
+]
+
+# sync fixtures
 redis_store_fixture = [(lazy_fixture("redis_store"))]
 books_fixture = [(lazy_fixture("redis_store"), book) for book in books]
 update_books_fixture = [
@@ -88,6 +155,23 @@ class Library(Model):
     (lazy_fixture("redis_store"), book.title) for book in books[-1:]
 ]
 
+# async fixtures
+async_redis_store_fixture = [(lazy_fixture("async_redis_store"))]
+async_books_fixture = [
+    (lazy_fixture("async_redis_store"), book) for book in async_books
+]
+async_update_books_fixture = [
+    (
+        lazy_fixture("async_redis_store"),
+        book.title,
+        {"author": authors["jane"], "in_stock": not book.in_stock},
+    )
+    for book in async_books[-1:]
+]
+async_delete_books_fixture = [
+    (lazy_fixture("async_redis_store"), book.title) for book in async_books[-1:]
+]
+
 
 @pytest.fixture()
 def unused_tcp_port():
@@ -110,13 +194,28 @@ def redis_server(unused_tcp_port):
 @pytest.fixture()
 def redis_store(redis_server):
     """Sets up a redis store using the redis_server fixture and adds the book model to it"""
-    store = Store(
+    store = syn.Store(
         name="sample",
-        redis_config=RedisConfig(port=redis_server, db=1),
+        redis_config=syn.RedisConfig(port=redis_server, db=1),
         life_span_in_seconds=3600,
     )
     store.register_model(Book)
     store.register_model(Author)
     store.register_model(Library)
     yield store
     store.redis_store.flushall()
+
+
+@pytest_asyncio.fixture
+async def async_redis_store(redis_server):
+    """Sets up a redis store using the redis_server fixture and adds the book model to it"""
+    store = asy.Store(
+        name="sample",
+        redis_config=syn.RedisConfig(port=redis_server, db=1),
+        life_span_in_seconds=3600,
+    )
+    store.register_model(AsyncBook)
+    store.register_model(AsyncAuthor)
+    store.register_model(AsyncLibrary)
+    yield store
+    await store.redis_store.flushall()