Skip to content
This repository was archived by the owner on May 6, 2026. It is now read-only.

fix: make optimized Query.count() work with the datastore emulator #528

Merged
chrisrossi merged 1 commit into from
Sep 2, 2020
Merged

fix: make optimized Query.count() work with the datastore emulator #528

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 15 additions & 4 deletions google/cloud/ndb/_datastore_query.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down
174 changes: 172 additions & 2 deletions tests/unit/test__datastore_query.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand All @@ -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",
Expand All @@ -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",),
Expand Down Expand Up @@ -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")
Expand Down