Skip to content

Commit

Permalink
feat: allow Query.fetch_page for queries with post filters
Browse files Browse the repository at this point in the history
  • Loading branch information
Chris Rossi committed Jun 4, 2020
1 parent 5d354e4 commit 1d6492e
Show file tree
Hide file tree
Showing 5 changed files with 49 additions and 51 deletions.
4 changes: 4 additions & 0 deletions google/cloud/ndb/_datastore_query.py
Original file line number Diff line number Diff line change
Expand Up @@ -494,6 +494,10 @@ def cursor_after(self):

return self._cursor_after

@property
def _more_results_after_limit(self):
return self._result_set._more_results_after_limit


class _MultiQueryIteratorImpl(QueryIterator):
"""Multiple Query Iterator
Expand Down
9 changes: 0 additions & 9 deletions google/cloud/ndb/query.py
Original file line number Diff line number Diff line change
Expand Up @@ -2319,15 +2319,6 @@ def fetch_page_async(self, page_size, **kwargs):
"that uses 'OR', '!=', or 'IN'."
)

post_filters = _options.filters._post_filters()
if post_filters:
raise TypeError(
"Can't use 'fetch_page' or 'fetch_page_async' with a "
"post-filter. (An in-memory filter.) This probably means "
"you're querying a repeated structured property which "
"requires post-filtering."
)

iterator = _datastore_query.iterate(_options, raw=True)
results = []
cursor = None
Expand Down
54 changes: 22 additions & 32 deletions tests/system/test_query.py
Original file line number Diff line number Diff line change
Expand Up @@ -1516,59 +1516,49 @@ def test_fetch_page_with_repeated_structured_property(dispose_of):
class OtherKind(ndb.Model):
one = ndb.StringProperty()
two = ndb.StringProperty()
three = ndb.StringProperty()
three = ndb.IntegerProperty()

class SomeKind(ndb.Model):
foo = ndb.IntegerProperty()
bar = ndb.StructuredProperty(OtherKind, repeated=True)

N = 30

@ndb.synctasklet
def make_entities():
entity1 = SomeKind(
foo=1,
bar=[
OtherKind(one="pish", two="posh", three="pash"),
OtherKind(one="bish", two="bosh", three="bash"),
],
)
entity2 = SomeKind(
foo=2,
bar=[
OtherKind(one="bish", two="bosh", three="bass"),
OtherKind(one="pish", two="posh", three="pass"),
],
)
entity3 = SomeKind(
foo=3,
bar=[
OtherKind(one="pish", two="fosh", three="fash"),
OtherKind(one="bish", two="posh", three="bash"),
],
)

keys = yield (
entity1.put_async(),
entity2.put_async(),
entity3.put_async(),
)
futures = [
SomeKind(
foo=i,
bar=[
OtherKind(one="pish", two="posh", three=i % 2),
OtherKind(one="bish", two="bosh", three=i % 2),
],
).put_async()
for i in range(N)
]

keys = yield futures
raise ndb.Return(keys)

keys = make_entities()
for key in keys:
dispose_of(key._key)

eventually(SomeKind.query().fetch, length_equals(3))
eventually(SomeKind.query().fetch, length_equals(N))
query = (
SomeKind.query()
.filter(
SomeKind.bar == OtherKind(one="pish", two="posh"),
SomeKind.bar == OtherKind(two="posh", three="pash"),
SomeKind.bar == OtherKind(two="bosh", three=0),
)
.order(SomeKind.foo)
)

with pytest.raises(TypeError):
query.fetch_page(page_size=10)
results, cursor, more = query.fetch_page(page_size=5)
assert [entity.foo for entity in results] == [0, 2, 4, 6, 8]

results, cursor, more = query.fetch_page(page_size=5, start_cursor=cursor)
assert [entity.foo for entity in results] == [10, 12, 14, 16, 18]


@pytest.mark.usefixtures("client_context")
Expand Down
23 changes: 23 additions & 0 deletions tests/unit/test__datastore_query.py
Original file line number Diff line number Diff line change
Expand Up @@ -680,6 +680,29 @@ def test_cursor_after_no_cursor():
with pytest.raises(exceptions.BadArgumentError):
iterator.cursor_after()

@staticmethod
def test__more_results_after_limit():
foo = model.StringProperty("foo")
query = query_module.QueryOptions(
offset=20, limit=10, filters=foo == u"this"
)
predicate = object()
iterator = _datastore_query._PostFilterQueryIteratorImpl(
query, predicate
)
assert iterator._result_set._query == query_module.QueryOptions(
filters=foo == u"this"
)
assert iterator._offset == 20
assert iterator._limit == 10
assert iterator._predicate is predicate

iterator._result_set._more_results_after_limit = False
assert iterator._more_results_after_limit is False

iterator._result_set._more_results_after_limit = True
assert iterator._more_results_after_limit is True


class Test_MultiQueryIteratorImpl:
@staticmethod
Expand Down
10 changes: 0 additions & 10 deletions tests/unit/test_query.py
Original file line number Diff line number Diff line change
Expand Up @@ -2198,16 +2198,6 @@ def test_fetch_page_multiquery():
with pytest.raises(TypeError):
query.fetch_page(5)

@staticmethod
@pytest.mark.usefixtures("in_context")
def test_fetch_page_post_filter():
query = query_module.Query()
query.filters = mock.Mock(
_multiquery=False, _post_filters=mock.Mock(return_value=True)
)
with pytest.raises(TypeError):
query.fetch_page(5)

@staticmethod
@pytest.mark.usefixtures("in_context")
@mock.patch("google.cloud.ndb._datastore_query")
Expand Down

0 comments on commit 1d6492e

Please sign in to comment.