diff --git a/google/cloud/ndb/_datastore_query.py b/google/cloud/ndb/_datastore_query.py index 3490c8e7..961251b8 100644 --- a/google/cloud/ndb/_datastore_query.py +++ b/google/cloud/ndb/_datastore_query.py @@ -165,14 +165,25 @@ def _count_by_skipping(query): response = yield _datastore_run_query(query) batch = response.batch - more_results = batch.more_results - count += batch.skipped_results - count += len(batch.entity_results) + # The Datastore emulator will never set more_results to NO_MORE_RESULTS, + # 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: + break + count += new_count if limit and count >= limit: break - cursor = Cursor(batch.end_cursor) + # The Datastore emulator won't set end_cursor to something useful if no results + # are returned, so the workaround is to use skipped_cursor in that case + if len(batch.entity_results): + cursor = Cursor(batch.end_cursor) + else: + cursor = Cursor(batch.skipped_cursor) + + more_results = batch.more_results raise tasklets.Return(count) diff --git a/tests/unit/test__datastore_query.py b/tests/unit/test__datastore_query.py index 59984f39..cadabf85 100644 --- a/tests/unit/test__datastore_query.py +++ b/tests/unit/test__datastore_query.py @@ -141,6 +141,73 @@ def next(self): raw=True, ) + @staticmethod + @pytest.mark.usefixtures("in_context") + @mock.patch("google.cloud.ndb._datastore_query._datastore_run_query") + def test_count_by_skipping_w_a_result(run_query): + # These results should technically be impossible, but better safe than sorry. + run_query.side_effect = utils.future_results( + mock.Mock( + batch=mock.Mock( + more_results=_datastore_query.NOT_FINISHED, + skipped_results=1000, + entity_results=[], + end_cursor=b"dontlookatme", + skipped_cursor=b"himom", + spec=( + "more_results", + "skipped_results", + "entity_results", + "end_cursor", + ), + ), + spec=("batch",), + ), + mock.Mock( + batch=mock.Mock( + more_results=_datastore_query.NO_MORE_RESULTS, + skipped_results=99, + entity_results=[object()], + end_cursor=b"ohhaithere", + skipped_cursor=b"hellodad", + spec=( + "more_results", + "skipped_results", + "entity_results", + "end_cursor", + "skipped_cursor", + ), + ), + spec=("batch",), + ), + ) + + query = query_module.QueryOptions() + future = _datastore_query.count(query) + assert future.result() == 1100 + + expected = [ + mock.call( + query_module.QueryOptions( + limit=1, + offset=10000, + projection=["__key__"], + ) + ), + ( + ( + query_module.QueryOptions( + limit=1, + offset=10000, + projection=["__key__"], + start_cursor=_datastore_query.Cursor(b"himom"), + ), + ), + {}, + ), + ] + assert run_query.call_args_list == expected + @staticmethod @pytest.mark.usefixtures("in_context") @mock.patch("google.cloud.ndb._datastore_query._datastore_run_query") @@ -151,7 +218,8 @@ def test_count_by_skipping(run_query): more_results=_datastore_query.NOT_FINISHED, skipped_results=1000, entity_results=[], - end_cursor=b"himom", + end_cursor=b"dontlookatme", + skipped_cursor=b"himom", spec=( "more_results", "skipped_results", @@ -166,12 +234,14 @@ def test_count_by_skipping(run_query): more_results=_datastore_query.NO_MORE_RESULTS, skipped_results=100, entity_results=[], - end_cursor=b"hellodad", + end_cursor=b"nopenuhuh", + skipped_cursor=b"hellodad", spec=( "more_results", "skipped_results", "entity_results", "end_cursor", + "skipped_cursor", ), ), spec=("batch",), @@ -204,6 +274,106 @@ def test_count_by_skipping(run_query): ] assert run_query.call_args_list == expected + @staticmethod + @pytest.mark.usefixtures("in_context") + @mock.patch("google.cloud.ndb._datastore_query._datastore_run_query") + def test_count_by_skipping_emulator(run_query): + """Regression test for #525 + + Test differences between emulator and the real Datastore. + + https://github.com/googleapis/python-ndb/issues/525 + """ + run_query.side_effect = utils.future_results( + mock.Mock( + batch=mock.Mock( + more_results=_datastore_query.MORE_RESULTS_AFTER_LIMIT, + skipped_results=1000, + entity_results=[], + end_cursor=b"dontlookatme", + skipped_cursor=b"himom", + spec=( + "more_results", + "skipped_results", + "entity_results", + "end_cursor", + ), + ), + spec=("batch",), + ), + mock.Mock( + batch=mock.Mock( + more_results=_datastore_query.MORE_RESULTS_AFTER_LIMIT, + skipped_results=100, + entity_results=[], + end_cursor=b"nopenuhuh", + skipped_cursor=b"hellodad", + spec=( + "more_results", + "skipped_results", + "entity_results", + "end_cursor", + "skipped_cursor", + ), + ), + spec=("batch",), + ), + mock.Mock( + batch=mock.Mock( + more_results=_datastore_query.MORE_RESULTS_AFTER_LIMIT, + skipped_results=0, + entity_results=[], + end_cursor=b"nopenuhuh", + skipped_cursor=b"hellodad", + spec=( + "more_results", + "skipped_results", + "entity_results", + "end_cursor", + "skipped_cursor", + ), + ), + spec=("batch",), + ), + ) + + query = query_module.QueryOptions() + future = _datastore_query.count(query) + assert future.result() == 1100 + + expected = [ + mock.call( + query_module.QueryOptions( + limit=1, + offset=10000, + projection=["__key__"], + ) + ), + ( + ( + query_module.QueryOptions( + limit=1, + offset=10000, + projection=["__key__"], + start_cursor=_datastore_query.Cursor(b"himom"), + ), + ), + {}, + ), + ( + ( + query_module.QueryOptions( + limit=1, + offset=10000, + projection=["__key__"], + start_cursor=_datastore_query.Cursor(b"hellodad"), + ), + ), + {}, + ), + ] + assert run_query.call_args_list == expected + @staticmethod @pytest.mark.usefixtures("in_context") @mock.patch("google.cloud.ndb._datastore_query._datastore_run_query")