feat: allow Query.fetch_page for queries with post filters (#463)

Chris Rossi · web-flow · commit 1111906887c7 · 2020-06-04T19:33:36.000-04:00
Closes #270
diff --git a/packages/google-cloud-ndb/google/cloud/ndb/_datastore_query.py b/packages/google-cloud-ndb/google/cloud/ndb/_datastore_query.py
@@ -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
diff --git a/packages/google-cloud-ndb/google/cloud/ndb/query.py b/packages/google-cloud-ndb/google/cloud/ndb/query.py
@@ -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
diff --git a/packages/google-cloud-ndb/tests/system/test_query.py b/packages/google-cloud-ndb/tests/system/test_query.py
@@ -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")
diff --git a/packages/google-cloud-ndb/tests/unit/test__datastore_query.py b/packages/google-cloud-ndb/tests/unit/test__datastore_query.py
@@ -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
diff --git a/packages/google-cloud-ndb/tests/unit/test_query.py b/packages/google-cloud-ndb/tests/unit/test_query.py
@@ -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")
