Skip to content

Commit 9a99eb8

Browse files
author
Chris Rossi
authored
fix: make optimized Query.count() work with the datastore emulator (#528)
The emulator is different enough from the real Datastore to require some special handling. Fixes #525
1 parent 0803d13 commit 9a99eb8

File tree

2 files changed

+187
-6
lines changed

2 files changed

+187
-6
lines changed

‎packages/google-cloud-ndb/google/cloud/ndb/_datastore_query.py‎

Lines changed: 15 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -165,14 +165,25 @@ def _count_by_skipping(query):
165165
response = yield _datastore_run_query(query)
166166
batch = response.batch
167167

168-
more_results = batch.more_results
169-
count += batch.skipped_results
170-
count += len(batch.entity_results)
168+
# The Datastore emulator will never set more_results to NO_MORE_RESULTS,
169+
# so for a workaround, just bail as soon as we neither skip nor retrieve any
170+
# results
171+
new_count = batch.skipped_results + len(batch.entity_results)
172+
if new_count == 0:
173+
break
171174

175+
count += new_count
172176
if limit and count >= limit:
173177
break
174178

175-
cursor = Cursor(batch.end_cursor)
179+
# The Datastore emulator won't set end_cursor to something useful if no results
180+
# are returned, so the workaround is to use skipped_cursor in that case
181+
if len(batch.entity_results):
182+
cursor = Cursor(batch.end_cursor)
183+
else:
184+
cursor = Cursor(batch.skipped_cursor)
185+
186+
more_results = batch.more_results
176187

177188
raise tasklets.Return(count)
178189

‎packages/google-cloud-ndb/tests/unit/test__datastore_query.py‎

Lines changed: 172 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -141,6 +141,73 @@ def next(self):
141141
raw=True,
142142
)
143143

144+
@staticmethod
145+
@pytest.mark.usefixtures("in_context")
146+
@mock.patch("google.cloud.ndb._datastore_query._datastore_run_query")
147+
def test_count_by_skipping_w_a_result(run_query):
148+
# These results should technically be impossible, but better safe than sorry.
149+
run_query.side_effect = utils.future_results(
150+
mock.Mock(
151+
batch=mock.Mock(
152+
more_results=_datastore_query.NOT_FINISHED,
153+
skipped_results=1000,
154+
entity_results=[],
155+
end_cursor=b"dontlookatme",
156+
skipped_cursor=b"himom",
157+
spec=(
158+
"more_results",
159+
"skipped_results",
160+
"entity_results",
161+
"end_cursor",
162+
),
163+
),
164+
spec=("batch",),
165+
),
166+
mock.Mock(
167+
batch=mock.Mock(
168+
more_results=_datastore_query.NO_MORE_RESULTS,
169+
skipped_results=99,
170+
entity_results=[object()],
171+
end_cursor=b"ohhaithere",
172+
skipped_cursor=b"hellodad",
173+
spec=(
174+
"more_results",
175+
"skipped_results",
176+
"entity_results",
177+
"end_cursor",
178+
"skipped_cursor",
179+
),
180+
),
181+
spec=("batch",),
182+
),
183+
)
184+
185+
query = query_module.QueryOptions()
186+
future = _datastore_query.count(query)
187+
assert future.result() == 1100
188+
189+
expected = [
190+
mock.call(
191+
query_module.QueryOptions(
192+
limit=1,
193+
offset=10000,
194+
projection=["__key__"],
195+
)
196+
),
197+
(
198+
(
199+
query_module.QueryOptions(
200+
limit=1,
201+
offset=10000,
202+
projection=["__key__"],
203+
start_cursor=_datastore_query.Cursor(b"himom"),
204+
),
205+
),
206+
{},
207+
),
208+
]
209+
assert run_query.call_args_list == expected
210+
144211
@staticmethod
145212
@pytest.mark.usefixtures("in_context")
146213
@mock.patch("google.cloud.ndb._datastore_query._datastore_run_query")
@@ -151,7 +218,8 @@ def test_count_by_skipping(run_query):
151218
more_results=_datastore_query.NOT_FINISHED,
152219
skipped_results=1000,
153220
entity_results=[],
154-
end_cursor=b"himom",
221+
end_cursor=b"dontlookatme",
222+
skipped_cursor=b"himom",
155223
spec=(
156224
"more_results",
157225
"skipped_results",
@@ -166,12 +234,14 @@ def test_count_by_skipping(run_query):
166234
more_results=_datastore_query.NO_MORE_RESULTS,
167235
skipped_results=100,
168236
entity_results=[],
169-
end_cursor=b"hellodad",
237+
end_cursor=b"nopenuhuh",
238+
skipped_cursor=b"hellodad",
170239
spec=(
171240
"more_results",
172241
"skipped_results",
173242
"entity_results",
174243
"end_cursor",
244+
"skipped_cursor",
175245
),
176246
),
177247
spec=("batch",),
@@ -204,6 +274,106 @@ def test_count_by_skipping(run_query):
204274
]
205275
assert run_query.call_args_list == expected
206276

277+
@staticmethod
278+
@pytest.mark.usefixtures("in_context")
279+
@mock.patch("google.cloud.ndb._datastore_query._datastore_run_query")
280+
def test_count_by_skipping_emulator(run_query):
281+
"""Regression test for #525
282+
283+
Test differences between emulator and the real Datastore.
284+
285+
https://github.com/googleapis/python-ndb/issues/525
286+
"""
287+
run_query.side_effect = utils.future_results(
288+
mock.Mock(
289+
batch=mock.Mock(
290+
more_results=_datastore_query.MORE_RESULTS_AFTER_LIMIT,
291+
skipped_results=1000,
292+
entity_results=[],
293+
end_cursor=b"dontlookatme",
294+
skipped_cursor=b"himom",
295+
spec=(
296+
"more_results",
297+
"skipped_results",
298+
"entity_results",
299+
"end_cursor",
300+
),
301+
),
302+
spec=("batch",),
303+
),
304+
mock.Mock(
305+
batch=mock.Mock(
306+
more_results=_datastore_query.MORE_RESULTS_AFTER_LIMIT,
307+
skipped_results=100,
308+
entity_results=[],
309+
end_cursor=b"nopenuhuh",
310+
skipped_cursor=b"hellodad",
311+
spec=(
312+
"more_results",
313+
"skipped_results",
314+
"entity_results",
315+
"end_cursor",
316+
"skipped_cursor",
317+
),
318+
),
319+
spec=("batch",),
320+
),
321+
mock.Mock(
322+
batch=mock.Mock(
323+
more_results=_datastore_query.MORE_RESULTS_AFTER_LIMIT,
324+
skipped_results=0,
325+
entity_results=[],
326+
end_cursor=b"nopenuhuh",
327+
skipped_cursor=b"hellodad",
328+
spec=(
329+
"more_results",
330+
"skipped_results",
331+
"entity_results",
332+
"end_cursor",
333+
"skipped_cursor",
334+
),
335+
),
336+
spec=("batch",),
337+
),
338+
)
339+
340+
query = query_module.QueryOptions()
341+
future = _datastore_query.count(query)
342+
assert future.result() == 1100
343+
344+
expected = [
345+
mock.call(
346+
query_module.QueryOptions(
347+
limit=1,
348+
offset=10000,
349+
projection=["__key__"],
350+
)
351+
),
352+
(
353+
(
354+
query_module.QueryOptions(
355+
limit=1,
356+
offset=10000,
357+
projection=["__key__"],
358+
start_cursor=_datastore_query.Cursor(b"himom"),
359+
),
360+
),
361+
{},
362+
),
363+
(
364+
(
365+
query_module.QueryOptions(
366+
limit=1,
367+
offset=10000,
368+
projection=["__key__"],
369+
start_cursor=_datastore_query.Cursor(b"hellodad"),
370+
),
371+
),
372+
{},
373+
),
374+
]
375+
assert run_query.call_args_list == expected
376+
207377
@staticmethod
208378
@pytest.mark.usefixtures("in_context")
209379
@mock.patch("google.cloud.ndb._datastore_query._datastore_run_query")

0 commit comments

Comments
 (0)