Skip to content

Commit f5c4a3b

Browse files
authored
feat(orm): Laravel-style chunk() and chunk_by_id() on the async ORM (#166)
* feat(orm): add Laravel-style chunk() and chunk_by_id() Add async-generator chunking to the QueryBuilder with Model classmethod entry points, both chainable after builder methods. - chunk(size): offset/limit paging, yields a Collection per batch and stops on an empty or short batch (no infinite loop). - chunk_by_id(size, column, alias): keyset pagination ordered by the primary key (or given column), filtering col > last_seen each pass so it stays correct when rows change mid-iteration. - Both guard against a non-positive size. Covered by sqlite-backed tests for batching, ordering, where-chaining, empty tables, delete-safety, and size validation. * refactor(orm): drop chunk_by_id docstring
1 parent 7afdbb6 commit f5c4a3b

3 files changed

Lines changed: 209 additions & 0 deletions

File tree

fastapi_startkit/src/fastapi_startkit/masoniteorm/models/builder.py

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -347,6 +347,51 @@ async def simple_paginate(self, per_page: int = 15, page: int = 1):
347347
results = await self.limit(per_page + 1).offset(offset).get()
348348
return SimplePaginator(results, per_page, page)
349349

350+
async def chunk(self, count: int):
351+
"""Yield results in batches of ``count`` using offset/limit paging.
352+
353+
Mirrors Laravel's ``chunk()``: iteration stops as soon as a batch
354+
comes back empty or shorter than ``count``, so it never loops forever.
355+
"""
356+
if count <= 0:
357+
raise ValueError("chunk() size must be a positive integer.")
358+
359+
page = 0
360+
while True:
361+
results = await self.limit(count).offset(page * count).get()
362+
if len(results) == 0:
363+
break
364+
365+
yield results
366+
367+
if len(results) < count:
368+
break
369+
page += 1
370+
371+
async def chunk_by_id(self, count: int, column: str = None, alias: str = None):
372+
if count <= 0:
373+
raise ValueError("chunk_by_id() size must be a positive integer.")
374+
375+
column = column or self._model.__primary_key__
376+
alias = alias or column
377+
base_wheres = list(self._wheres)
378+
last_id = None
379+
while True:
380+
self._wheres = list(base_wheres)
381+
if last_id is not None:
382+
self.where(column, ">", last_id)
383+
384+
self._order_by = ()
385+
results = await self.order_by(column, "asc").limit(count).get()
386+
if len(results) == 0:
387+
break
388+
389+
yield results
390+
391+
if len(results) < count:
392+
break
393+
last_id = getattr(results.last(), alias)
394+
350395
def new(self):
351396
return self.connection.query()
352397

fastapi_startkit/src/fastapi_startkit/masoniteorm/models/model.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -251,6 +251,14 @@ async def all(cls):
251251
async def count(cls, column: str = "*"):
252252
return await cls.query().count(column)
253253

254+
@classmethod
255+
def chunk(cls, count: int):
256+
return cls.query().chunk(count)
257+
258+
@classmethod
259+
def chunk_by_id(cls, count: int, column: str = None, alias: str = None):
260+
return cls.query().chunk_by_id(count, column, alias)
261+
254262
def set_connection(self, connection: str):
255263
self.connection = connection
256264

Lines changed: 156 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,156 @@
1+
"""
2+
Tests for Laravel-style chunk() and chunk_by_id() on the async QueryBuilder
3+
and their Model classmethod entry points.
4+
5+
All tests run against an in-memory SQLite database so they are
6+
self-contained and require no external services.
7+
"""
8+
9+
import pytest
10+
11+
from fastapi_startkit.masoniteorm.connections.factory import ConnectionFactory
12+
from fastapi_startkit.masoniteorm.connections.manager import DatabaseManager
13+
from fastapi_startkit.masoniteorm.models.model import Model
14+
15+
SQLITE_CONFIG = {
16+
"default": "sqlite",
17+
"connections": {
18+
"sqlite": {
19+
"driver": "sqlite",
20+
"url": "sqlite+aiosqlite:///:memory:",
21+
},
22+
},
23+
}
24+
25+
26+
@pytest.fixture
27+
async def db():
28+
manager = DatabaseManager(ConnectionFactory(), SQLITE_CONFIG)
29+
yield manager
30+
await manager.clear()
31+
32+
33+
@pytest.fixture
34+
def ProjectModel(db):
35+
class Project(Model):
36+
id: int
37+
name: str
38+
active: bool
39+
40+
Project.db_manager = db
41+
return Project
42+
43+
44+
@pytest.fixture
45+
async def projects_table(db, ProjectModel):
46+
schema = db.get_schema_builder()
47+
await schema.drop_table_if_exists("projects")
48+
async with await schema.on("default").create("projects") as table:
49+
table.id()
50+
table.string("name")
51+
table.boolean("active").default(True)
52+
table.timestamps()
53+
yield schema
54+
await schema.drop_table_if_exists("projects")
55+
56+
57+
@pytest.fixture
58+
async def seeded_projects(ProjectModel, projects_table):
59+
"""Insert 10 projects (ids 1..10); even-numbered ones are inactive."""
60+
await ProjectModel.query().insert([{"name": f"Project {i}", "active": i % 2 == 1} for i in range(1, 11)])
61+
return ProjectModel
62+
63+
64+
# ---------------------------------------------------------------------------
65+
# chunk()
66+
# ---------------------------------------------------------------------------
67+
68+
69+
class TestChunk:
70+
async def test_yields_batches_covering_all_rows(self, ProjectModel, seeded_projects):
71+
batches = []
72+
async for batch in ProjectModel.chunk(3):
73+
batches.append(batch)
74+
75+
assert [len(b) for b in batches] == [3, 3, 3, 1]
76+
ids = [p.id for batch in batches for p in batch]
77+
assert ids == list(range(1, 11))
78+
79+
async def test_batches_are_collections_of_hydrated_models(self, ProjectModel, seeded_projects):
80+
async for batch in ProjectModel.chunk(4):
81+
assert isinstance(batch, ProjectModel().new_collection([]).__class__)
82+
assert all(isinstance(p, ProjectModel) for p in batch)
83+
break
84+
85+
async def test_exact_multiple_does_not_loop_forever(self, ProjectModel, seeded_projects):
86+
batches = [b async for b in ProjectModel.chunk(5)]
87+
assert [len(b) for b in batches] == [5, 5]
88+
89+
async def test_chainable_after_where(self, ProjectModel, seeded_projects):
90+
ids = []
91+
async for batch in ProjectModel.where("active", True).chunk(2):
92+
ids.extend(p.id for p in batch)
93+
assert ids == [1, 3, 5, 7, 9]
94+
95+
async def test_empty_table_yields_nothing(self, ProjectModel, projects_table):
96+
batches = [b async for b in ProjectModel.chunk(3)]
97+
assert batches == []
98+
99+
async def test_size_zero_raises(self, ProjectModel, seeded_projects):
100+
with pytest.raises(ValueError):
101+
async for _ in ProjectModel.chunk(0):
102+
pass
103+
104+
async def test_negative_size_raises(self, ProjectModel, seeded_projects):
105+
with pytest.raises(ValueError):
106+
async for _ in ProjectModel.chunk(-1):
107+
pass
108+
109+
110+
# ---------------------------------------------------------------------------
111+
# chunk_by_id()
112+
# ---------------------------------------------------------------------------
113+
114+
115+
class TestChunkById:
116+
async def test_yields_all_rows_ordered_by_id(self, ProjectModel, seeded_projects):
117+
ids = []
118+
async for batch in ProjectModel.chunk_by_id(3):
119+
ids.extend(p.id for p in batch)
120+
assert ids == list(range(1, 11))
121+
122+
async def test_batches_are_collections_of_hydrated_models(self, ProjectModel, seeded_projects):
123+
async for batch in ProjectModel.chunk_by_id(4):
124+
assert isinstance(batch, ProjectModel().new_collection([]).__class__)
125+
assert all(isinstance(p, ProjectModel) for p in batch)
126+
break
127+
128+
async def test_chainable_after_where(self, ProjectModel, seeded_projects):
129+
ids = []
130+
async for batch in ProjectModel.where("active", True).chunk_by_id(2):
131+
ids.extend(p.id for p in batch)
132+
assert ids == [1, 3, 5, 7, 9]
133+
134+
async def test_custom_column(self, ProjectModel, seeded_projects):
135+
ids = []
136+
async for batch in ProjectModel.chunk_by_id(4, column="id"):
137+
ids.extend(p.id for p in batch)
138+
assert ids == list(range(1, 11))
139+
140+
async def test_safe_when_rows_deleted_mid_iteration(self, ProjectModel, seeded_projects):
141+
seen = []
142+
async for batch in ProjectModel.chunk_by_id(2):
143+
seen.extend(p.id for p in batch)
144+
# Delete a not-yet-seen row to prove keyset pagination doesn't skip.
145+
await ProjectModel.where("id", batch.last().id + 1).delete()
146+
# ids 1,2 -> delete 3; 4,5 -> delete 6; 7,8 -> delete 9; 10
147+
assert seen == [1, 2, 4, 5, 7, 8, 10]
148+
149+
async def test_empty_table_yields_nothing(self, ProjectModel, projects_table):
150+
batches = [b async for b in ProjectModel.chunk_by_id(3)]
151+
assert batches == []
152+
153+
async def test_size_zero_raises(self, ProjectModel, seeded_projects):
154+
with pytest.raises(ValueError):
155+
async for _ in ProjectModel.chunk_by_id(0):
156+
pass

0 commit comments

Comments
 (0)