Iterators QuerySet Method (#688)

* feat: add iterator function with sample docstring

* feat: implement the iterator queryset method

* feat: completed the docstring of iterator method

* test: write test function to check iterator result

* refactor: use iterate method instead fetch_all

* fix: debuging syntax error in kwargs of iterator

* feat: write a base sample doc for iterator method

* refactor: add ouput comment iterator docs example

* refactor: check change pk yield ormar model

* test: write new test to coverage iterator queryset

* fix: debuging new user model class 3 number

* fix: iterate on user3 model class

* fix: debug id field in user3 model by main user

* fix: remove prefetch_related for iterator method

* fix: debug mypy test for type annotation

* fix: added type annotation for rows variable

* simplify row checks as processing is expensive, raise exception on prefetch_related with iterator

* fix coverage

* fix mypy, bump mypy in pre-commit to newest version

* refactor: update document and test uuid pk type

* feat: write docs of iterate in quesrysetproxy

* feat: write iterate method querysetproxy tests

* fix: debuging new test written uuid pk

* refactor: seperate iterate test modules

* refactor: change description and handle empty set

* feat: added iterate method in readme files

* fix: set pragma: no cover for raised test

Co-authored-by: collerek <collerek@gmail.com>
This commit is contained in:
Sepehr Bazyar
2022-07-04 15:11:28 +04:30
committed by GitHub
parent 20cddde2f4
commit 06c3bdb5eb
8 changed files with 403 additions and 3 deletions

View File

@ -1,6 +1,6 @@
repos:
- repo: https://github.com/psf/black
rev: 22.1.0
rev: 22.3.0
hooks:
- id: black
exclude: ^(docs_src/|examples/)
@ -11,7 +11,7 @@ repos:
exclude: ^(docs_src/|examples/|tests/)
args: [ '--max-line-length=88' ]
- repo: https://github.com/pre-commit/mirrors-mypy
rev: v0.931
rev: v0.961
hooks:
- id: mypy
exclude: ^(docs_src/|examples/)

View File

@ -595,6 +595,7 @@ metadata.drop_all(engine)
* `bulk_update(objects: List[Model], columns: List[str] = None) -> None`
* `delete(*args, each: bool = False, **kwargs) -> int`
* `all(*args, **kwargs) -> List[Optional[Model]]`
* `iterate(*args, **kwargs) -> AsyncGenerator[Model]`
* `filter(*args, **kwargs) -> QuerySet`
* `exclude(*args, **kwargs) -> QuerySet`
* `select_related(related: Union[List, str]) -> QuerySet`

View File

@ -604,6 +604,7 @@ metadata.drop_all(engine)
* `bulk_update(objects: List[Model], columns: List[str] = None) -> None`
* `delete(*args, each: bool = False, **kwargs) -> int`
* `all(*args, **kwargs) -> List[Optional[Model]]`
* `iterate(*args, **kwargs) -> AsyncGenerator[Model]`
* `filter(*args, **kwargs) -> QuerySet`
* `exclude(*args, **kwargs) -> QuerySet`
* `select_related(related: Union[List, str]) -> QuerySet`

View File

@ -173,6 +173,45 @@ tracks = await Track.objects.all()
```
## iterate
`iterate(*args, **kwargs) -> AsyncGenerator["Model"]`
Return async iterable generator for all rows from a database for given model.
Passing args and/or kwargs is a shortcut and equals to calling `filter(*args, **kwargs).iterate()`.
If there are no rows meeting the criteria an empty async generator is returned.
```python
class Album(ormar.Model):
class Meta:
tablename = "album"
metadata = metadata
database = database
id: int = ormar.Integer(primary_key=True)
name: str = ormar.String(max_length=100)
```
```python
await Album.objects.create(name='The Cat')
await Album.objects.create(name='The Dog')
# will asynchronously iterate all Album models yielding one main model at a time from the generator
async for album in Album.objects.iterate():
print(album.name)
# The Cat
# The Dog
```
!!!warning
Use of `iterate()` causes previous `prefetch_related()` calls to be ignored;
since these two optimizations do not make sense together.
If `iterate()` & `prefetch_related()` are used together the `QueryDefinitionError` exception is raised.
## Model methods
Each model instance have a set of methods to `save`, `update` or `load` itself.

View File

@ -84,6 +84,23 @@ assert news_posts[0].author == guido
!!!tip
Read more in queries documentation [all][all]
### iterate
`iterate(**kwargs) -> AsyncGenerator["Model"]`
To iterate on related models use `iterate()` method.
Note that you can filter the queryset, select related, exclude fields etc. like in normal query.
```python
# iterate on categories of this post with an async generator
async for category in post.categories.iterate():
print(category.name)
```
!!!tip
Read more in queries documentation [iterate][iterate]
## Insert/ update data into database
### create
@ -294,6 +311,7 @@ Returns a bool value to confirm if there are rows matching the given criteria (a
[queries]: ../queries/index.md
[get]: ../queries/read.md#get
[all]: ../queries/read.md#all
[iterate]: ../queries/read.md#iterate
[create]: ../queries/create.md#create
[get_or_create]: ../queries/read.md#get_or_create
[update_or_create]: ../queries/update.md#update_or_create

View File

@ -12,6 +12,7 @@ from typing import (
TypeVar,
Union,
cast,
AsyncGenerator,
)
import databases
@ -1055,6 +1056,55 @@ class QuerySet(Generic[T]):
return result_rows
async def iterate( # noqa: A003
self,
*args: Any,
**kwargs: Any,
) -> AsyncGenerator["T", None]:
"""
Return async iterable generator for all rows from a database for given model.
Passing args and/or kwargs is a shortcut and equals to calling
`filter(*args, **kwargs).iterate()`.
If there are no rows meeting the criteria an empty async generator is returned.
:param kwargs: fields names and proper value types
:type kwargs: Any
:return: asynchronous iterable generator of returned models
:rtype: AsyncGenerator[Model]
"""
if self._prefetch_related:
raise QueryDefinitionError(
"Prefetch related queries are not supported in iterators"
)
if kwargs or args:
async for result in self.filter(*args, **kwargs).iterate():
yield result
return
expr = self.build_select_expression()
rows: list = []
last_primary_key = None
pk_alias = self.model.get_column_alias(self.model_meta.pkname)
async for row in self.database.iterate(query=expr):
current_primary_key = row[pk_alias]
if last_primary_key == current_primary_key or last_primary_key is None:
last_primary_key = current_primary_key
rows.append(row)
continue
yield self._process_query_result_rows(rows)[0]
last_primary_key = current_primary_key
rows = [row]
if rows:
yield self._process_query_result_rows(rows)[0]
async def create(self, **kwargs: Any) -> "T":
"""
Creates the model instance, saves it in a database and returns the updates model

View File

@ -14,6 +14,7 @@ from typing import ( # noqa: I100, I201
TypeVar,
Union,
cast,
AsyncGenerator,
)
import ormar # noqa: I100, I202
@ -431,6 +432,28 @@ class QuerysetProxy(Generic[T]):
self._register_related(all_items)
return all_items
async def iterate( # noqa: A003
self,
*args: Any,
**kwargs: Any,
) -> AsyncGenerator["T", None]:
"""
Return async iterable generator for all rows from a database for given model.
Passing args and/or kwargs is a shortcut and equals to calling
`filter(*args, **kwargs).iterate()`.
If there are no rows meeting the criteria an empty async generator is returned.
:param kwargs: fields names and proper value types
:type kwargs: Any
:return: asynchronous iterable generator of returned models
:rtype: AsyncGenerator[Model]
"""
async for item in self.queryset.iterate(*args, **kwargs):
yield item
async def create(self, **kwargs: Any) -> "T":
"""
Creates the model instance, saves it in a database and returns the updates model

View File

@ -0,0 +1,268 @@
import uuid
import databases
import pytest
import sqlalchemy
import ormar
from ormar.exceptions import QueryDefinitionError
from tests.settings import DATABASE_URL
database = databases.Database(DATABASE_URL, force_rollback=True)
metadata = sqlalchemy.MetaData()
class User(ormar.Model):
class Meta:
tablename = "users3"
metadata = metadata
database = database
id: int = ormar.Integer(primary_key=True)
name: str = ormar.String(max_length=100, default="")
class User2(ormar.Model):
class Meta:
tablename = "users4"
metadata = metadata
database = database
id: uuid.UUID = ormar.UUID(
uuid_format="string", primary_key=True, default=uuid.uuid4
)
name: str = ormar.String(max_length=100, default="")
class Task(ormar.Model):
class Meta:
tablename = "tasks"
metadata = metadata
database = database
id: int = ormar.Integer(primary_key=True)
name: str = ormar.String(max_length=100, default="")
user: User = ormar.ForeignKey(to=User)
class Task2(ormar.Model):
class Meta:
tablename = "tasks2"
metadata = metadata
database = database
id: uuid.UUID = ormar.UUID(
uuid_format="string", primary_key=True, default=uuid.uuid4
)
name: str = ormar.String(max_length=100, default="")
user: User2 = ormar.ForeignKey(to=User2)
@pytest.fixture(autouse=True, scope="module")
def create_test_database():
engine = sqlalchemy.create_engine(DATABASE_URL)
metadata.drop_all(engine)
metadata.create_all(engine)
yield
metadata.drop_all(engine)
@pytest.mark.asyncio
async def test_empty_result():
async with database:
async with database.transaction(force_rollback=True):
async for user in User.objects.iterate():
pass # pragma: no cover
@pytest.mark.asyncio
async def test_model_iterator():
async with database:
async with database.transaction(force_rollback=True):
tom = await User.objects.create(name="Tom")
jane = await User.objects.create(name="Jane")
lucy = await User.objects.create(name="Lucy")
async for user in User.objects.iterate():
assert user in (tom, jane, lucy)
@pytest.mark.asyncio
async def test_model_iterator_filter():
async with database:
async with database.transaction(force_rollback=True):
tom = await User.objects.create(name="Tom")
await User.objects.create(name="Jane")
await User.objects.create(name="Lucy")
async for user in User.objects.iterate(name="Tom"):
assert user.name == tom.name
@pytest.mark.asyncio
async def test_model_iterator_relations():
async with database:
async with database.transaction(force_rollback=True):
tom = await User.objects.create(name="Tom")
jane = await User.objects.create(name="Jane")
lucy = await User.objects.create(name="Lucy")
for user in tom, jane, lucy:
await Task.objects.create(name="task1", user=user)
await Task.objects.create(name="task2", user=user)
results = []
async for user in User.objects.select_related(User.tasks).iterate():
assert len(user.tasks) == 2
results.append(user)
assert len(results) == 3
@pytest.mark.asyncio
async def test_model_iterator_relations_queryset_proxy():
async with database:
async with database.transaction(force_rollback=True):
tom = await User.objects.create(name="Tom")
jane = await User.objects.create(name="Jane")
for user in tom, jane:
await Task.objects.create(name="task1", user=user)
await Task.objects.create(name="task2", user=user)
tom_tasks = []
async for task in tom.tasks.iterate():
assert task.name in ("task1", "task2")
tom_tasks.append(task)
assert len(tom_tasks) == 2
jane_tasks = []
async for task in jane.tasks.iterate():
assert task.name in ("task1", "task2")
jane_tasks.append(task)
assert len(jane_tasks) == 2
@pytest.mark.asyncio
async def test_model_iterator_uneven_number_of_relations():
async with database:
async with database.transaction(force_rollback=True):
tom = await User.objects.create(name="Tom")
jane = await User.objects.create(name="Jane")
lucy = await User.objects.create(name="Lucy")
for user in tom, jane:
await Task.objects.create(name="task1", user=user)
await Task.objects.create(name="task2", user=user)
await Task.objects.create(name="task3", user=lucy)
expected_counts = {"Tom": 2, "Jane": 2, "Lucy": 1}
results = []
async for user in User.objects.select_related(User.tasks).iterate():
assert len(user.tasks) == expected_counts[user.name]
results.append(user)
assert len(results) == 3
@pytest.mark.asyncio
async def test_model_iterator_uuid_pk():
async with database:
async with database.transaction(force_rollback=True):
tom = await User2.objects.create(name="Tom")
jane = await User2.objects.create(name="Jane")
lucy = await User2.objects.create(name="Lucy")
async for user in User2.objects.iterate():
assert user in (tom, jane, lucy)
@pytest.mark.asyncio
async def test_model_iterator_filter_uuid_pk():
async with database:
async with database.transaction(force_rollback=True):
tom = await User2.objects.create(name="Tom")
await User2.objects.create(name="Jane")
await User2.objects.create(name="Lucy")
async for user in User2.objects.iterate(name="Tom"):
assert user.name == tom.name
@pytest.mark.asyncio
async def test_model_iterator_relations_uuid_pk():
async with database:
async with database.transaction(force_rollback=True):
tom = await User2.objects.create(name="Tom")
jane = await User2.objects.create(name="Jane")
lucy = await User2.objects.create(name="Lucy")
for user in tom, jane, lucy:
await Task2.objects.create(name="task1", user=user)
await Task2.objects.create(name="task2", user=user)
results = []
async for user in User2.objects.select_related(User2.task2s).iterate():
assert len(user.task2s) == 2
results.append(user)
assert len(results) == 3
@pytest.mark.asyncio
async def test_model_iterator_relations_queryset_proxy_uuid_pk():
async with database:
async with database.transaction(force_rollback=True):
tom = await User2.objects.create(name="Tom")
jane = await User2.objects.create(name="Jane")
for user in tom, jane:
await Task2.objects.create(name="task1", user=user)
await Task2.objects.create(name="task2", user=user)
tom_tasks = []
async for task in tom.task2s.iterate():
assert task.name in ("task1", "task2")
tom_tasks.append(task)
assert len(tom_tasks) == 2
jane_tasks = []
async for task in jane.task2s.iterate():
assert task.name in ("task1", "task2")
jane_tasks.append(task)
assert len(jane_tasks) == 2
@pytest.mark.asyncio
async def test_model_iterator_uneven_number_of_relations_uuid_pk():
async with database:
async with database.transaction(force_rollback=True):
tom = await User2.objects.create(name="Tom")
jane = await User2.objects.create(name="Jane")
lucy = await User2.objects.create(name="Lucy")
for user in tom, jane:
await Task2.objects.create(name="task1", user=user)
await Task2.objects.create(name="task2", user=user)
await Task2.objects.create(name="task3", user=lucy)
expected_counts = {"Tom": 2, "Jane": 2, "Lucy": 1}
results = []
async for user in User2.objects.select_related(User2.task2s).iterate():
assert len(user.task2s) == expected_counts[user.name]
results.append(user)
assert len(results) == 3
@pytest.mark.asyncio
async def test_model_iterator_with_prefetch_raises_error():
async with database:
with pytest.raises(QueryDefinitionError):
async for user in User.objects.prefetch_related(User.tasks).iterate():
pass # pragma: no cover