diff --git a/google/cloud/ndb/_datastore_query.py b/google/cloud/ndb/_datastore_query.py index 961251b8..43916f15 100644 --- a/google/cloud/ndb/_datastore_query.py +++ b/google/cloud/ndb/_datastore_query.py @@ -169,7 +169,7 @@ def _count_by_skipping(query): # so for a workaround, just bail as soon as we neither skip nor retrieve any # results new_count = batch.skipped_results + len(batch.entity_results) - if new_count == 0: + if new_count == 0 and more_results != NOT_FINISHED: break count += new_count diff --git a/tests/unit/test__datastore_query.py b/tests/unit/test__datastore_query.py index 86d052db..cbe0c49e 100644 --- a/tests/unit/test__datastore_query.py +++ b/tests/unit/test__datastore_query.py @@ -163,6 +163,22 @@ def test_count_by_skipping_w_a_result(run_query): ), spec=("batch",), ), + mock.Mock( + batch=mock.Mock( + more_results=_datastore_query.NOT_FINISHED, + skipped_results=0, + entity_results=[], + end_cursor=b"secondCursor", + skipped_cursor=b"skiptomylou", + spec=( + "more_results", + "skipped_results", + "entity_results", + "end_cursor", + ), + ), + spec=("batch",), + ), mock.Mock( batch=mock.Mock( more_results=_datastore_query.NO_MORE_RESULTS, @@ -205,6 +221,17 @@ def test_count_by_skipping_w_a_result(run_query): ), {}, ), + ( + ( + query_module.QueryOptions( + limit=1, + offset=10000, + projection=["__key__"], + start_cursor=_datastore_query.Cursor(b"skiptomylou"), + ), + ), + {}, + ), ] assert run_query.call_args_list == expected