From da988ec77a327f155a05f5d67f445baf3c4a6bf4 Mon Sep 17 00:00:00 2001 From: Martin Date: Sat, 17 Dec 2022 14:34:24 +0300 Subject: [PATCH] Close #20: Add pagination to select --- CHANGELOG.md | 10 ++ README.md | 143 +++++++++++++------- pydantic_redis/__init__.py | 2 +- pydantic_redis/asyncio/model.py | 35 ++--- pydantic_redis/shared/lua_scripts.py | 137 +++++++++++++++++++ pydantic_redis/shared/model/delete_utils.py | 32 +++++ pydantic_redis/shared/model/insert_utils.py | 10 +- pydantic_redis/shared/model/select_utils.py | 77 +++++++++-- pydantic_redis/shared/store.py | 18 +++ pydantic_redis/syncio/model.py | 32 ++--- setup.py | 2 +- test/test_async_pydantic_redis.py | 64 ++++++++- test/test_benchmarks.py | 14 ++ test/test_pydantic_redis.py | 60 +++++++- 14 files changed, 530 insertions(+), 106 deletions(-) create mode 100644 pydantic_redis/shared/model/delete_utils.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 985cd862..0fbbd67f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,16 @@ and this project adheres to [Semantic Versioning](http://semver.org/). ## [Unreleased] +## [0.4.0] - 2022-12-17 + +### Added + +- Added pagination + +### Changed + +- Changed redis index to use sorted sets instead of ordinary sets + ## [0.3.0] - 2022-12-15 ### Added diff --git a/README.md b/README.md index df9256ae..ab2ecbf5 100644 --- a/README.md +++ b/README.md @@ -38,7 +38,8 @@ from pydantic_redis import RedisConfig, Model, Store class Author(Model): """ - An Author model, just like a pydantic model with appropriate type annotations + An Author model, just like a pydantic model with appropriate + type annotations NOTE: The `_primary_key_field` is mandatory """ _primary_key_field: str = 'name' @@ -51,18 +52,24 @@ class Book(Model): A Book model. Models can have the following field types - - The usual i.e. float, int, dict, list, date, str, dict, Optional etc as long as they are serializable by orjson + - The usual i.e. float, int, dict, list, date, str, dict, Optional etc + as long as they are serializable by orjson - Nested models e.g. `author: Author` or `author: Optional[Author]` - - List of nested models e.g. `authors: List[Author]` or `authors: Optional[List[Author]]` - - Tuples including nested models e.g. `access_log: Tuple[Author, date]` or `access_log: Optional[Tuple[Author, date]]]` + - List of nested models e.g. `authors: List[Author]` + or `authors: Optional[List[Author]]` + - Tuples including nested models e.g. `access_log: Tuple[Author, date]` + or `access_log: Optional[Tuple[Author, date]]]` - NOTE: 1. Any nested model whether plain or in a list or tuple will automatically inserted into the redis store - when the parent model is inserted. e.g. a Book with an author field, when inserted, will also insert - the author. The author can then be queried directly if that's something one wishes to do. + NOTE: 1. Any nested model whether plain or in a list or tuple will automatically + inserted into the redis store when the parent model is inserted. + e.g. a Book with an author field, when inserted, will also insert + the author. The author can then be queried directly if that's something + one wishes to do. - 2. When a parent model is inserted with a nested model instance that already exists, the older nested model - instance is overwritten. This is one way of updating nested models. All parent models that contain that nested - model instance will see the change. + 2. When a parent model is inserted with a nested model instance that + already exists, the older nested model instance is overwritten. + This is one way of updating nested models. + All parent models that contain that nested model instance will see the change. """ _primary_key_field: str = 'title' title: str @@ -81,8 +88,10 @@ class Library(Model): About Nested Model Performance --- - To minimize the performance penalty for nesting models, we use REDIS EVALSHA to eagerly load the nested models - before the response is returned to the client. This ensures that only ONE network call is made every time. + To minimize the performance penalty for nesting models, + we use REDIS EVALSHA to eagerly load the nested models + before the response is returned to the client. + This ensures that only ONE network call is made every time. """ _primary_key_field: str = 'name' name: str @@ -180,15 +189,26 @@ print(all_books) # Prints [Book(title="Oliver Twist", author="Charles Dickens", published_on=date(year=1215, month=4, day=4), # in_stock=False), Book(...] +# or paginate i.e. skip some books and return only upto a given number +paginated_books = Book.select(skip=2, limit=2) +print(paginated_books) + # Get some, with all fields shown. Data returned is a list of models instances. some_books = Book.select(ids=["Oliver Twist", "Jane Eyre"]) print(some_books) +# Note: Pagination does not work when ids are provided i.e. +assert some_books == Book.select(ids=["Oliver Twist", "Jane Eyre"], skip=100, limit=10) + # Get all, with only a few fields shown. Data returned is a list of dictionaries. books_with_few_fields = Book.select(columns=["author", "in_stock"]) print(books_with_few_fields) # Prints [{"author": "'Charles Dickens", "in_stock": "True"},...] +# or paginate i.e. skip some books and return only upto a given number +paginated_books_with_few_fields = Book.select(columns=["author", "in_stock"], skip=2, limit=2) +print(paginated_books_with_few_fields) + # Get some, with only some fields shown. Data returned is a list of dictionaries. some_books_with_few_fields = Book.select(ids=["Oliver Twist", "Jane Eyre"], columns=["author", "in_stock"]) print(some_books_with_few_fields) @@ -217,13 +237,15 @@ from datetime import date from typing import Tuple, List, Optional from pydantic_redis.asyncio import RedisConfig, Model, Store -# The features are exactly the same as the synchronous version, except for the ability -# to return coroutines when `insert`, `update`, `select` or `delete` are called. +# The features are exactly the same as the synchronous version, +# except for the ability to return coroutines when `insert`, +# `update`, `select` or `delete` are called. class Author(Model): """ - An Author model, just like a pydantic model with appropriate type annotations + An Author model, just like a pydantic model with appropriate + type annotations NOTE: The `_primary_key_field` is mandatory """ _primary_key_field: str = 'name' @@ -236,18 +258,24 @@ class Book(Model): A Book model. Models can have the following field types - - The usual i.e. float, int, dict, list, date, str, dict, Optional etc as long as they are serializable by orjson + - The usual i.e. float, int, dict, list, date, str, dict, Optional etc + as long as they are serializable by orjson - Nested models e.g. `author: Author` or `author: Optional[Author]` - - List of nested models e.g. `authors: List[Author]` or `authors: Optional[List[Author]]` - - Tuples including nested models e.g. `access_log: Tuple[Author, date]` or `access_log: Optional[Tuple[Author, date]]]` + - List of nested models e.g. `authors: List[Author]` + or `authors: Optional[List[Author]]` + - Tuples including nested models e.g. `access_log: Tuple[Author, date]` + or `access_log: Optional[Tuple[Author, date]]]` - NOTE: 1. Any nested model whether plain or in a list or tuple will automatically inserted into the redis store - when the parent model is inserted. e.g. a Book with an author field, when inserted, will also insert - the author. The author can then be queried directly if that's something one wishes to do. + NOTE: 1. Any nested model whether plain or in a list or tuple will automatically + inserted into the redis store when the parent model is inserted. + e.g. a Book with an author field, when inserted, will also insert + the author. The author can then be queried directly if that's something + one wishes to do. - 2. When a parent model is inserted with a nested model instance that already exists, the older nested model - instance is overwritten. This is one way of updating nested models. All parent models that contain that nested - model instance will see the change. + 2. When a parent model is inserted with a nested model instance that + already exists, the older nested model instance is overwritten. + This is one way of updating nested models. + All parent models that contain that nested model instance will see the change. """ _primary_key_field: str = 'title' title: str @@ -266,8 +294,10 @@ class Library(Model): About Nested Model Performance --- - To minimize the performance penalty for nesting models, we use REDIS EVALSHA to eagerly load the nested models - before the response is returned to the client. This ensures that only ONE network call is made every time. + To minimize the performance penalty for nesting models, + we use REDIS EVALSHA to eagerly load the nested models + before the response is returned to the client. + This ensures that only ONE network call is made every time. """ _primary_key_field: str = 'name' name: str @@ -367,15 +397,26 @@ async def run_async(): # Prints [Book(title="Oliver Twist", author="Charles Dickens", published_on=date(year=1215, month=4, day=4), # in_stock=False), Book(...] + # or paginate i.e. skip some books and return only upto a given number + paginated_books = await Book.select(skip=2, limit=2) + print(paginated_books) + # Get some, with all fields shown. Data returned is a list of models instances. some_books = await Book.select(ids=["Oliver Twist", "Jane Eyre"]) print(some_books) + # Note: Pagination does not work when ids are provided i.e. + assert some_books == await Book.select(ids=["Oliver Twist", "Jane Eyre"], skip=100, limit=10) + # Get all, with only a few fields shown. Data returned is a list of dictionaries. books_with_few_fields = await Book.select(columns=["author", "in_stock"]) print(books_with_few_fields) # Prints [{"author": "'Charles Dickens", "in_stock": "True"},...] + # or paginate i.e. skip some books and return only upto a given number + paginated_books_with_few_fields = await Book.select(columns=["author", "in_stock"], skip=2, limit=2) + print(paginated_books_with_few_fields) + # Get some, with only some fields shown. Data returned is a list of dictionaries. some_books_with_few_fields = await Book.select(ids=["Oliver Twist", "Jane Eyre"], columns=["author", "in_stock"]) print(some_books_with_few_fields) @@ -440,30 +481,32 @@ asyncio.run(run_async()) On an average PC ~16GB RAM, i7 Core ``` -------------------------------------------------- benchmark: 20 tests ------------------------------------------------- -Name (time in us) Mean Min Max ------------------------------------------------------------------------------------------------------------------------ -benchmark_select_columns_for_one_id[redis_store-book1] 143.5316 (1.08) 117.4340 (1.0) 347.5900 (1.0) -benchmark_select_columns_for_one_id[redis_store-book3] 151.6032 (1.14) 117.6690 (1.00) 405.4620 (1.17) -benchmark_select_columns_for_one_id[redis_store-book0] 133.0856 (1.0) 117.8720 (1.00) 403.9400 (1.16) -benchmark_select_columns_for_one_id[redis_store-book2] 156.8152 (1.18) 118.7220 (1.01) 569.9800 (1.64) -benchmark_select_columns_for_some_items[redis_store] 138.0488 (1.04) 120.1550 (1.02) 350.7040 (1.01) -benchmark_delete[redis_store-Wuthering Heights] 199.9205 (1.50) 127.6990 (1.09) 1,092.2190 (3.14) -benchmark_bulk_delete[redis_store] 178.4756 (1.34) 143.7480 (1.22) 647.6660 (1.86) -benchmark_select_all_for_one_id[redis_store-book1] 245.7787 (1.85) 195.2030 (1.66) 528.9250 (1.52) -benchmark_select_all_for_one_id[redis_store-book0] 239.1152 (1.80) 199.4360 (1.70) 767.2540 (2.21) -benchmark_select_all_for_one_id[redis_store-book3] 243.8724 (1.83) 200.8060 (1.71) 535.3640 (1.54) -benchmark_select_all_for_one_id[redis_store-book2] 256.1625 (1.92) 202.4630 (1.72) 701.3000 (2.02) -benchmark_update[redis_store-Wuthering Heights-data0] 329.1363 (2.47) 266.9700 (2.27) 742.1360 (2.14) -benchmark_select_some_items[redis_store] 301.0471 (2.26) 268.9410 (2.29) 551.1060 (1.59) -benchmark_select_columns[redis_store] 313.4356 (2.36) 281.4460 (2.40) 578.7730 (1.67) -benchmark_single_insert[redis_store-book2] 348.5624 (2.62) 297.3610 (2.53) 580.8780 (1.67) -benchmark_single_insert[redis_store-book1] 342.1879 (2.57) 297.5410 (2.53) 650.5420 (1.87) -benchmark_single_insert[redis_store-book0] 366.4513 (2.75) 310.1640 (2.64) 660.5380 (1.90) -benchmark_single_insert[redis_store-book3] 377.6208 (2.84) 327.5290 (2.79) 643.4090 (1.85) -benchmark_select_default[redis_store] 486.6931 (3.66) 428.8810 (3.65) 1,181.9620 (3.40) -benchmark_bulk_insert[redis_store] 897.7862 (6.75) 848.7410 (7.23) 1,188.5160 (3.42) ------------------------------------------------------------------------------------------------------------------------ +-------------------------------------------------- benchmark: 22 tests -------------------------------------------------- +Name (time in us) Mean Min Max +------------------------------------------------------------------------------------------------------------------------- +benchmark_select_columns_for_one_id[redis_store-book2] 124.2687 (1.00) 115.4530 (1.0) 331.8030 (1.26) +benchmark_select_columns_for_one_id[redis_store-book0] 123.7213 (1.0) 115.6680 (1.00) 305.7170 (1.16) +benchmark_select_columns_for_one_id[redis_store-book3] 124.4495 (1.01) 115.9580 (1.00) 263.4370 (1.0) +benchmark_select_columns_for_one_id[redis_store-book1] 124.8431 (1.01) 117.4770 (1.02) 310.3140 (1.18) +benchmark_select_columns_for_some_items[redis_store] 128.0657 (1.04) 118.6380 (1.03) 330.2680 (1.25) +benchmark_delete[redis_store-Wuthering Heights] 131.8713 (1.07) 125.9920 (1.09) 328.9660 (1.25) +benchmark_bulk_delete[redis_store] 148.6963 (1.20) 142.3190 (1.23) 347.4750 (1.32) +benchmark_select_all_for_one_id[redis_store-book3] 211.6941 (1.71) 195.6520 (1.69) 422.8840 (1.61) +benchmark_select_all_for_one_id[redis_store-book2] 212.3612 (1.72) 195.9020 (1.70) 447.4910 (1.70) +benchmark_select_all_for_one_id[redis_store-book1] 212.9524 (1.72) 197.7530 (1.71) 423.3030 (1.61) +benchmark_select_all_for_one_id[redis_store-book0] 214.9924 (1.74) 198.8280 (1.72) 402.6310 (1.53) +benchmark_select_columns_paginated[redis_store] 227.9248 (1.84) 211.0610 (1.83) 425.8390 (1.62) +benchmark_select_some_items[redis_store] 297.5700 (2.41) 271.1510 (2.35) 572.1220 (2.17) +benchmark_select_default_paginated[redis_store] 301.7495 (2.44) 282.6500 (2.45) 490.3450 (1.86) +benchmark_select_columns[redis_store] 316.2119 (2.56) 290.6110 (2.52) 578.0310 (2.19) +benchmark_update[redis_store-Wuthering Heights-data0] 346.5816 (2.80) 304.5420 (2.64) 618.0250 (2.35) +benchmark_single_insert[redis_store-book2] 378.0613 (3.06) 337.8070 (2.93) 616.4930 (2.34) +benchmark_single_insert[redis_store-book0] 396.6513 (3.21) 347.1000 (3.01) 696.1350 (2.64) +benchmark_single_insert[redis_store-book3] 395.9082 (3.20) 361.0980 (3.13) 623.8630 (2.37) +benchmark_single_insert[redis_store-book1] 401.1377 (3.24) 363.5890 (3.15) 610.4400 (2.32) +benchmark_select_default[redis_store] 498.4673 (4.03) 428.1350 (3.71) 769.7640 (2.92) +benchmark_bulk_insert[redis_store] 1,025.0436 (8.29) 962.2230 (8.33) 1,200.3840 (4.56) +------------------------------------------------------------------------------------------------------------------------- ``` ## Contributions diff --git a/pydantic_redis/__init__.py b/pydantic_redis/__init__.py index da72f1de..9603a70b 100644 --- a/pydantic_redis/__init__.py +++ b/pydantic_redis/__init__.py @@ -5,4 +5,4 @@ __all__ = [Store, RedisConfig, Model, asyncio] -__version__ = "0.3.0" +__version__ = "0.4.0" diff --git a/pydantic_redis/asyncio/model.py b/pydantic_redis/asyncio/model.py index 76f10033..ea04ddea 100644 --- a/pydantic_redis/asyncio/model.py +++ b/pydantic_redis/asyncio/model.py @@ -15,6 +15,7 @@ ) from .store import Store +from ..shared.model.delete_utils import delete_on_pipeline class Model(AbstractModel): @@ -92,21 +93,7 @@ async def delete(cls, ids: Union[Any, List[Any]]): 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) + delete_on_pipeline(model=cls, pipeline=pipeline, ids=ids) return await pipeline.execute() @classmethod @@ -114,19 +101,33 @@ async def select( cls, columns: Optional[List[str]] = None, ids: Optional[List[Any]] = None, + skip: int = 0, + limit: Optional[int] = None, **kwargs, ): """ Selects given rows or sets of rows in the table + + + However, if `limit` is set, the number of items + returned will be less or equal to `limit`. + `skip` defaults to 0. It is the number of items to skip. + `skip` is only relevant when limit is specified. + + `skip` and `limit` are irrelevant when `ids` are provided. """ if columns is None and ids is None: - response = await select_all_fields_all_ids(model=cls) + response = await select_all_fields_all_ids( + model=cls, skip=skip, limit=limit + ) 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) + response = await select_some_fields_all_ids( + model=cls, fields=columns, skip=skip, limit=limit + ) elif isinstance(columns, list) and isinstance(ids, list): response = await select_some_fields_some_ids( diff --git a/pydantic_redis/shared/lua_scripts.py b/pydantic_redis/shared/lua_scripts.py index d1896bfe..8221e501 100644 --- a/pydantic_redis/shared/lua_scripts.py +++ b/pydantic_redis/shared/lua_scripts.py @@ -59,6 +59,67 @@ return filtered """ +PAGINATED_SELECT_ALL_FIELDS_FOR_ALL_IDS_SCRIPT = """ +local s_find = string.find +local s_gmatch = string.gmatch +local ipairs = ipairs +local table_insert = table.insert +local next = next +local redis_call = redis.call + +local function startswith(s, prefix) + return s_find(s, prefix, 1, true) == 1 +end + +local function trim_dunder(s) + return s:match '^_*(.-)$' +end + +local function get_obj(id) + local value = redis_call('HGETALL', id) + + for i, k in ipairs(value) do + if not (i % 2 == 0) then + if startswith(k, '___') or startswith(k, '____') then + local nested = {} + + for v in s_gmatch(value[i + 1], '([^%[^,^%]^\"]+)') do + table_insert(nested, get_obj(v)) + end + + value[i + 1] = nested + value[i] = trim_dunder(k) + elseif startswith(k, '__') then + value[i + 1] = get_obj(value[i + 1]) + value[i] = trim_dunder(k) + end + end + end + + if next(value) == nil then + return id + end + + return value +end + +local table_index_key = ARGV[1] +local start = ARGV[2] +local stop = ARGV[3] + start - 1 +local result = {} + +local ids = redis_call('ZRANGE', table_index_key, start, stop) + +for _, key in ipairs(ids) do + local value = get_obj(key) + if type(value) == 'table' then + table_insert(result, value) + end +end + +return result +""" + SELECT_ALL_FIELDS_FOR_SOME_IDS_SCRIPT = """ local s_find = string.find local s_gmatch = string.gmatch @@ -190,6 +251,82 @@ return filtered """ +PAGINATED_SELECT_SOME_FIELDS_FOR_ALL_IDS_SCRIPT = """ +local s_find = string.find +local s_gmatch = string.gmatch +local ipairs = ipairs +local table_insert = table.insert +local next = next +local redis_call = redis.call +local table_unpack = table.unpack or unpack + +local function startswith(s, prefix) + return s_find(s, prefix, 1, true) == 1 +end + +local function trim_dunder(s) + return s:match '^_*(.-)$' +end + +local function get_obj(id) + local value = redis_call('HGETALL', id) + + for i, k in ipairs(value) do + if not (i % 2 == 0) then + if startswith(k, '___') or startswith(k, '____') then + local nested = {} + + for v in s_gmatch(value[i + 1], '([^%[^,^%]^\"]+)') do + table_insert(nested, get_obj(v)) + end + + value[i + 1] = nested + value[i] = trim_dunder(k) + elseif startswith(k, '__') then + value[i + 1] = get_obj(value[i + 1]) + value[i] = trim_dunder(k) + end + end + end + + if next(value) == nil then + return id + end + + return value +end + +local result = {} +local columns = {} +local table_index_key = ARGV[1] +local start = ARGV[2] +local stop = ARGV[3] + start - 1 + +for i, k in ipairs(ARGV) do + if i > 3 then + table_insert(columns, k) + end +end + +local ids = redis_call('ZRANGE', table_index_key, start, stop) + +for _, key in ipairs(ids) do + local data = redis_call('HMGET', key, table_unpack(columns)) + local parsed_data = {} + + for i, v in ipairs(data) do + if v then + table_insert(parsed_data, trim_dunder(columns[i])) + table_insert(parsed_data, get_obj(v)) + end + end + + table_insert(result, parsed_data) +end + +return result +""" + SELECT_SOME_FIELDS_FOR_SOME_IDS_SCRIPT = """ local s_find = string.find local s_gmatch = string.gmatch diff --git a/pydantic_redis/shared/model/delete_utils.py b/pydantic_redis/shared/model/delete_utils.py new file mode 100644 index 00000000..19709111 --- /dev/null +++ b/pydantic_redis/shared/model/delete_utils.py @@ -0,0 +1,32 @@ +"""Module containing common functionality for deleting records""" +from typing import Type, Union, List + +from redis.client import Pipeline +from redis.asyncio.client import Pipeline as AioPipeline + +from pydantic_redis.shared.model import AbstractModel +from pydantic_redis.shared.model.prop_utils import get_primary_key, get_table_index_key + + +def delete_on_pipeline( + model: Type[AbstractModel], pipeline: Union[Pipeline, AioPipeline], ids: List[str] +): + """ + Pipelines the deletion of the given ids, so that when pipeline.execute + is called later, deletion occurs + """ + primary_keys = [] + + if isinstance(ids, list): + primary_keys = ids + elif ids is not None: + primary_keys = [ids] + + names = [ + get_primary_key(model=model, primary_key_value=primary_key_value) + for primary_key_value in primary_keys + ] + pipeline.delete(*names) + # remove the primary keys from the indexz + table_index_key = get_table_index_key(model=model) + pipeline.zrem(table_index_key, *names) diff --git a/pydantic_redis/shared/model/insert_utils.py b/pydantic_redis/shared/model/insert_utils.py index 7c70b8f9..2492d6db 100644 --- a/pydantic_redis/shared/model/insert_utils.py +++ b/pydantic_redis/shared/model/insert_utils.py @@ -1,4 +1,5 @@ """Module containing the mixin for insert functionality in model""" +from datetime import datetime from typing import Union, Optional, Any, Dict, Tuple, List, Type from redis.asyncio.client import Pipeline as AioPipeline @@ -37,9 +38,10 @@ def insert_on_pipeline( if life_span is not None: pipeline.expire(name=name, time=life_span) - # save the primary key in an index + # save the primary key in an index: a sorted set, whose score is current timestamp table_index_key = get_table_index_key(model) - pipeline.sadd(table_index_key, name) + timestamp = datetime.utcnow().timestamp() + pipeline.zadd(table_index_key, {name: timestamp}) if life_span is not None: pipeline.expire(table_index_key, time=life_span) @@ -113,8 +115,8 @@ def _serialize_nested_model_tuple_field( record=item, life_span=life_span, ) - # FIXME: The reference to AbstractModel here. Is it right? - if issubclass(field_type, AbstractModel) else item + if issubclass(field_type, AbstractModel) + else item for field_type, item in zip(field_types, value) ] key = f"{NESTED_MODEL_TUPLE_FIELD_PREFIX}{key}" diff --git a/pydantic_redis/shared/model/select_utils.py b/pydantic_redis/shared/model/select_utils.py index ac152ec2..5911f505 100644 --- a/pydantic_redis/shared/model/select_utils.py +++ b/pydantic_redis/shared/model/select_utils.py @@ -1,5 +1,5 @@ """Module containing the mixin functionality for selecting""" -from typing import List, Any, Type, Union, Awaitable +from typing import List, Any, Type, Union, Awaitable, Optional from pydantic_redis.shared.model.prop_utils import ( NESTED_MODEL_PREFIX, @@ -7,6 +7,7 @@ NESTED_MODEL_TUPLE_FIELD_PREFIX, get_table_keys_regex, get_table_prefix, + get_table_index_key, ) @@ -39,12 +40,24 @@ def get_select_fields(model: Type[AbstractModel], columns: List[str]) -> List[st def select_all_fields_all_ids( model: Type[AbstractModel], + skip: int = 0, + limit: Optional[int] = None, ) -> Union[List[List[Any]], Awaitable[List[List[Any]]]]: - """Selects all items in the database, returning all their fields""" - table_keys_regex = get_table_keys_regex(model=model) - args = [table_keys_regex] - store = model.get_store() - return store.select_all_fields_for_all_ids_script(args=args) + """ + Selects all items in the database, returning all their fields + + However, if `limit` is set, the number of items + returned will be less or equal to `limit`. + `skip` defaults to 0. It is the number of items to skip. + `skip` is only relevant when limit is specified. + """ + if isinstance(limit, int): + return _select_all_ids_all_fields_paginated(model=model, limit=limit, skip=skip) + else: + table_keys_regex = get_table_keys_regex(model=model) + args = [table_keys_regex] + store = model.get_store() + return store.select_all_fields_for_all_ids_script(args=args) def select_all_fields_some_ids( @@ -58,14 +71,30 @@ def select_all_fields_some_ids( def select_some_fields_all_ids( - model: Type[AbstractModel], fields: List[str] + model: Type[AbstractModel], + fields: List[str], + skip: int = 0, + limit: Optional[int] = None, ) -> Union[List[List[Any]], Awaitable[List[List[Any]]]]: - """Selects all items in the database, returning only the specified fields""" - table_keys_regex = get_table_keys_regex(model=model) + """ + Selects all items in the database, returning only the specified fields. + + However, if `limit` is set, the number of items + returned will be less or equal to `limit`. + `skip` defaults to 0. It is the number of items to skip. + `skip` is only relevant when limit is specified. + """ columns = get_select_fields(model=model, columns=fields) - args = [table_keys_regex, *columns] - store = model.get_store() - return store.select_some_fields_for_all_ids_script(args=args) + + if isinstance(limit, int): + return _select_some_fields_all_ids_paginated( + model=model, columns=columns, limit=limit, skip=skip + ) + else: + table_keys_regex = get_table_keys_regex(model=model) + args = [table_keys_regex, *columns] + store = model.get_store() + return store.select_some_fields_for_all_ids_script(args=args) def select_some_fields_some_ids( @@ -100,3 +129,27 @@ def parse_select_response( ] return [model.deserialize_partially(record) for record in response if record != []] + + +def _select_all_ids_all_fields_paginated( + model: Type[AbstractModel], limit: int, skip: Optional[int] +): + """Selects all fields for at most `limit` number of items after skipping `skip` items""" + if skip is None: + skip = 0 + table_index_key = get_table_index_key(model) + args = [table_index_key, skip, limit] + store = model.get_store() + return store.paginated_select_all_fields_for_all_ids_script(args=args) + + +def _select_some_fields_all_ids_paginated( + model: Type[AbstractModel], columns: List[str], limit: int, skip: int +): + """Selects some fields for at most `limit` number of items after skipping `skip` items""" + if skip is None: + skip = 0 + table_index_key = get_table_index_key(model) + args = [table_index_key, skip, limit, *columns] + store = model.get_store() + return store.paginated_select_some_fields_for_all_ids_script(args=args) diff --git a/pydantic_redis/shared/store.py b/pydantic_redis/shared/store.py index 5d193563..65e10c68 100644 --- a/pydantic_redis/shared/store.py +++ b/pydantic_redis/shared/store.py @@ -12,6 +12,8 @@ SELECT_ALL_FIELDS_FOR_SOME_IDS_SCRIPT, SELECT_SOME_FIELDS_FOR_ALL_IDS_SCRIPT, SELECT_SOME_FIELDS_FOR_SOME_IDS_SCRIPT, + PAGINATED_SELECT_ALL_FIELDS_FOR_ALL_IDS_SCRIPT, + PAGINATED_SELECT_SOME_FIELDS_FOR_ALL_IDS_SCRIPT, ) from typing import TYPE_CHECKING @@ -30,8 +32,14 @@ class AbstractStore(BaseModel): redis_store: Optional[Union[Redis, AioRedis]] = None life_span_in_seconds: Optional[int] = None select_all_fields_for_all_ids_script: Optional[Union[AsyncScript, Script]] = None + paginated_select_all_fields_for_all_ids_script: Optional[ + Union[AsyncScript, Script] + ] = None select_all_fields_for_some_ids_script: Optional[Union[AsyncScript, Script]] = None select_some_fields_for_all_ids_script: Optional[Union[AsyncScript, Script]] = None + paginated_select_some_fields_for_all_ids_script: Optional[ + Union[AsyncScript, Script] + ] = None select_some_fields_for_some_ids_script: Optional[Union[AsyncScript, Script]] = None models: Dict[str, Type["AbstractModel"]] = {} @@ -67,12 +75,22 @@ def _register_lua_scripts(self): self.select_all_fields_for_all_ids_script = self.redis_store.register_script( SELECT_ALL_FIELDS_FOR_ALL_IDS_SCRIPT ) + self.paginated_select_all_fields_for_all_ids_script = ( + self.redis_store.register_script( + PAGINATED_SELECT_ALL_FIELDS_FOR_ALL_IDS_SCRIPT + ) + ) self.select_all_fields_for_some_ids_script = self.redis_store.register_script( SELECT_ALL_FIELDS_FOR_SOME_IDS_SCRIPT ) self.select_some_fields_for_all_ids_script = self.redis_store.register_script( SELECT_SOME_FIELDS_FOR_ALL_IDS_SCRIPT ) + self.paginated_select_some_fields_for_all_ids_script = ( + self.redis_store.register_script( + PAGINATED_SELECT_SOME_FIELDS_FOR_ALL_IDS_SCRIPT + ) + ) self.select_some_fields_for_some_ids_script = self.redis_store.register_script( SELECT_SOME_FIELDS_FOR_SOME_IDS_SCRIPT ) diff --git a/pydantic_redis/syncio/model.py b/pydantic_redis/syncio/model.py index 0999e939..25cbe545 100644 --- a/pydantic_redis/syncio/model.py +++ b/pydantic_redis/syncio/model.py @@ -13,6 +13,7 @@ ) from .store import Store +from ..shared.model.delete_utils import delete_on_pipeline class Model(AbstractModel): @@ -89,21 +90,7 @@ def delete(cls, ids: Union[Any, List[Any]]): """ store = cls.get_store() 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) + delete_on_pipeline(model=cls, pipeline=pipeline, ids=ids) return pipeline.execute() @classmethod @@ -111,19 +98,30 @@ def select( cls, columns: Optional[List[str]] = None, ids: Optional[List[Any]] = None, + skip: int = 0, + limit: Optional[int] = None, **kwargs, ): """ Selects given rows or sets of rows in the table + + However, if `limit` is set, the number of items + returned will be less or equal to `limit`. + `skip` defaults to 0. It is the number of items to skip. + `skip` is only relevant when limit is specified. + + `skip` and `limit` are irrelevant when `ids` are provided. """ if columns is None and ids is None: - response = select_all_fields_all_ids(model=cls) + response = select_all_fields_all_ids(model=cls, skip=skip, limit=limit) elif columns is None and isinstance(ids, list): response = select_all_fields_some_ids(model=cls, ids=ids) elif isinstance(columns, list) and ids is None: - response = select_some_fields_all_ids(model=cls, fields=columns) + response = select_some_fields_all_ids( + model=cls, fields=columns, skip=skip, limit=limit + ) elif isinstance(columns, list) and isinstance(ids, list): response = select_some_fields_some_ids(model=cls, fields=columns, ids=ids) diff --git a/setup.py b/setup.py index 7ba54f91..26d00ecf 100644 --- a/setup.py +++ b/setup.py @@ -11,7 +11,7 @@ # This call to setup() does all the work setup( name="pydantic-redis", - version="0.3.0", + version="0.4.0", description="This package provides a simple ORM for redis using pydantic-like models.", long_description=README, long_description_content_type="text/markdown", diff --git a/test/test_async_pydantic_redis.py b/test/test_async_pydantic_redis.py index 3d177ff4..4208075f 100644 --- a/test/test_async_pydantic_redis.py +++ b/test/test_async_pydantic_redis.py @@ -1,5 +1,5 @@ """Tests for the redis orm""" - +from collections import namedtuple from typing import Dict, Any import pytest @@ -229,6 +229,28 @@ async def test_select_default(store): assert sorted_async_books == sorted_response +@pytest.mark.asyncio +@pytest.mark.parametrize("store", async_redis_store_fixture) +async def test_select_default_paginated(store): + """ + Selecting without arguments returns the book models after + skipping `skip` number of models and returning upto `limit` number of items + """ + await AsyncBook.insert(async_books) + Record = namedtuple("Record", ["skip", "limit", "expected"]) + test_data = [ + Record(0, 2, sorted(async_books[:2], key=lambda x: x.title)), + Record(None, 2, sorted(async_books[:2], key=lambda x: x.title)), + Record(2, 2, sorted(async_books[2:4], key=lambda x: x.title)), + Record(3, 2, sorted(async_books[3:5], key=lambda x: x.title)), + Record(0, 3, sorted(async_books[:3], key=lambda x: x.title)), + ] + for record in test_data: + response = await AsyncBook.select(skip=record.skip, limit=record.limit) + sorted_response = sorted(response, key=lambda x: x.title) + assert record.expected == sorted_response + + @pytest.mark.asyncio @pytest.mark.parametrize("store", async_redis_store_fixture) async def test_select_some_columns(store): @@ -238,7 +260,7 @@ async def test_select_some_columns(store): await AsyncBook.insert(async_books) async_books_dict = {book.title: book for book in async_books} columns = ["title", "author", "in_stock"] - response = await AsyncBook.select(columns=["title", "author", "in_stock"]) + response = await AsyncBook.select(columns=columns) response_dict = {book["title"]: book for book in response} for title, book in async_books_dict.items(): @@ -253,6 +275,44 @@ async def test_select_some_columns(store): assert f"{book_in_response[column]}" == f"{getattr(book, column)}" +@pytest.mark.asyncio +@pytest.mark.parametrize("store", async_redis_store_fixture) +async def test_select_some_columns_paginated(store): + """ + Selecting some columns returns a list of dictionaries of all books models with only those columns + skipping `skip` number of models and returning upto `limit` number of items + """ + await AsyncBook.insert(async_books) + columns = ["title", "author", "in_stock"] + + Record = namedtuple("Record", ["skip", "limit", "expected"]) + test_data = [ + Record(0, 2, sorted(async_books[:2], key=lambda x: x.title)), + Record(None, 2, sorted(async_books[:2], key=lambda x: x.title)), + Record(2, 2, sorted(async_books[2:4], key=lambda x: x.title)), + Record(3, 2, sorted(async_books[3:5], key=lambda x: x.title)), + Record(0, 3, sorted(async_books[:3], key=lambda x: x.title)), + ] + for record in test_data: + response = await AsyncBook.select( + columns=columns, skip=record.skip, limit=record.limit + ) + response_dict = {book["title"]: book for book in response} + books_dict = {book.title: book for book in record.expected} + assert len(record.expected) == len(response_dict) + + for title, book in books_dict.items(): + book_in_response = response_dict[title] + assert isinstance(book_in_response, dict) + assert sorted(book_in_response.keys()) == sorted(columns) + + for column in columns: + if column == "author": + assert book_in_response[column] == getattr(book, column) + else: + assert f"{book_in_response[column]}" == f"{getattr(book, column)}" + + @pytest.mark.asyncio @pytest.mark.parametrize("store", async_redis_store_fixture) async def test_select_some_ids(store): diff --git a/test/test_benchmarks.py b/test/test_benchmarks.py index 4428268d..3a24d7f5 100644 --- a/test/test_benchmarks.py +++ b/test/test_benchmarks.py @@ -30,6 +30,13 @@ def test_benchmark_select_default(benchmark, store): benchmark(Book.select) +@pytest.mark.parametrize("store", redis_store_fixture) +def test_benchmark_select_default_paginated(benchmark, store): + """Benchmarks the select default operation when paginated""" + Book.insert(books) + benchmark(Book.select, skip=2, limit=2) + + @pytest.mark.parametrize("store", redis_store_fixture) def test_benchmark_select_columns(benchmark, store): """Benchmarks the select columns operation""" @@ -37,6 +44,13 @@ def test_benchmark_select_columns(benchmark, store): benchmark(Book.select, columns=["title", "author", "in_stock"]) +@pytest.mark.parametrize("store", redis_store_fixture) +def test_benchmark_select_columns_paginated(benchmark, store): + """Benchmarks the select columns operation, when paginated""" + Book.insert(books) + benchmark(Book.select, columns=["title", "author", "in_stock"], skip=2, limit=2) + + @pytest.mark.parametrize("store", redis_store_fixture) def test_benchmark_select_some_items(benchmark, store): """Benchmarks the select some items operation""" diff --git a/test/test_pydantic_redis.py b/test/test_pydantic_redis.py index d8d52eaf..99f7ca0a 100644 --- a/test/test_pydantic_redis.py +++ b/test/test_pydantic_redis.py @@ -1,5 +1,5 @@ """Tests for the redis orm""" - +from collections import namedtuple from typing import Dict, Any import pytest @@ -207,6 +207,27 @@ def test_select_default(store): assert sorted_books == sorted_response +@pytest.mark.parametrize("store", redis_store_fixture) +def test_select_default_paginated(store): + """ + Selecting without arguments returns the book models after + skipping `skip` number of models and returning upto `limit` number of items + """ + Book.insert(books) + Record = namedtuple("Record", ["skip", "limit", "expected"]) + test_data = [ + Record(0, 2, sorted(books[:2], key=lambda x: x.title)), + Record(None, 2, sorted(books[:2], key=lambda x: x.title)), + Record(2, 2, sorted(books[2:4], key=lambda x: x.title)), + Record(3, 2, sorted(books[3:5], key=lambda x: x.title)), + Record(0, 3, sorted(books[:3], key=lambda x: x.title)), + ] + for record in test_data: + response = Book.select(skip=record.skip, limit=record.limit) + sorted_response = sorted(response, key=lambda x: x.title) + assert record.expected == sorted_response + + @pytest.mark.parametrize("store", redis_store_fixture) def test_select_some_columns(store): """ @@ -215,7 +236,7 @@ def test_select_some_columns(store): Book.insert(books) books_dict = {book.title: book for book in books} columns = ["title", "author", "in_stock"] - response = Book.select(columns=["title", "author", "in_stock"]) + response = Book.select(columns=columns) response_dict = {book["title"]: book for book in response} for title, book in books_dict.items(): @@ -230,6 +251,41 @@ def test_select_some_columns(store): assert f"{book_in_response[column]}" == f"{getattr(book, column)}" +@pytest.mark.parametrize("store", redis_store_fixture) +def test_select_some_columns_paginated(store): + """ + Selecting some columns returns a list of dictionaries of all books models with only those columns + skipping `skip` number of models and returning upto `limit` number of items + """ + Book.insert(books) + columns = ["title", "author", "in_stock"] + + Record = namedtuple("Record", ["skip", "limit", "expected"]) + test_data = [ + Record(0, 2, sorted(books[:2], key=lambda x: x.title)), + Record(None, 2, sorted(books[:2], key=lambda x: x.title)), + Record(2, 2, sorted(books[2:4], key=lambda x: x.title)), + Record(3, 2, sorted(books[3:5], key=lambda x: x.title)), + Record(0, 3, sorted(books[:3], key=lambda x: x.title)), + ] + for record in test_data: + response = Book.select(columns=columns, skip=record.skip, limit=record.limit) + response_dict = {book["title"]: book for book in response} + books_dict = {book.title: book for book in record.expected} + assert len(record.expected) == len(response_dict) + + for title, book in books_dict.items(): + book_in_response = response_dict[title] + assert isinstance(book_in_response, dict) + assert sorted(book_in_response.keys()) == sorted(columns) + + for column in columns: + if column == "author": + assert book_in_response[column] == getattr(book, column) + else: + assert f"{book_in_response[column]}" == f"{getattr(book, column)}" + + @pytest.mark.parametrize("store", redis_store_fixture) def test_select_some_ids(store): """