From 813a57b1070a1f6ac41d02897fab33f8039b83f9 Mon Sep 17 00:00:00 2001 From: Craig Labenz Date: Mon, 16 Aug 2021 15:11:56 -0700 Subject: [PATCH] feat: add recursive delete (#420) * feat: add recursive delete * made chunkify private Co-authored-by: Christopher Wilcox --- google/cloud/firestore_v1/async_client.py | 84 ++++- google/cloud/firestore_v1/async_collection.py | 4 + google/cloud/firestore_v1/async_query.py | 44 ++- google/cloud/firestore_v1/base_client.py | 13 +- google/cloud/firestore_v1/base_document.py | 4 +- google/cloud/firestore_v1/base_query.py | 6 + google/cloud/firestore_v1/client.py | 88 +++++- google/cloud/firestore_v1/collection.py | 3 + google/cloud/firestore_v1/query.py | 44 ++- tests/system/test_system.py | 274 +++++++++++++---- tests/system/test_system_async.py | 287 +++++++++++++----- tests/unit/v1/test_async_client.py | 106 +++++++ tests/unit/v1/test_async_collection.py | 38 +++ tests/unit/v1/test_async_query.py | 24 ++ tests/unit/v1/test_client.py | 100 ++++++ tests/unit/v1/test_collection.py | 37 +++ tests/unit/v1/test_query.py | 36 +++ 17 files changed, 1046 insertions(+), 146 deletions(-) diff --git a/google/cloud/firestore_v1/async_client.py b/google/cloud/firestore_v1/async_client.py index 68cb676f2..a4be11002 100644 --- a/google/cloud/firestore_v1/async_client.py +++ b/google/cloud/firestore_v1/async_client.py @@ -43,13 +43,17 @@ DocumentSnapshot, ) from google.cloud.firestore_v1.async_transaction import AsyncTransaction +from google.cloud.firestore_v1.field_path import FieldPath from google.cloud.firestore_v1.services.firestore import ( async_client as firestore_client, ) from google.cloud.firestore_v1.services.firestore.transports import ( grpc_asyncio as firestore_grpc_transport, ) -from typing import Any, AsyncGenerator, Iterable, List +from typing import Any, AsyncGenerator, Iterable, List, Optional, Union, TYPE_CHECKING + +if TYPE_CHECKING: + from google.cloud.firestore_v1.bulk_writer import BulkWriter # pragma: NO COVER class AsyncClient(BaseClient): @@ -300,6 +304,84 @@ async def collections( async for collection_id in iterator: yield self.collection(collection_id) + async def recursive_delete( + self, + reference: Union[AsyncCollectionReference, AsyncDocumentReference], + *, + bulk_writer: Optional["BulkWriter"] = None, + chunk_size: Optional[int] = 5000, + ): + """Deletes documents and their subcollections, regardless of collection + name. + + Passing an AsyncCollectionReference leads to each document in the + collection getting deleted, as well as all of their descendents. + + Passing an AsyncDocumentReference deletes that one document and all of + its descendents. + + Args: + reference (Union[ + :class:`@google.cloud.firestore_v1.async_collection.CollectionReference`, + :class:`@google.cloud.firestore_v1.async_document.DocumentReference`, + ]) + The reference to be deleted. + + bulk_writer (Optional[:class:`@google.cloud.firestore_v1.bulk_writer.BulkWriter`]) + The BulkWriter used to delete all matching documents. Supply this + if you want to override the default throttling behavior. + """ + return await self._recursive_delete( + reference, bulk_writer=bulk_writer, chunk_size=chunk_size, + ) + + async def _recursive_delete( + self, + reference: Union[AsyncCollectionReference, AsyncDocumentReference], + *, + bulk_writer: Optional["BulkWriter"] = None, # type: ignore + chunk_size: Optional[int] = 5000, + depth: Optional[int] = 0, + ) -> int: + """Recursion helper for `recursive_delete.""" + from google.cloud.firestore_v1.bulk_writer import BulkWriter + + bulk_writer = bulk_writer or BulkWriter() + + num_deleted: int = 0 + + if isinstance(reference, AsyncCollectionReference): + chunk: List[DocumentSnapshot] + async for chunk in reference.recursive().select( + [FieldPath.document_id()] + )._chunkify(chunk_size): + doc_snap: DocumentSnapshot + for doc_snap in chunk: + num_deleted += 1 + bulk_writer.delete(doc_snap.reference) + + elif isinstance(reference, AsyncDocumentReference): + col_ref: AsyncCollectionReference + async for col_ref in reference.collections(): + num_deleted += await self._recursive_delete( + col_ref, + bulk_writer=bulk_writer, + depth=depth + 1, + chunk_size=chunk_size, + ) + num_deleted += 1 + bulk_writer.delete(reference) + + else: + raise TypeError( + f"Unexpected type for reference: {reference.__class__.__name__}" + ) + + if depth == 0: + bulk_writer.close() + + return num_deleted + def batch(self) -> AsyncWriteBatch: """Get a batch instance from this client. diff --git a/google/cloud/firestore_v1/async_collection.py b/google/cloud/firestore_v1/async_collection.py index ca4ec8b0f..d06405127 100644 --- a/google/cloud/firestore_v1/async_collection.py +++ b/google/cloud/firestore_v1/async_collection.py @@ -72,6 +72,10 @@ def _query(self) -> async_query.AsyncQuery: """ return async_query.AsyncQuery(self) + async def _chunkify(self, chunk_size: int): + async for page in self._query()._chunkify(chunk_size): + yield page + async def add( self, document_data: dict, diff --git a/google/cloud/firestore_v1/async_query.py b/google/cloud/firestore_v1/async_query.py index 2f94b5f7c..0444b92bc 100644 --- a/google/cloud/firestore_v1/async_query.py +++ b/google/cloud/firestore_v1/async_query.py @@ -33,7 +33,8 @@ ) from google.cloud.firestore_v1 import async_document -from typing import AsyncGenerator, Type +from google.cloud.firestore_v1.base_document import DocumentSnapshot +from typing import AsyncGenerator, List, Optional, Type # Types needed only for Type Hints from google.cloud.firestore_v1.transaction import Transaction @@ -126,6 +127,47 @@ def __init__( recursive=recursive, ) + async def _chunkify( + self, chunk_size: int + ) -> AsyncGenerator[List[DocumentSnapshot], None]: + # Catch the edge case where a developer writes the following: + # `my_query.limit(500)._chunkify(1000)`, which ultimately nullifies any + # need to yield chunks. + if self._limit and chunk_size > self._limit: + yield await self.get() + return + + max_to_return: Optional[int] = self._limit + num_returned: int = 0 + original: AsyncQuery = self._copy() + last_document: Optional[DocumentSnapshot] = None + + while True: + # Optionally trim the `chunk_size` down to honor a previously + # applied limit as set by `self.limit()` + _chunk_size: int = original._resolve_chunk_size(num_returned, chunk_size) + + # Apply the optionally pruned limit and the cursor, if we are past + # the first page. + _q = original.limit(_chunk_size) + if last_document: + _q = _q.start_after(last_document) + + snapshots = await _q.get() + last_document = snapshots[-1] + num_returned += len(snapshots) + + yield snapshots + + # Terminate the iterator if we have reached either of two end + # conditions: + # 1. There are no more documents, or + # 2. We have reached the desired overall limit + if len(snapshots) < _chunk_size or ( + max_to_return and num_returned >= max_to_return + ): + return + async def get( self, transaction: Transaction = None, diff --git a/google/cloud/firestore_v1/base_client.py b/google/cloud/firestore_v1/base_client.py index e68031ed4..17068a974 100644 --- a/google/cloud/firestore_v1/base_client.py +++ b/google/cloud/firestore_v1/base_client.py @@ -37,11 +37,9 @@ from google.cloud.firestore_v1 import __version__ from google.cloud.firestore_v1 import types from google.cloud.firestore_v1.base_document import DocumentSnapshot -from google.cloud.firestore_v1.bulk_writer import ( - BulkWriter, - BulkWriterOptions, -) + from google.cloud.firestore_v1.field_path import render_field_path +from google.cloud.firestore_v1.bulk_writer import BulkWriter, BulkWriterOptions from typing import ( Any, AsyncGenerator, @@ -312,6 +310,13 @@ def _document_path_helper(self, *document_path) -> List[str]: joined_path = joined_path[len(base_path) :] return joined_path.split(_helpers.DOCUMENT_PATH_DELIMITER) + def recursive_delete( + self, + reference: Union[BaseCollectionReference, BaseDocumentReference], + bulk_writer: Optional["BulkWriter"] = None, # type: ignore + ) -> int: + raise NotImplementedError + @staticmethod def field_path(*field_names: str) -> str: """Create a **field path** from a list of nested field names. diff --git a/google/cloud/firestore_v1/base_document.py b/google/cloud/firestore_v1/base_document.py index 32694ac47..9e15b108c 100644 --- a/google/cloud/firestore_v1/base_document.py +++ b/google/cloud/firestore_v1/base_document.py @@ -315,10 +315,10 @@ def _prep_collections( def collections( self, page_size: int = None, retry: retries.Retry = None, timeout: float = None, - ) -> NoReturn: + ) -> None: raise NotImplementedError - def on_snapshot(self, callback) -> NoReturn: + def on_snapshot(self, callback) -> None: raise NotImplementedError diff --git a/google/cloud/firestore_v1/base_query.py b/google/cloud/firestore_v1/base_query.py index 1812cfca0..4f3ee101f 100644 --- a/google/cloud/firestore_v1/base_query.py +++ b/google/cloud/firestore_v1/base_query.py @@ -424,6 +424,12 @@ def limit_to_last(self, count: int) -> "BaseQuery": """ return self._copy(limit=count, limit_to_last=True) + def _resolve_chunk_size(self, num_loaded: int, chunk_size: int) -> int: + """Utility function for chunkify.""" + if self._limit is not None and (num_loaded + chunk_size) > self._limit: + return max(self._limit - num_loaded, 0) + return chunk_size + def offset(self, num_to_skip: int) -> "BaseQuery": """Skip to an offset in a query. diff --git a/google/cloud/firestore_v1/client.py b/google/cloud/firestore_v1/client.py index 20ef5055f..750acb0be 100644 --- a/google/cloud/firestore_v1/client.py +++ b/google/cloud/firestore_v1/client.py @@ -39,17 +39,22 @@ from google.cloud.firestore_v1.batch import WriteBatch from google.cloud.firestore_v1.collection import CollectionReference from google.cloud.firestore_v1.document import DocumentReference +from google.cloud.firestore_v1.field_path import FieldPath from google.cloud.firestore_v1.transaction import Transaction from google.cloud.firestore_v1.services.firestore import client as firestore_client from google.cloud.firestore_v1.services.firestore.transports import ( grpc as firestore_grpc_transport, ) -from typing import Any, Generator, Iterable +from typing import Any, Generator, Iterable, List, Optional, Union, TYPE_CHECKING # Types needed only for Type Hints from google.cloud.firestore_v1.base_document import DocumentSnapshot +if TYPE_CHECKING: + from google.cloud.firestore_v1.bulk_writer import BulkWriter # pragma: NO COVER + + class Client(BaseClient): """Client for interacting with Google Cloud Firestore API. @@ -286,6 +291,87 @@ def collections( for collection_id in iterator: yield self.collection(collection_id) + def recursive_delete( + self, + reference: Union[CollectionReference, DocumentReference], + *, + bulk_writer: Optional["BulkWriter"] = None, + chunk_size: Optional[int] = 5000, + ) -> int: + """Deletes documents and their subcollections, regardless of collection + name. + + Passing a CollectionReference leads to each document in the collection + getting deleted, as well as all of their descendents. + + Passing a DocumentReference deletes that one document and all of its + descendents. + + Args: + reference (Union[ + :class:`@google.cloud.firestore_v1.collection.CollectionReference`, + :class:`@google.cloud.firestore_v1.document.DocumentReference`, + ]) + The reference to be deleted. + + bulk_writer (Optional[:class:`@google.cloud.firestore_v1.bulk_writer.BulkWriter`]) + The BulkWriter used to delete all matching documents. Supply this + if you want to override the default throttling behavior. + + """ + return self._recursive_delete( + reference, bulk_writer=bulk_writer, chunk_size=chunk_size, + ) + + def _recursive_delete( + self, + reference: Union[CollectionReference, DocumentReference], + *, + bulk_writer: Optional["BulkWriter"] = None, + chunk_size: Optional[int] = 5000, + depth: Optional[int] = 0, + ) -> int: + """Recursion helper for `recursive_delete.""" + from google.cloud.firestore_v1.bulk_writer import BulkWriter + + bulk_writer = bulk_writer or BulkWriter() + + num_deleted: int = 0 + + if isinstance(reference, CollectionReference): + chunk: List[DocumentSnapshot] + for chunk in ( + reference.recursive() + .select([FieldPath.document_id()]) + ._chunkify(chunk_size) + ): + doc_snap: DocumentSnapshot + for doc_snap in chunk: + num_deleted += 1 + bulk_writer.delete(doc_snap.reference) + + elif isinstance(reference, DocumentReference): + col_ref: CollectionReference + for col_ref in reference.collections(): + num_deleted += self._recursive_delete( + col_ref, + bulk_writer=bulk_writer, + chunk_size=chunk_size, + depth=depth + 1, + ) + num_deleted += 1 + bulk_writer.delete(reference) + + else: + raise TypeError( + f"Unexpected type for reference: {reference.__class__.__name__}" + ) + + if depth == 0: + bulk_writer.close() + + return num_deleted + def batch(self) -> WriteBatch: """Get a batch instance from this client. diff --git a/google/cloud/firestore_v1/collection.py b/google/cloud/firestore_v1/collection.py index 96d076e2c..643e2d7ef 100644 --- a/google/cloud/firestore_v1/collection.py +++ b/google/cloud/firestore_v1/collection.py @@ -137,6 +137,9 @@ def list_documents( ) return (_item_to_document_ref(self, i) for i in iterator) + def _chunkify(self, chunk_size: int): + return self._query()._chunkify(chunk_size) + def get( self, transaction: Transaction = None, diff --git a/google/cloud/firestore_v1/query.py b/google/cloud/firestore_v1/query.py index f1e044cbd..50c5559b1 100644 --- a/google/cloud/firestore_v1/query.py +++ b/google/cloud/firestore_v1/query.py @@ -18,7 +18,6 @@ a :class:`~google.cloud.firestore_v1.collection.Collection` and that can be a more common way to create a query than direct usage of the constructor. """ - from google.cloud import firestore_v1 from google.cloud.firestore_v1.base_document import DocumentSnapshot from google.api_core import gapic_v1 # type: ignore @@ -35,7 +34,7 @@ from google.cloud.firestore_v1 import document from google.cloud.firestore_v1.watch import Watch -from typing import Any, Callable, Generator, List, Type +from typing import Any, Callable, Generator, List, Optional, Type class Query(BaseQuery): @@ -168,6 +167,47 @@ def get( return list(result) + def _chunkify( + self, chunk_size: int + ) -> Generator[List[DocumentSnapshot], None, None]: + # Catch the edge case where a developer writes the following: + # `my_query.limit(500)._chunkify(1000)`, which ultimately nullifies any + # need to yield chunks. + if self._limit and chunk_size > self._limit: + yield self.get() + return + + max_to_return: Optional[int] = self._limit + num_returned: int = 0 + original: Query = self._copy() + last_document: Optional[DocumentSnapshot] = None + + while True: + # Optionally trim the `chunk_size` down to honor a previously + # applied limits as set by `self.limit()` + _chunk_size: int = original._resolve_chunk_size(num_returned, chunk_size) + + # Apply the optionally pruned limit and the cursor, if we are past + # the first page. + _q = original.limit(_chunk_size) + if last_document: + _q = _q.start_after(last_document) + + snapshots = _q.get() + last_document = snapshots[-1] + num_returned += len(snapshots) + + yield snapshots + + # Terminate the iterator if we have reached either of two end + # conditions: + # 1. There are no more documents, or + # 2. We have reached the desired overall limit + if len(snapshots) < _chunk_size or ( + max_to_return and num_returned >= max_to_return + ): + return + def stream( self, transaction=None, diff --git a/tests/system/test_system.py b/tests/system/test_system.py index 0975a73d0..109029ced 100644 --- a/tests/system/test_system.py +++ b/tests/system/test_system.py @@ -29,6 +29,7 @@ from google.cloud import firestore_v1 as firestore from time import sleep +from typing import Callable, Dict, List, Optional from tests.system.test__helpers import ( FIRESTORE_CREDS, @@ -1235,65 +1236,157 @@ def test_array_union(client, cleanup): assert doc_ref.get().to_dict() == expected -def test_recursive_query(client, cleanup): +def _persist_documents( + client: firestore.Client, + collection_name: str, + documents: List[Dict], + cleanup: Optional[Callable] = None, +): + """Assuming `documents` is a recursive list of dictionaries representing + documents and subcollections, this method writes all of those through + `client.collection(...).document(...).create()`. - philosophers = [ + `documents` must be of this structure: + ```py + documents = [ { - "data": {"name": "Socrates", "favoriteCity": "Athens"}, - "subcollections": { - "pets": [{"name": "Scruffy"}, {"name": "Snowflake"}], - "hobbies": [{"name": "pontificating"}, {"name": "journaling"}], - "philosophers": [{"name": "Aristotle"}, {"name": "Plato"}], - }, + # Required key + "data": , + + # Optional key + "subcollections": , }, - { - "data": {"name": "Aristotle", "favoriteCity": "Sparta"}, - "subcollections": { - "pets": [{"name": "Floof-Boy"}, {"name": "Doggy-Dog"}], - "hobbies": [{"name": "questioning-stuff"}, {"name": "meditation"}], - }, + ... + ] + ``` + """ + for block in documents: + col_ref = client.collection(collection_name) + document_id: str = block["data"]["name"] + doc_ref = col_ref.document(document_id) + doc_ref.set(block["data"]) + if cleanup is not None: + cleanup(doc_ref.delete) + + if "subcollections" in block: + for subcollection_name, inner_blocks in block["subcollections"].items(): + _persist_documents( + client, + f"{collection_name}/{document_id}/{subcollection_name}", + inner_blocks, + ) + + +# documents compatible with `_persist_documents` +philosophers_data_set = [ + { + "data": {"name": "Socrates", "favoriteCity": "Athens"}, + "subcollections": { + "pets": [{"data": {"name": "Scruffy"}}, {"data": {"name": "Snowflake"}}], + "hobbies": [ + {"data": {"name": "pontificating"}}, + {"data": {"name": "journaling"}}, + ], + "philosophers": [ + {"data": {"name": "Aristotle"}}, + {"data": {"name": "Plato"}}, + ], }, - { - "data": {"name": "Plato", "favoriteCity": "Corinth"}, - "subcollections": { - "pets": [{"name": "Cuddles"}, {"name": "Sergeant-Puppers"}], - "hobbies": [{"name": "abstraction"}, {"name": "hypotheticals"}], - }, + }, + { + "data": {"name": "Aristotle", "favoriteCity": "Sparta"}, + "subcollections": { + "pets": [{"data": {"name": "Floof-Boy"}}, {"data": {"name": "Doggy-Dog"}}], + "hobbies": [ + {"data": {"name": "questioning-stuff"}}, + {"data": {"name": "meditation"}}, + ], + }, + }, + { + "data": {"name": "Plato", "favoriteCity": "Corinth"}, + "subcollections": { + "pets": [ + {"data": {"name": "Cuddles"}}, + {"data": {"name": "Sergeant-Puppers"}}, + ], + "hobbies": [ + {"data": {"name": "abstraction"}}, + {"data": {"name": "hypotheticals"}}, + ], }, + }, +] + + +def _do_recursive_delete_with_bulk_writer(client, bulk_writer): + philosophers = [philosophers_data_set[0]] + _persist_documents(client, f"philosophers{UNIQUE_RESOURCE_ID}", philosophers) + + doc_paths = [ + "", + "/pets/Scruffy", + "/pets/Snowflake", + "/hobbies/pontificating", + "/hobbies/journaling", + "/philosophers/Aristotle", + "/philosophers/Plato", ] - db = client - collection_ref = db.collection("philosophers") - for philosopher in philosophers: - ref = collection_ref.document( - f"{philosopher['data']['name']}{UNIQUE_RESOURCE_ID}" - ) - ref.set(philosopher["data"]) - cleanup(ref.delete) - for col_name, entries in philosopher["subcollections"].items(): - sub_col = ref.collection(col_name) - for entry in entries: - inner_doc_ref = sub_col.document(entry["name"]) - inner_doc_ref.set(entry) - cleanup(inner_doc_ref.delete) + # Assert all documents were created so that when they're missing after the + # delete, we're actually testing something. + collection_ref = client.collection(f"philosophers{UNIQUE_RESOURCE_ID}") + for path in doc_paths: + snapshot = collection_ref.document(f"Socrates{path}").get() + assert snapshot.exists, f"Snapshot at Socrates{path} should have been created" + + # Now delete. + num_deleted = client.recursive_delete(collection_ref, bulk_writer=bulk_writer) + assert num_deleted == len(doc_paths) + + # Now they should all be missing + for path in doc_paths: + snapshot = collection_ref.document(f"Socrates{path}").get() + assert ( + not snapshot.exists + ), f"Snapshot at Socrates{path} should have been deleted" + + +def test_recursive_delete_parallelized(client, cleanup): + from google.cloud.firestore_v1.bulk_writer import BulkWriterOptions, SendMode + + bw = client.bulk_writer(options=BulkWriterOptions(mode=SendMode.parallel)) + _do_recursive_delete_with_bulk_writer(client, bw) + - ids = [doc.id for doc in db.collection_group("philosophers").recursive().get()] +def test_recursive_delete_serialized(client, cleanup): + from google.cloud.firestore_v1.bulk_writer import BulkWriterOptions, SendMode + + bw = client.bulk_writer(options=BulkWriterOptions(mode=SendMode.serial)) + _do_recursive_delete_with_bulk_writer(client, bw) + + +def test_recursive_query(client, cleanup): + col_id: str = f"philosophers-recursive-query{UNIQUE_RESOURCE_ID}" + _persist_documents(client, col_id, philosophers_data_set, cleanup) + + ids = [doc.id for doc in client.collection_group(col_id).recursive().get()] expected_ids = [ # Aristotle doc and subdocs - f"Aristotle{UNIQUE_RESOURCE_ID}", + "Aristotle", "meditation", "questioning-stuff", "Doggy-Dog", "Floof-Boy", # Plato doc and subdocs - f"Plato{UNIQUE_RESOURCE_ID}", + "Plato", "abstraction", "hypotheticals", "Cuddles", "Sergeant-Puppers", # Socrates doc and subdocs - f"Socrates{UNIQUE_RESOURCE_ID}", + "Socrates", "journaling", "pontificating", "Scruffy", @@ -1312,34 +1405,12 @@ def test_recursive_query(client, cleanup): def test_nested_recursive_query(client, cleanup): + col_id: str = f"philosophers-nested-recursive-query{UNIQUE_RESOURCE_ID}" + _persist_documents(client, col_id, philosophers_data_set, cleanup) - philosophers = [ - { - "data": {"name": "Aristotle", "favoriteCity": "Sparta"}, - "subcollections": { - "pets": [{"name": "Floof-Boy"}, {"name": "Doggy-Dog"}], - "hobbies": [{"name": "questioning-stuff"}, {"name": "meditation"}], - }, - }, - ] - - db = client - collection_ref = db.collection("philosophers") - for philosopher in philosophers: - ref = collection_ref.document( - f"{philosopher['data']['name']}{UNIQUE_RESOURCE_ID}" - ) - ref.set(philosopher["data"]) - cleanup(ref.delete) - for col_name, entries in philosopher["subcollections"].items(): - sub_col = ref.collection(col_name) - for entry in entries: - inner_doc_ref = sub_col.document(entry["name"]) - inner_doc_ref.set(entry) - cleanup(inner_doc_ref.delete) - - aristotle = collection_ref.document(f"Aristotle{UNIQUE_RESOURCE_ID}") - ids = [doc.id for doc in aristotle.collection("pets")._query().recursive().get()] + collection_ref = client.collection(col_id) + aristotle = collection_ref.document("Aristotle") + ids = [doc.id for doc in aristotle.collection("pets").recursive().get()] expected_ids = [ # Aristotle pets @@ -1356,6 +1427,79 @@ def test_nested_recursive_query(client, cleanup): assert ids[index] == expected_ids[index], error_msg +def test_chunked_query(client, cleanup): + col = client.collection(f"chunked-test{UNIQUE_RESOURCE_ID}") + for index in range(10): + doc_ref = col.document(f"document-{index + 1}") + doc_ref.set({"index": index}) + cleanup(doc_ref.delete) + + iter = col._chunkify(3) + assert len(next(iter)) == 3 + assert len(next(iter)) == 3 + assert len(next(iter)) == 3 + assert len(next(iter)) == 1 + + +def test_chunked_query_smaller_limit(client, cleanup): + col = client.collection(f"chunked-test-smaller-limit{UNIQUE_RESOURCE_ID}") + for index in range(10): + doc_ref = col.document(f"document-{index + 1}") + doc_ref.set({"index": index}) + cleanup(doc_ref.delete) + + iter = col.limit(5)._chunkify(9) + assert len(next(iter)) == 5 + + +def test_chunked_and_recursive(client, cleanup): + col_id = f"chunked-recursive-test{UNIQUE_RESOURCE_ID}" + documents = [ + { + "data": {"name": "Root-1"}, + "subcollections": { + "children": [ + {"data": {"name": f"Root-1--Child-{index + 1}"}} + for index in range(5) + ] + }, + }, + { + "data": {"name": "Root-2"}, + "subcollections": { + "children": [ + {"data": {"name": f"Root-2--Child-{index + 1}"}} + for index in range(5) + ] + }, + }, + ] + _persist_documents(client, col_id, documents, cleanup) + collection_ref = client.collection(col_id) + iter = collection_ref.recursive()._chunkify(5) + + page_1_ids = [ + "Root-1", + "Root-1--Child-1", + "Root-1--Child-2", + "Root-1--Child-3", + "Root-1--Child-4", + ] + assert [doc.id for doc in next(iter)] == page_1_ids + + page_2_ids = [ + "Root-1--Child-5", + "Root-2", + "Root-2--Child-1", + "Root-2--Child-2", + "Root-2--Child-3", + ] + assert [doc.id for doc in next(iter)] == page_2_ids + + page_3_ids = ["Root-2--Child-4", "Root-2--Child-5"] + assert [doc.id for doc in next(iter)] == page_3_ids + + def test_watch_query_order(client, cleanup): db = client collection_ref = db.collection("users") diff --git a/tests/system/test_system_async.py b/tests/system/test_system_async.py index a4db4e75f..b7c562fd3 100644 --- a/tests/system/test_system_async.py +++ b/tests/system/test_system_async.py @@ -18,6 +18,7 @@ import math import pytest import operator +from typing import Callable, Dict, List, Optional from google.oauth2 import service_account @@ -1094,67 +1095,159 @@ async def test_batch(client, cleanup): assert not (await document3.get()).exists -async def test_recursive_query(client, cleanup): +async def _persist_documents( + client: firestore.AsyncClient, + collection_name: str, + documents: List[Dict], + cleanup: Optional[Callable] = None, +): + """Assuming `documents` is a recursive list of dictionaries representing + documents and subcollections, this method writes all of those through + `client.collection(...).document(...).create()`. - philosophers = [ + `documents` must be of this structure: + ```py + documents = [ { - "data": {"name": "Socrates", "favoriteCity": "Athens"}, - "subcollections": { - "pets": [{"name": "Scruffy"}, {"name": "Snowflake"}], - "hobbies": [{"name": "pontificating"}, {"name": "journaling"}], - "philosophers": [{"name": "Aristotle"}, {"name": "Plato"}], - }, + # Required key + "data": , + + # Optional key + "subcollections": , }, - { - "data": {"name": "Aristotle", "favoriteCity": "Sparta"}, - "subcollections": { - "pets": [{"name": "Floof-Boy"}, {"name": "Doggy-Dog"}], - "hobbies": [{"name": "questioning-stuff"}, {"name": "meditation"}], - }, + ... + ] + ``` + """ + for block in documents: + col_ref = client.collection(collection_name) + document_id: str = block["data"]["name"] + doc_ref = col_ref.document(document_id) + await doc_ref.set(block["data"]) + if cleanup is not None: + cleanup(doc_ref.delete) + + if "subcollections" in block: + for subcollection_name, inner_blocks in block["subcollections"].items(): + await _persist_documents( + client, + f"{collection_name}/{document_id}/{subcollection_name}", + inner_blocks, + ) + + +# documents compatible with `_persist_documents` +philosophers_data_set = [ + { + "data": {"name": "Socrates", "favoriteCity": "Athens"}, + "subcollections": { + "pets": [{"data": {"name": "Scruffy"}}, {"data": {"name": "Snowflake"}}], + "hobbies": [ + {"data": {"name": "pontificating"}}, + {"data": {"name": "journaling"}}, + ], + "philosophers": [ + {"data": {"name": "Aristotle"}}, + {"data": {"name": "Plato"}}, + ], }, - { - "data": {"name": "Plato", "favoriteCity": "Corinth"}, - "subcollections": { - "pets": [{"name": "Cuddles"}, {"name": "Sergeant-Puppers"}], - "hobbies": [{"name": "abstraction"}, {"name": "hypotheticals"}], - }, + }, + { + "data": {"name": "Aristotle", "favoriteCity": "Sparta"}, + "subcollections": { + "pets": [{"data": {"name": "Floof-Boy"}}, {"data": {"name": "Doggy-Dog"}}], + "hobbies": [ + {"data": {"name": "questioning-stuff"}}, + {"data": {"name": "meditation"}}, + ], }, - ] + }, + { + "data": {"name": "Plato", "favoriteCity": "Corinth"}, + "subcollections": { + "pets": [ + {"data": {"name": "Cuddles"}}, + {"data": {"name": "Sergeant-Puppers"}}, + ], + "hobbies": [ + {"data": {"name": "abstraction"}}, + {"data": {"name": "hypotheticals"}}, + ], + }, + }, +] - db = client - collection_ref = db.collection("philosophers") - for philosopher in philosophers: - ref = collection_ref.document( - f"{philosopher['data']['name']}{UNIQUE_RESOURCE_ID}-async" - ) - await ref.set(philosopher["data"]) - cleanup(ref.delete) - for col_name, entries in philosopher["subcollections"].items(): - sub_col = ref.collection(col_name) - for entry in entries: - inner_doc_ref = sub_col.document(entry["name"]) - await inner_doc_ref.set(entry) - cleanup(inner_doc_ref.delete) - - ids = [ - doc.id for doc in await db.collection_group("philosophers").recursive().get() + +async def _do_recursive_delete_with_bulk_writer(client, bulk_writer): + philosophers = [philosophers_data_set[0]] + await _persist_documents( + client, f"philosophers-async{UNIQUE_RESOURCE_ID}", philosophers + ) + + doc_paths = [ + "", + "/pets/Scruffy", + "/pets/Snowflake", + "/hobbies/pontificating", + "/hobbies/journaling", + "/philosophers/Aristotle", + "/philosophers/Plato", ] + # Assert all documents were created so that when they're missing after the + # delete, we're actually testing something. + collection_ref = client.collection(f"philosophers-async{UNIQUE_RESOURCE_ID}") + for path in doc_paths: + snapshot = await collection_ref.document(f"Socrates{path}").get() + assert snapshot.exists, f"Snapshot at Socrates{path} should have been created" + + # Now delete. + num_deleted = await client.recursive_delete(collection_ref, bulk_writer=bulk_writer) + assert num_deleted == len(doc_paths) + + # Now they should all be missing + for path in doc_paths: + snapshot = await collection_ref.document(f"Socrates{path}").get() + assert ( + not snapshot.exists + ), f"Snapshot at Socrates{path} should have been deleted" + + +async def test_async_recursive_delete_parallelized(client, cleanup): + from google.cloud.firestore_v1.bulk_writer import BulkWriterOptions, SendMode + + bw = client.bulk_writer(options=BulkWriterOptions(mode=SendMode.parallel)) + await _do_recursive_delete_with_bulk_writer(client, bw) + + +async def test_async_recursive_delete_serialized(client, cleanup): + from google.cloud.firestore_v1.bulk_writer import BulkWriterOptions, SendMode + + bw = client.bulk_writer(options=BulkWriterOptions(mode=SendMode.serial)) + await _do_recursive_delete_with_bulk_writer(client, bw) + + +async def test_recursive_query(client, cleanup): + col_id: str = f"philosophers-recursive-async-query{UNIQUE_RESOURCE_ID}" + await _persist_documents(client, col_id, philosophers_data_set, cleanup) + + ids = [doc.id for doc in await client.collection_group(col_id).recursive().get()] + expected_ids = [ # Aristotle doc and subdocs - f"Aristotle{UNIQUE_RESOURCE_ID}-async", + "Aristotle", "meditation", "questioning-stuff", "Doggy-Dog", "Floof-Boy", # Plato doc and subdocs - f"Plato{UNIQUE_RESOURCE_ID}-async", + "Plato", "abstraction", "hypotheticals", "Cuddles", "Sergeant-Puppers", # Socrates doc and subdocs - f"Socrates{UNIQUE_RESOURCE_ID}-async", + "Socrates", "journaling", "pontificating", "Scruffy", @@ -1173,36 +1266,12 @@ async def test_recursive_query(client, cleanup): async def test_nested_recursive_query(client, cleanup): + col_id: str = f"philosophers-nested-recursive-async-query{UNIQUE_RESOURCE_ID}" + await _persist_documents(client, col_id, philosophers_data_set, cleanup) - philosophers = [ - { - "data": {"name": "Aristotle", "favoriteCity": "Sparta"}, - "subcollections": { - "pets": [{"name": "Floof-Boy"}, {"name": "Doggy-Dog"}], - "hobbies": [{"name": "questioning-stuff"}, {"name": "meditation"}], - }, - }, - ] - - db = client - collection_ref = db.collection("philosophers") - for philosopher in philosophers: - ref = collection_ref.document( - f"{philosopher['data']['name']}{UNIQUE_RESOURCE_ID}-async" - ) - await ref.set(philosopher["data"]) - cleanup(ref.delete) - for col_name, entries in philosopher["subcollections"].items(): - sub_col = ref.collection(col_name) - for entry in entries: - inner_doc_ref = sub_col.document(entry["name"]) - await inner_doc_ref.set(entry) - cleanup(inner_doc_ref.delete) - - aristotle = collection_ref.document(f"Aristotle{UNIQUE_RESOURCE_ID}-async") - ids = [ - doc.id for doc in await aristotle.collection("pets")._query().recursive().get() - ] + collection_ref = client.collection(col_id) + aristotle = collection_ref.document("Aristotle") + ids = [doc.id for doc in await aristotle.collection("pets").recursive().get()] expected_ids = [ # Aristotle pets @@ -1219,6 +1288,84 @@ async def test_nested_recursive_query(client, cleanup): assert ids[index] == expected_ids[index], error_msg +async def test_chunked_query(client, cleanup): + col = client.collection(f"async-chunked-test{UNIQUE_RESOURCE_ID}") + for index in range(10): + doc_ref = col.document(f"document-{index + 1}") + await doc_ref.set({"index": index}) + cleanup(doc_ref.delete) + + lengths: List[int] = [len(chunk) async for chunk in col._chunkify(3)] + assert len(lengths) == 4 + assert lengths[0] == 3 + assert lengths[1] == 3 + assert lengths[2] == 3 + assert lengths[3] == 1 + + +async def test_chunked_query_smaller_limit(client, cleanup): + col = client.collection(f"chunked-test-smaller-limit{UNIQUE_RESOURCE_ID}") + for index in range(10): + doc_ref = col.document(f"document-{index + 1}") + await doc_ref.set({"index": index}) + cleanup(doc_ref.delete) + + lengths: List[int] = [len(chunk) async for chunk in col.limit(5)._chunkify(9)] + assert len(lengths) == 1 + assert lengths[0] == 5 + + +async def test_chunked_and_recursive(client, cleanup): + col_id = f"chunked-async-recursive-test{UNIQUE_RESOURCE_ID}" + documents = [ + { + "data": {"name": "Root-1"}, + "subcollections": { + "children": [ + {"data": {"name": f"Root-1--Child-{index + 1}"}} + for index in range(5) + ] + }, + }, + { + "data": {"name": "Root-2"}, + "subcollections": { + "children": [ + {"data": {"name": f"Root-2--Child-{index + 1}"}} + for index in range(5) + ] + }, + }, + ] + await _persist_documents(client, col_id, documents, cleanup) + collection_ref = client.collection(col_id) + iter = collection_ref.recursive()._chunkify(5) + + pages = [page async for page in iter] + doc_ids = [[doc.id for doc in page] for page in pages] + + page_1_ids = [ + "Root-1", + "Root-1--Child-1", + "Root-1--Child-2", + "Root-1--Child-3", + "Root-1--Child-4", + ] + assert doc_ids[0] == page_1_ids + + page_2_ids = [ + "Root-1--Child-5", + "Root-2", + "Root-2--Child-1", + "Root-2--Child-2", + "Root-2--Child-3", + ] + assert doc_ids[1] == page_2_ids + + page_3_ids = ["Root-2--Child-4", "Root-2--Child-5"] + assert doc_ids[2] == page_3_ids + + async def _chain(*iterators): """Asynchronous reimplementation of `itertools.chain`.""" for iterator in iterators: diff --git a/tests/unit/v1/test_async_client.py b/tests/unit/v1/test_async_client.py index bb7a51dd8..598da81ea 100644 --- a/tests/unit/v1/test_async_client.py +++ b/tests/unit/v1/test_async_client.py @@ -18,6 +18,8 @@ import aiounittest import mock +from google.cloud.firestore_v1.types.document import Document +from google.cloud.firestore_v1.types.firestore import RunQueryResponse from tests.unit.v1.test__helpers import AsyncIter, AsyncMock @@ -388,6 +390,110 @@ def test_sync_copy(self): # Multiple calls to this method should return the same cached instance. self.assertIs(client._to_sync_copy(), client._to_sync_copy()) + @pytest.mark.asyncio + async def test_recursive_delete(self): + client = self._make_default_one() + client._firestore_api_internal = AsyncMock(spec=["run_query"]) + collection_ref = client.collection("my_collection") + + results = [] + for index in range(10): + results.append( + RunQueryResponse(document=Document(name=f"{collection_ref.id}/{index}")) + ) + + chunks = [ + results[:3], + results[3:6], + results[6:9], + results[9:], + ] + + def _get_chunk(*args, **kwargs): + return AsyncIter(items=chunks.pop(0)) + + client._firestore_api_internal.run_query.side_effect = _get_chunk + + bulk_writer = mock.MagicMock() + bulk_writer.mock_add_spec(spec=["delete", "close"]) + + num_deleted = await client.recursive_delete( + collection_ref, bulk_writer=bulk_writer, chunk_size=3 + ) + self.assertEqual(num_deleted, len(results)) + + @pytest.mark.asyncio + async def test_recursive_delete_from_document(self): + client = self._make_default_one() + client._firestore_api_internal = mock.Mock( + spec=["run_query", "list_collection_ids"] + ) + collection_ref = client.collection("my_collection") + + collection_1_id: str = "collection_1_id" + collection_2_id: str = "collection_2_id" + + parent_doc = collection_ref.document("parent") + + collection_1_results = [] + collection_2_results = [] + + for index in range(10): + collection_1_results.append( + RunQueryResponse(document=Document(name=f"{collection_1_id}/{index}"),), + ) + + collection_2_results.append( + RunQueryResponse(document=Document(name=f"{collection_2_id}/{index}"),), + ) + + col_1_chunks = [ + collection_1_results[:3], + collection_1_results[3:6], + collection_1_results[6:9], + collection_1_results[9:], + ] + + col_2_chunks = [ + collection_2_results[:3], + collection_2_results[3:6], + collection_2_results[6:9], + collection_2_results[9:], + ] + + async def _get_chunk(*args, **kwargs): + start_at = ( + kwargs["request"]["structured_query"].start_at.values[0].reference_value + ) + + if collection_1_id in start_at: + return AsyncIter(col_1_chunks.pop(0)) + return AsyncIter(col_2_chunks.pop(0)) + + async def _get_collections(*args, **kwargs): + return AsyncIter([collection_1_id, collection_2_id]) + + client._firestore_api_internal.run_query.side_effect = _get_chunk + client._firestore_api_internal.list_collection_ids.side_effect = ( + _get_collections + ) + + bulk_writer = mock.MagicMock() + bulk_writer.mock_add_spec(spec=["delete", "close"]) + + num_deleted = await client.recursive_delete( + parent_doc, bulk_writer=bulk_writer, chunk_size=3 + ) + + expected_len = len(collection_1_results) + len(collection_2_results) + 1 + self.assertEqual(num_deleted, expected_len) + + @pytest.mark.asyncio + async def test_recursive_delete_raises(self): + client = self._make_default_one() + with self.assertRaises(TypeError): + await client.recursive_delete(object()) + def test_batch(self): from google.cloud.firestore_v1.async_batch import AsyncWriteBatch diff --git a/tests/unit/v1/test_async_collection.py b/tests/unit/v1/test_async_collection.py index 33006e254..1955ca52d 100644 --- a/tests/unit/v1/test_async_collection.py +++ b/tests/unit/v1/test_async_collection.py @@ -12,6 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +from google.cloud.firestore_v1.types.document import Document +from google.cloud.firestore_v1.types.firestore import RunQueryResponse import pytest import types import aiounittest @@ -204,6 +206,42 @@ async def test_add_w_retry_timeout(self): timeout = 123.0 await self._add_helper(retry=retry, timeout=timeout) + @pytest.mark.asyncio + async def test_chunkify(self): + client = _make_client() + col = client.collection("my-collection") + + client._firestore_api_internal = mock.Mock(spec=["run_query"]) + + results = [] + for index in range(10): + results.append( + RunQueryResponse( + document=Document( + name=f"projects/project-project/databases/(default)/documents/my-collection/{index}", + ), + ), + ) + + chunks = [ + results[:3], + results[3:6], + results[6:9], + results[9:], + ] + + async def _get_chunk(*args, **kwargs): + return AsyncIter(chunks.pop(0)) + + client._firestore_api_internal.run_query.side_effect = _get_chunk + + counter = 0 + expected_lengths = [3, 3, 3, 1] + async for chunk in col._chunkify(3): + msg = f"Expected chunk of length {expected_lengths[counter]} at index {counter}. Saw {len(chunk)}." + self.assertEqual(len(chunk), expected_lengths[counter], msg) + counter += 1 + @pytest.mark.asyncio async def _list_documents_helper(self, page_size=None, retry=None, timeout=None): from google.cloud.firestore_v1 import _helpers diff --git a/tests/unit/v1/test_async_query.py b/tests/unit/v1/test_async_query.py index 64feddaf4..4d18d551b 100644 --- a/tests/unit/v1/test_async_query.py +++ b/tests/unit/v1/test_async_query.py @@ -12,6 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +from google.cloud.firestore_v1.types.document import Document +from google.cloud.firestore_v1.types.firestore import RunQueryResponse import pytest import types import aiounittest @@ -469,6 +471,28 @@ async def test_stream_w_collection_group(self): metadata=client._rpc_metadata, ) + @pytest.mark.asyncio + async def test_unnecessary_chunkify(self): + client = _make_client() + + firestore_api = AsyncMock(spec=["run_query"]) + firestore_api.run_query.return_value = AsyncIter( + [ + RunQueryResponse( + document=Document( + name=f"projects/project-project/databases/(default)/documents/asdf/{index}", + ), + ) + for index in range(5) + ] + ) + client._firestore_api_internal = firestore_api + + query = client.collection("asdf")._query() + + async for chunk in query.limit(5)._chunkify(10): + self.assertEqual(len(chunk), 5) + class TestCollectionGroup(aiounittest.AsyncTestCase): @staticmethod diff --git a/tests/unit/v1/test_client.py b/tests/unit/v1/test_client.py index a46839ac5..5fbc73793 100644 --- a/tests/unit/v1/test_client.py +++ b/tests/unit/v1/test_client.py @@ -17,6 +17,8 @@ import unittest import mock +from google.cloud.firestore_v1.types.document import Document +from google.cloud.firestore_v1.types.firestore import RunQueryResponse class TestClient(unittest.TestCase): @@ -360,6 +362,104 @@ def test_get_all_unknown_result(self): metadata=client._rpc_metadata, ) + def test_recursive_delete(self): + client = self._make_default_one() + client._firestore_api_internal = mock.Mock(spec=["run_query"]) + collection_ref = client.collection("my_collection") + + results = [] + for index in range(10): + results.append( + RunQueryResponse(document=Document(name=f"{collection_ref.id}/{index}")) + ) + + chunks = [ + results[:3], + results[3:6], + results[6:9], + results[9:], + ] + + def _get_chunk(*args, **kwargs): + return iter(chunks.pop(0)) + + client._firestore_api_internal.run_query.side_effect = _get_chunk + + bulk_writer = mock.MagicMock() + bulk_writer.mock_add_spec(spec=["delete", "close"]) + + num_deleted = client.recursive_delete( + collection_ref, bulk_writer=bulk_writer, chunk_size=3 + ) + self.assertEqual(num_deleted, len(results)) + + def test_recursive_delete_from_document(self): + client = self._make_default_one() + client._firestore_api_internal = mock.Mock( + spec=["run_query", "list_collection_ids"] + ) + collection_ref = client.collection("my_collection") + + collection_1_id: str = "collection_1_id" + collection_2_id: str = "collection_2_id" + + parent_doc = collection_ref.document("parent") + + collection_1_results = [] + collection_2_results = [] + + for index in range(10): + collection_1_results.append( + RunQueryResponse(document=Document(name=f"{collection_1_id}/{index}"),), + ) + + collection_2_results.append( + RunQueryResponse(document=Document(name=f"{collection_2_id}/{index}"),), + ) + + col_1_chunks = [ + collection_1_results[:3], + collection_1_results[3:6], + collection_1_results[6:9], + collection_1_results[9:], + ] + + col_2_chunks = [ + collection_2_results[:3], + collection_2_results[3:6], + collection_2_results[6:9], + collection_2_results[9:], + ] + + def _get_chunk(*args, **kwargs): + start_at = ( + kwargs["request"]["structured_query"].start_at.values[0].reference_value + ) + + if collection_1_id in start_at: + return iter(col_1_chunks.pop(0)) + return iter(col_2_chunks.pop(0)) + + client._firestore_api_internal.run_query.side_effect = _get_chunk + client._firestore_api_internal.list_collection_ids.return_value = [ + collection_1_id, + collection_2_id, + ] + + bulk_writer = mock.MagicMock() + bulk_writer.mock_add_spec(spec=["delete", "close"]) + + num_deleted = client.recursive_delete( + parent_doc, bulk_writer=bulk_writer, chunk_size=3 + ) + + expected_len = len(collection_1_results) + len(collection_2_results) + 1 + self.assertEqual(num_deleted, expected_len) + + def test_recursive_delete_raises(self): + client = self._make_default_one() + self.assertRaises(TypeError, client.recursive_delete, object()) + def test_batch(self): from google.cloud.firestore_v1.batch import WriteBatch diff --git a/tests/unit/v1/test_collection.py b/tests/unit/v1/test_collection.py index 5885a29d9..cfefeb9e6 100644 --- a/tests/unit/v1/test_collection.py +++ b/tests/unit/v1/test_collection.py @@ -12,6 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +from google.cloud.firestore_v1.types.document import Document +from google.cloud.firestore_v1.types.firestore import RunQueryResponse import types import unittest @@ -355,3 +357,38 @@ def test_recursive(self): col = self._make_one("collection") self.assertIsInstance(col.recursive(), Query) + + def test_chunkify(self): + client = _test_helpers.make_client() + col = client.collection("my-collection") + + client._firestore_api_internal = mock.Mock(spec=["run_query"]) + + results = [] + for index in range(10): + results.append( + RunQueryResponse( + document=Document( + name=f"projects/project-project/databases/(default)/documents/my-collection/{index}", + ), + ), + ) + + chunks = [ + results[:3], + results[3:6], + results[6:9], + results[9:], + ] + + def _get_chunk(*args, **kwargs): + return iter(chunks.pop(0)) + + client._firestore_api_internal.run_query.side_effect = _get_chunk + + counter = 0 + expected_lengths = [3, 3, 3, 1] + for chunk in col._chunkify(3): + msg = f"Expected chunk of length {expected_lengths[counter]} at index {counter}. Saw {len(chunk)}." + self.assertEqual(len(chunk), expected_lengths[counter], msg) + counter += 1 diff --git a/tests/unit/v1/test_query.py b/tests/unit/v1/test_query.py index 91172b120..ea28969a8 100644 --- a/tests/unit/v1/test_query.py +++ b/tests/unit/v1/test_query.py @@ -12,6 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +from google.cloud.firestore_v1.types.document import Document +from google.cloud.firestore_v1.types.firestore import RunQueryResponse import types import unittest @@ -460,6 +462,40 @@ def test_on_snapshot(self, watch): query.on_snapshot(None) watch.for_query.assert_called_once() + def test_unnecessary_chunkify(self): + client = _make_client() + + firestore_api = mock.Mock(spec=["run_query"]) + firestore_api.run_query.return_value = iter( + [ + RunQueryResponse( + document=Document( + name=f"projects/project-project/databases/(default)/documents/asdf/{index}", + ), + ) + for index in range(5) + ] + ) + client._firestore_api_internal = firestore_api + + query = client.collection("asdf")._query() + + for chunk in query.limit(5)._chunkify(10): + self.assertEqual(len(chunk), 5) + + def test__resolve_chunk_size(self): + # With a global limit + query = _make_client().collection("asdf").limit(5) + self.assertEqual(query._resolve_chunk_size(3, 10), 2) + self.assertEqual(query._resolve_chunk_size(3, 1), 1) + self.assertEqual(query._resolve_chunk_size(3, 2), 2) + + # With no limit + query = _make_client().collection("asdf")._query() + self.assertEqual(query._resolve_chunk_size(3, 10), 10) + self.assertEqual(query._resolve_chunk_size(3, 1), 1) + self.assertEqual(query._resolve_chunk_size(3, 2), 2) + class TestCollectionGroup(unittest.TestCase): @staticmethod