diff --git a/fastapi_startkit/src/fastapi_startkit/masoniteorm/models/builder.py b/fastapi_startkit/src/fastapi_startkit/masoniteorm/models/builder.py index b91bed5a..5d818105 100644 --- a/fastapi_startkit/src/fastapi_startkit/masoniteorm/models/builder.py +++ b/fastapi_startkit/src/fastapi_startkit/masoniteorm/models/builder.py @@ -368,29 +368,58 @@ async def chunk(self, count: int): break page += 1 - async def chunk_by_id(self, count: int, column: str = None, alias: str = None): + async def chunk_by_id(self, count: int, column: str = None, alias: str = None, descending: bool = False): if count <= 0: raise ValueError("chunk_by_id() size must be a positive integer.") column = column or self._model.__primary_key__ alias = alias or column + operator = "<" if descending else ">" + direction = "desc" if descending else "asc" + base_wheres = list(self._wheres) + offset = self._offset or 0 + remaining = self._limit if self._limit is not False else None + last_id = None + page = 1 while True: + limit = count if remaining is None else min(count, remaining) + if limit == 0: + break + self._wheres = list(base_wheres) + self._order_by = () + # The starting offset only applies to the first page; keyset + # filtering drives every page after that. + self._offset = offset if page == 1 else False if last_id is not None: - self.where(column, ">", last_id) + self.where(column, operator, last_id) - self._order_by = () - results = await self.order_by(column, "asc").limit(count).get() - if len(results) == 0: + results = await self.order_by(column, direction).limit(limit).get() + count_results = len(results) + if count_results == 0: break + if remaining is not None: + remaining = max(remaining - count_results, 0) + yield results - if len(results) < count: + last_id = results.last().get_attributes().get(alias) + if last_id is None: + raise RuntimeError( + f"The chunk_by_id operation was aborted because the [{alias}] " + "column is not present in the query result." + ) + + if count_results != count: break - last_id = getattr(results.last(), alias) + page += 1 + + async def chunk_by_id_desc(self, count: int, column: str = None, alias: str = None): + async for results in self.chunk_by_id(count, column, alias, descending=True): + yield results def new(self): return self.connection.query() diff --git a/fastapi_startkit/src/fastapi_startkit/masoniteorm/models/model.py b/fastapi_startkit/src/fastapi_startkit/masoniteorm/models/model.py index 9cf5c0a1..79c3614c 100644 --- a/fastapi_startkit/src/fastapi_startkit/masoniteorm/models/model.py +++ b/fastapi_startkit/src/fastapi_startkit/masoniteorm/models/model.py @@ -256,8 +256,12 @@ def chunk(cls, count: int): return cls.query().chunk(count) @classmethod - def chunk_by_id(cls, count: int, column: str = None, alias: str = None): - return cls.query().chunk_by_id(count, column, alias) + def chunk_by_id(cls, count: int, column: str = None, alias: str = None, descending: bool = False): + return cls.query().chunk_by_id(count, column, alias, descending) + + @classmethod + def chunk_by_id_desc(cls, count: int, column: str = None, alias: str = None): + return cls.query().chunk_by_id_desc(count, column, alias) def set_connection(self, connection: str): self.connection = connection diff --git a/fastapi_startkit/tests/masoniteorm/models/test_model_chunk.py b/fastapi_startkit/tests/masoniteorm/models/test_model_chunk.py index 52e33bb2..398077e4 100644 --- a/fastapi_startkit/tests/masoniteorm/models/test_model_chunk.py +++ b/fastapi_startkit/tests/masoniteorm/models/test_model_chunk.py @@ -1,156 +1,174 @@ """ -Tests for Laravel-style chunk() and chunk_by_id() on the async QueryBuilder -and their Model classmethod entry points. +Tests for Laravel-style chunk(), chunk_by_id() and chunk_by_id_desc() on the +async QueryBuilder and their Model classmethod entry points. -All tests run against an in-memory SQLite database so they are -self-contained and require no external services. +A dedicated ``chunk_items`` table (ids 1..10, odd ids active) is created and +seeded on top of the shared sqlite TestCase so the batching maths stay simple +and deterministic. The seeded ``countries`` table (country_id 10/20/30/40 on +the ``dev`` connection) exercises keyset paging on a non-``id`` primary key. """ -import pytest +from fastapi_startkit.masoniteorm import Model +from tests.masoniteorm.fixtures.model import Country +from tests.masoniteorm.sqlite.test_case import TestCase -from fastapi_startkit.masoniteorm.connections.factory import ConnectionFactory -from fastapi_startkit.masoniteorm.connections.manager import DatabaseManager -from fastapi_startkit.masoniteorm.models.model import Model -SQLITE_CONFIG = { - "default": "sqlite", - "connections": { - "sqlite": { - "driver": "sqlite", - "url": "sqlite+aiosqlite:///:memory:", - }, - }, -} +class ChunkItem(Model): + __table__ = "chunk_items" + __timestamps__ = False + id: int + name: str + active: bool -@pytest.fixture -async def db(): - manager = DatabaseManager(ConnectionFactory(), SQLITE_CONFIG) - yield manager - await manager.clear() +class ChunkTestCase(TestCase): + async def asyncSetUp(self): + await super().asyncSetUp() + await self.schema.drop_table_if_exists("chunk_items") + async with await self.schema.on("default").create("chunk_items") as table: + table.id() + table.string("name") + table.boolean("active").default(True) + await ChunkItem.query().insert([{"name": f"Item {i}", "active": i % 2 == 1} for i in range(1, 11)]) -@pytest.fixture -def ProjectModel(db): - class Project(Model): - id: int - name: str - active: bool + async def asyncTearDown(self): + await self.schema.drop_table_if_exists("chunk_items") + await super().asyncTearDown() - Project.db_manager = db - return Project +class TestChunk(ChunkTestCase): + async def test_yields_batches_covering_all_rows(self): + batches = [] + async for batch in ChunkItem.chunk(3): + batches.append(batch) -@pytest.fixture -async def projects_table(db, ProjectModel): - schema = db.get_schema_builder() - await schema.drop_table_if_exists("projects") - async with await schema.on("default").create("projects") as table: - table.id() - table.string("name") - table.boolean("active").default(True) - table.timestamps() - yield schema - await schema.drop_table_if_exists("projects") + self.assertEqual([len(b) for b in batches], [3, 3, 3, 1]) + ids = [item.id for batch in batches for item in batch] + self.assertEqual(ids, list(range(1, 11))) + async def test_batches_are_collections_of_hydrated_models(self): + async for batch in ChunkItem.chunk(4): + self.assertIsInstance(batch, ChunkItem().new_collection([]).__class__) + self.assertTrue(all(isinstance(item, ChunkItem) for item in batch)) + break -@pytest.fixture -async def seeded_projects(ProjectModel, projects_table): - """Insert 10 projects (ids 1..10); even-numbered ones are inactive.""" - await ProjectModel.query().insert([{"name": f"Project {i}", "active": i % 2 == 1} for i in range(1, 11)]) - return ProjectModel + async def test_exact_multiple_does_not_loop_forever(self): + batches = [len(b) async for b in ChunkItem.chunk(5)] + self.assertEqual(batches, [5, 5]) + async def test_chainable_after_where(self): + ids = [] + async for batch in ChunkItem.where("active", True).chunk(2): + ids.extend(item.id for item in batch) + self.assertEqual(ids, [1, 3, 5, 7, 9]) -# --------------------------------------------------------------------------- -# chunk() -# --------------------------------------------------------------------------- + async def test_empty_result_yields_nothing(self): + batches = [b async for b in ChunkItem.where("id", ">", 100).chunk(3)] + self.assertEqual(batches, []) + async def test_size_zero_raises(self): + with self.assertRaises(ValueError): + async for _ in ChunkItem.chunk(0): + pass -class TestChunk: - async def test_yields_batches_covering_all_rows(self, ProjectModel, seeded_projects): - batches = [] - async for batch in ProjectModel.chunk(3): - batches.append(batch) + async def test_negative_size_raises(self): + with self.assertRaises(ValueError): + async for _ in ChunkItem.chunk(-1): + pass - assert [len(b) for b in batches] == [3, 3, 3, 1] - ids = [p.id for batch in batches for p in batch] - assert ids == list(range(1, 11)) - async def test_batches_are_collections_of_hydrated_models(self, ProjectModel, seeded_projects): - async for batch in ProjectModel.chunk(4): - assert isinstance(batch, ProjectModel().new_collection([]).__class__) - assert all(isinstance(p, ProjectModel) for p in batch) +class TestChunkById(ChunkTestCase): + async def test_yields_all_rows_ordered_by_id(self): + ids = [] + async for batch in ChunkItem.chunk_by_id(3): + ids.extend(item.id for item in batch) + self.assertEqual(ids, list(range(1, 11))) + + async def test_batches_are_collections_of_hydrated_models(self): + async for batch in ChunkItem.chunk_by_id(4): + self.assertIsInstance(batch, ChunkItem().new_collection([]).__class__) + self.assertTrue(all(isinstance(item, ChunkItem) for item in batch)) break - async def test_exact_multiple_does_not_loop_forever(self, ProjectModel, seeded_projects): - batches = [b async for b in ProjectModel.chunk(5)] - assert [len(b) for b in batches] == [5, 5] + async def test_chainable_after_where(self): + ids = [] + async for batch in ChunkItem.where("active", True).chunk_by_id(2): + ids.extend(item.id for item in batch) + self.assertEqual(ids, [1, 3, 5, 7, 9]) - async def test_chainable_after_where(self, ProjectModel, seeded_projects): + async def test_custom_column_on_non_id_primary_key(self): ids = [] - async for batch in ProjectModel.where("active", True).chunk(2): - ids.extend(p.id for p in batch) - assert ids == [1, 3, 5, 7, 9] + async for batch in Country.chunk_by_id(2, column="country_id"): + ids.extend(c.country_id for c in batch) + self.assertEqual(ids, [10, 20, 30, 40]) - async def test_empty_table_yields_nothing(self, ProjectModel, projects_table): - batches = [b async for b in ProjectModel.chunk(3)] - assert batches == [] + async def test_safe_when_rows_deleted_mid_iteration(self): + seen = [] + async for batch in ChunkItem.chunk_by_id(2): + seen.extend(item.id for item in batch) + # Delete a not-yet-seen row to prove keyset paging never skips. + await ChunkItem.where("id", batch.last().id + 1).delete() + # 1,2 -> del 3; 4,5 -> del 6; 7,8 -> del 9; 10 + self.assertEqual(seen, [1, 2, 4, 5, 7, 8, 10]) + + async def test_respects_existing_limit(self): + batches = [] + async for batch in ChunkItem.limit(4).chunk_by_id(3): + batches.append([item.id for item in batch]) + # limit(4) chunked by 3 -> 3 + 1 + self.assertEqual(batches, [[1, 2, 3], [4]]) - async def test_size_zero_raises(self, ProjectModel, seeded_projects): - with pytest.raises(ValueError): - async for _ in ProjectModel.chunk(0): - pass + async def test_respects_existing_offset(self): + ids = [] + async for batch in ChunkItem.offset(2).chunk_by_id(3): + ids.extend(item.id for item in batch) + self.assertEqual(ids, [3, 4, 5, 6, 7, 8, 9, 10]) - async def test_negative_size_raises(self, ProjectModel, seeded_projects): - with pytest.raises(ValueError): - async for _ in ProjectModel.chunk(-1): + async def test_missing_alias_column_raises(self): + with self.assertRaises(RuntimeError): + async for _ in ChunkItem.select("name").chunk_by_id(3, alias="id"): pass + async def test_alias_reads_last_seen_from_given_column(self): + ids = [] + async for batch in ChunkItem.chunk_by_id(4, column="id", alias="id"): + ids.extend(item.id for item in batch) + self.assertEqual(ids, list(range(1, 11))) -# --------------------------------------------------------------------------- -# chunk_by_id() -# --------------------------------------------------------------------------- + async def test_size_zero_raises(self): + with self.assertRaises(ValueError): + async for _ in ChunkItem.chunk_by_id(0): + pass -class TestChunkById: - async def test_yields_all_rows_ordered_by_id(self, ProjectModel, seeded_projects): +class TestChunkByIdDesc(ChunkTestCase): + async def test_yields_all_rows_in_descending_id_order(self): ids = [] - async for batch in ProjectModel.chunk_by_id(3): - ids.extend(p.id for p in batch) - assert ids == list(range(1, 11)) - - async def test_batches_are_collections_of_hydrated_models(self, ProjectModel, seeded_projects): - async for batch in ProjectModel.chunk_by_id(4): - assert isinstance(batch, ProjectModel().new_collection([]).__class__) - assert all(isinstance(p, ProjectModel) for p in batch) - break + async for batch in ChunkItem.chunk_by_id_desc(3): + ids.extend(item.id for item in batch) + self.assertEqual(ids, list(range(10, 0, -1))) - async def test_chainable_after_where(self, ProjectModel, seeded_projects): - ids = [] - async for batch in ProjectModel.where("active", True).chunk_by_id(2): - ids.extend(p.id for p in batch) - assert ids == [1, 3, 5, 7, 9] + async def test_batch_sizes(self): + batches = [len(b) async for b in ChunkItem.chunk_by_id_desc(3)] + self.assertEqual(batches, [3, 3, 3, 1]) - async def test_custom_column(self, ProjectModel, seeded_projects): + async def test_chainable_after_where(self): ids = [] - async for batch in ProjectModel.chunk_by_id(4, column="id"): - ids.extend(p.id for p in batch) - assert ids == list(range(1, 11)) + async for batch in ChunkItem.where("active", True).chunk_by_id_desc(2): + ids.extend(item.id for item in batch) + self.assertEqual(ids, [9, 7, 5, 3, 1]) - async def test_safe_when_rows_deleted_mid_iteration(self, ProjectModel, seeded_projects): + async def test_safe_when_rows_deleted_mid_iteration(self): seen = [] - async for batch in ProjectModel.chunk_by_id(2): - seen.extend(p.id for p in batch) - # Delete a not-yet-seen row to prove keyset pagination doesn't skip. - await ProjectModel.where("id", batch.last().id + 1).delete() - # ids 1,2 -> delete 3; 4,5 -> delete 6; 7,8 -> delete 9; 10 - assert seen == [1, 2, 4, 5, 7, 8, 10] - - async def test_empty_table_yields_nothing(self, ProjectModel, projects_table): - batches = [b async for b in ProjectModel.chunk_by_id(3)] - assert batches == [] - - async def test_size_zero_raises(self, ProjectModel, seeded_projects): - with pytest.raises(ValueError): - async for _ in ProjectModel.chunk_by_id(0): + async for batch in ChunkItem.chunk_by_id_desc(2): + seen.extend(item.id for item in batch) + # Delete the next lower, not-yet-seen row. + await ChunkItem.where("id", batch.last().id - 1).delete() + # 10,9 -> del 8; 7,6 -> del 5; 4,3 -> del 2; 1 + self.assertEqual(seen, [10, 9, 7, 6, 4, 3, 1]) + + async def test_size_zero_raises(self): + with self.assertRaises(ValueError): + async for _ in ChunkItem.chunk_by_id_desc(0): pass