Skip to content

Commit d596c04

Browse files
tseavertseaver
authored
fix: handle empty last chunk correctly in 'Query._chunkify' (#489)
Closes #487. Supersedes #488.
1 parent 24eca62 commit d596c04

File tree

11 files changed

+298
-120
lines changed

11 files changed

+298
-120
lines changed

‎packages/google-cloud-firestore/google/cloud/firestore_v1/async_client.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -331,22 +331,22 @@ async def recursive_delete(
331331
The BulkWriter used to delete all matching documents. Supply this
332332
if you want to override the default throttling behavior.
333333
"""
334+
if bulk_writer is None:
335+
bulk_writer = self.bulk_writer()
336+
334337
return await self._recursive_delete(
335338
reference, bulk_writer=bulk_writer, chunk_size=chunk_size,
336339
)
337340

338341
async def _recursive_delete(
339342
self,
340343
reference: Union[AsyncCollectionReference, AsyncDocumentReference],
344+
bulk_writer: "BulkWriter",
341345
*,
342-
bulk_writer: Optional["BulkWriter"] = None, # type: ignore
343346
chunk_size: Optional[int] = 5000,
344347
depth: Optional[int] = 0,
345348
) -> int:
346349
"""Recursion helper for `recursive_delete."""
347-
from google.cloud.firestore_v1.bulk_writer import BulkWriter
348-
349-
bulk_writer = bulk_writer or BulkWriter()
350350

351351
num_deleted: int = 0
352352

‎packages/google-cloud-firestore/google/cloud/firestore_v1/async_query.py‎

Lines changed: 5 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -130,13 +130,6 @@ def __init__(
130130
async def _chunkify(
131131
self, chunk_size: int
132132
) -> AsyncGenerator[List[DocumentSnapshot], None]:
133-
# Catch the edge case where a developer writes the following:
134-
# `my_query.limit(500)._chunkify(1000)`, which ultimately nullifies any
135-
# need to yield chunks.
136-
if self._limit and chunk_size > self._limit:
137-
yield await self.get()
138-
return
139-
140133
max_to_return: Optional[int] = self._limit
141134
num_returned: int = 0
142135
original: AsyncQuery = self._copy()
@@ -150,11 +143,15 @@ async def _chunkify(
150143
# Apply the optionally pruned limit and the cursor, if we are past
151144
# the first page.
152145
_q = original.limit(_chunk_size)
146+
153147
if last_document:
154148
_q = _q.start_after(last_document)
155149

156150
snapshots = await _q.get()
157-
last_document = snapshots[-1]
151+
152+
if snapshots:
153+
last_document = snapshots[-1]
154+
158155
num_returned += len(snapshots)
159156

160157
yield snapshots

‎packages/google-cloud-firestore/google/cloud/firestore_v1/bulk_writer.py‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -252,7 +252,7 @@ class BulkWriter(AsyncBulkWriterMixin):
252252

253253
def __init__(
254254
self,
255-
client: Optional["BaseClient"] = None,
255+
client: "BaseClient" = None,
256256
options: Optional["BulkWriterOptions"] = None,
257257
):
258258
# Because `BulkWriter` instances are all synchronous/blocking on the
@@ -895,6 +895,11 @@ def __init__(
895895
self.mode = mode
896896
self.retry = retry
897897

898+
def __eq__(self, other):
899+
if not isinstance(other, self.__class__): # pragma: NO COVER
900+
return NotImplemented
901+
return self.__dict__ == other.__dict__
902+
898903
class BulkWriteFailure:
899904
def __init__(
900905
self,

‎packages/google-cloud-firestore/google/cloud/firestore_v1/client.py‎

Lines changed: 6 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,6 @@
5050
# Types needed only for Type Hints
5151
from google.cloud.firestore_v1.base_document import DocumentSnapshot
5252

53-
5453
if TYPE_CHECKING:
5554
from google.cloud.firestore_v1.bulk_writer import BulkWriter # pragma: NO COVER
5655

@@ -319,22 +318,20 @@ def recursive_delete(
319318
if you want to override the default throttling behavior.
320319
321320
"""
322-
return self._recursive_delete(
323-
reference, bulk_writer=bulk_writer, chunk_size=chunk_size,
324-
)
321+
if bulk_writer is None:
322+
bulk_writer or self.bulk_writer()
323+
324+
return self._recursive_delete(reference, bulk_writer, chunk_size=chunk_size,)
325325

326326
def _recursive_delete(
327327
self,
328328
reference: Union[CollectionReference, DocumentReference],
329+
bulk_writer: "BulkWriter",
329330
*,
330-
bulk_writer: Optional["BulkWriter"] = None,
331331
chunk_size: Optional[int] = 5000,
332332
depth: Optional[int] = 0,
333333
) -> int:
334334
"""Recursion helper for `recursive_delete."""
335-
from google.cloud.firestore_v1.bulk_writer import BulkWriter
336-
337-
bulk_writer = bulk_writer or BulkWriter()
338335

339336
num_deleted: int = 0
340337

@@ -354,10 +351,7 @@ def _recursive_delete(
354351
col_ref: CollectionReference
355352
for col_ref in reference.collections():
356353
num_deleted += self._recursive_delete(
357-
col_ref,
358-
bulk_writer=bulk_writer,
359-
chunk_size=chunk_size,
360-
depth=depth + 1,
354+
col_ref, bulk_writer, chunk_size=chunk_size, depth=depth + 1,
361355
)
362356
num_deleted += 1
363357
bulk_writer.delete(reference)

‎packages/google-cloud-firestore/google/cloud/firestore_v1/query.py‎

Lines changed: 5 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -171,12 +171,6 @@ def get(
171171
def _chunkify(
172172
self, chunk_size: int
173173
) -> Generator[List[DocumentSnapshot], None, None]:
174-
# Catch the edge case where a developer writes the following:
175-
# `my_query.limit(500)._chunkify(1000)`, which ultimately nullifies any
176-
# need to yield chunks.
177-
if self._limit and chunk_size > self._limit:
178-
yield self.get()
179-
return
180174

181175
max_to_return: Optional[int] = self._limit
182176
num_returned: int = 0
@@ -191,11 +185,15 @@ def _chunkify(
191185
# Apply the optionally pruned limit and the cursor, if we are past
192186
# the first page.
193187
_q = original.limit(_chunk_size)
188+
194189
if last_document:
195190
_q = _q.start_after(last_document)
196191

197192
snapshots = _q.get()
198-
last_document = snapshots[-1]
193+
194+
if snapshots:
195+
last_document = snapshots[-1]
196+
199197
num_returned += len(snapshots)
200198

201199
yield snapshots

‎packages/google-cloud-firestore/tests/system/test_system.py‎

Lines changed: 33 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1325,19 +1325,24 @@ def _persist_documents(
13251325
]
13261326

13271327

1328-
def _do_recursive_delete_with_bulk_writer(client, bulk_writer):
1329-
philosophers = [philosophers_data_set[0]]
1330-
_persist_documents(client, f"philosophers{UNIQUE_RESOURCE_ID}", philosophers)
1328+
def _do_recursive_delete(client, bulk_writer, empty_philosophers=False):
13311329

1332-
doc_paths = [
1333-
"",
1334-
"/pets/Scruffy",
1335-
"/pets/Snowflake",
1336-
"/hobbies/pontificating",
1337-
"/hobbies/journaling",
1338-
"/philosophers/Aristotle",
1339-
"/philosophers/Plato",
1340-
]
1330+
if empty_philosophers:
1331+
doc_paths = philosophers = []
1332+
else:
1333+
philosophers = [philosophers_data_set[0]]
1334+
1335+
doc_paths = [
1336+
"",
1337+
"/pets/Scruffy",
1338+
"/pets/Snowflake",
1339+
"/hobbies/pontificating",
1340+
"/hobbies/journaling",
1341+
"/philosophers/Aristotle",
1342+
"/philosophers/Plato",
1343+
]
1344+
1345+
_persist_documents(client, f"philosophers{UNIQUE_RESOURCE_ID}", philosophers)
13411346

13421347
# Assert all documents were created so that when they're missing after the
13431348
# delete, we're actually testing something.
@@ -1362,14 +1367,28 @@ def test_recursive_delete_parallelized(client, cleanup):
13621367
from google.cloud.firestore_v1.bulk_writer import BulkWriterOptions, SendMode
13631368

13641369
bw = client.bulk_writer(options=BulkWriterOptions(mode=SendMode.parallel))
1365-
_do_recursive_delete_with_bulk_writer(client, bw)
1370+
_do_recursive_delete(client, bw)
13661371

13671372

13681373
def test_recursive_delete_serialized(client, cleanup):
13691374
from google.cloud.firestore_v1.bulk_writer import BulkWriterOptions, SendMode
13701375

13711376
bw = client.bulk_writer(options=BulkWriterOptions(mode=SendMode.serial))
1372-
_do_recursive_delete_with_bulk_writer(client, bw)
1377+
_do_recursive_delete(client, bw)
1378+
1379+
1380+
def test_recursive_delete_parallelized_empty(client, cleanup):
1381+
from google.cloud.firestore_v1.bulk_writer import BulkWriterOptions, SendMode
1382+
1383+
bw = client.bulk_writer(options=BulkWriterOptions(mode=SendMode.parallel))
1384+
_do_recursive_delete(client, bw, empty_philosophers=True)
1385+
1386+
1387+
def test_recursive_delete_serialized_empty(client, cleanup):
1388+
from google.cloud.firestore_v1.bulk_writer import BulkWriterOptions, SendMode
1389+
1390+
bw = client.bulk_writer(options=BulkWriterOptions(mode=SendMode.serial))
1391+
_do_recursive_delete(client, bw, empty_philosophers=True)
13731392

13741393

13751394
def test_recursive_query(client, cleanup):

‎packages/google-cloud-firestore/tests/system/test_system_async.py‎

Lines changed: 32 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1184,22 +1184,26 @@ async def _persist_documents(
11841184
]
11851185

11861186

1187-
async def _do_recursive_delete_with_bulk_writer(client, bulk_writer):
1188-
philosophers = [philosophers_data_set[0]]
1187+
async def _do_recursive_delete(client, bulk_writer, empty_philosophers=False):
1188+
1189+
if empty_philosophers:
1190+
philosophers = doc_paths = []
1191+
else:
1192+
philosophers = [philosophers_data_set[0]]
1193+
doc_paths = [
1194+
"",
1195+
"/pets/Scruffy",
1196+
"/pets/Snowflake",
1197+
"/hobbies/pontificating",
1198+
"/hobbies/journaling",
1199+
"/philosophers/Aristotle",
1200+
"/philosophers/Plato",
1201+
]
1202+
11891203
await _persist_documents(
11901204
client, f"philosophers-async{UNIQUE_RESOURCE_ID}", philosophers
11911205
)
11921206

1193-
doc_paths = [
1194-
"",
1195-
"/pets/Scruffy",
1196-
"/pets/Snowflake",
1197-
"/hobbies/pontificating",
1198-
"/hobbies/journaling",
1199-
"/philosophers/Aristotle",
1200-
"/philosophers/Plato",
1201-
]
1202-
12031207
# Assert all documents were created so that when they're missing after the
12041208
# delete, we're actually testing something.
12051209
collection_ref = client.collection(f"philosophers-async{UNIQUE_RESOURCE_ID}")
@@ -1223,14 +1227,28 @@ async def test_async_recursive_delete_parallelized(client, cleanup):
12231227
from google.cloud.firestore_v1.bulk_writer import BulkWriterOptions, SendMode
12241228

12251229
bw = client.bulk_writer(options=BulkWriterOptions(mode=SendMode.parallel))
1226-
await _do_recursive_delete_with_bulk_writer(client, bw)
1230+
await _do_recursive_delete(client, bw)
12271231

12281232

12291233
async def test_async_recursive_delete_serialized(client, cleanup):
12301234
from google.cloud.firestore_v1.bulk_writer import BulkWriterOptions, SendMode
12311235

12321236
bw = client.bulk_writer(options=BulkWriterOptions(mode=SendMode.serial))
1233-
await _do_recursive_delete_with_bulk_writer(client, bw)
1237+
await _do_recursive_delete(client, bw)
1238+
1239+
1240+
async def test_async_recursive_delete_parallelized_empty(client, cleanup):
1241+
from google.cloud.firestore_v1.bulk_writer import BulkWriterOptions, SendMode
1242+
1243+
bw = client.bulk_writer(options=BulkWriterOptions(mode=SendMode.parallel))
1244+
await _do_recursive_delete(client, bw, empty_philosophers=True)
1245+
1246+
1247+
async def test_async_recursive_delete_serialized_empty(client, cleanup):
1248+
from google.cloud.firestore_v1.bulk_writer import BulkWriterOptions, SendMode
1249+
1250+
bw = client.bulk_writer(options=BulkWriterOptions(mode=SendMode.serial))
1251+
await _do_recursive_delete(client, bw, empty_philosophers=True)
12341252

12351253

12361254
async def test_recursive_query(client, cleanup):

0 commit comments

Comments
 (0)