From 2646cef3e2687461174a11c45f29de7b84d1fcdb Mon Sep 17 00:00:00 2001 From: Jim Morrison Date: Wed, 28 Feb 2024 09:14:57 -0800 Subject: [PATCH] feat: Allow queries using server side IN. (#954) * feat: Allow queries using server side IN. * Rename force_server to server_op. --- google/cloud/ndb/_datastore_query.py | 2 +- google/cloud/ndb/model.py | 4 ++-- google/cloud/ndb/query.py | 18 ++++----------- tests/system/test_query.py | 34 ++++++++++++++++++++++++++++ tests/unit/test_model.py | 29 +++++++++++++++++++++++- tests/unit/test_query.py | 7 ------ 6 files changed, 69 insertions(+), 25 deletions(-) diff --git a/google/cloud/ndb/_datastore_query.py b/google/cloud/ndb/_datastore_query.py index 553c8285..480a2a68 100644 --- a/google/cloud/ndb/_datastore_query.py +++ b/google/cloud/ndb/_datastore_query.py @@ -57,7 +57,7 @@ ">": query_pb2.PropertyFilter.Operator.GREATER_THAN, ">=": query_pb2.PropertyFilter.Operator.GREATER_THAN_OR_EQUAL, "!=": query_pb2.PropertyFilter.Operator.NOT_EQUAL, - "IN": query_pb2.PropertyFilter.Operator.IN, + "in": query_pb2.PropertyFilter.Operator.IN, } _KEY_NOT_IN_CACHE = object() diff --git a/google/cloud/ndb/model.py b/google/cloud/ndb/model.py index 224c6deb..b43d4163 100644 --- a/google/cloud/ndb/model.py +++ b/google/cloud/ndb/model.py @@ -1258,7 +1258,7 @@ def __ge__(self, value): """FilterNode: Represents the ``>=`` comparison.""" return self._comparison(">=", value) - def _IN(self, value): + def _IN(self, value, server_op=False): """For the ``in`` comparison operator. The ``in`` operator cannot be overloaded in the way we want @@ -1315,7 +1315,7 @@ def _IN(self, value): sub_value = self._datastore_type(sub_value) values.append(sub_value) - return query.FilterNode(self._name, "in", values) + return query.FilterNode(self._name, "in", values, server_op=server_op) IN = _IN """Used to check if a property value is contained in a set of values. diff --git a/google/cloud/ndb/query.py b/google/cloud/ndb/query.py index 7fa46706..6109fe11 100644 --- a/google/cloud/ndb/query.py +++ b/google/cloud/ndb/query.py @@ -619,6 +619,7 @@ class FilterNode(Node): opsymbol (str): The comparison operator. One of ``=``, ``!=``, ``<``, ``<=``, ``>``, ``>=`` or ``in``. value (Any): The value to filter on / relative to. + server_op (bool): Force the operator to use a server side filter. Raises: TypeError: If ``opsymbol`` is ``"in"`` but ``value`` is not a @@ -630,7 +631,7 @@ class FilterNode(Node): _opsymbol = None _value = None - def __new__(cls, name, opsymbol, value): + def __new__(cls, name, opsymbol, value, server_op=False): # Avoid circular import in Python 2.7 from google.cloud.ndb import model @@ -648,7 +649,8 @@ def __new__(cls, name, opsymbol, value): return FalseNode() if len(nodes) == 1: return nodes[0] - return DisjunctionNode(*nodes) + if not server_op: + return DisjunctionNode(*nodes) instance = super(FilterNode, cls).__new__(cls) instance._name = name @@ -695,24 +697,12 @@ def _to_filter(self, post=False): Optional[query_pb2.PropertyFilter]: Returns :data:`None`, if this is a post-filter, otherwise returns the protocol buffer representation of the filter. - - Raises: - NotImplementedError: If the ``opsymbol`` is ``in``, since - they should correspond to a composite filter. This should - never occur since the constructor will create ``OR`` nodes for - ``in`` """ # Avoid circular import in Python 2.7 from google.cloud.ndb import _datastore_query if post: return None - if self._opsymbol in (_IN_OP): - raise NotImplementedError( - "Inequality filters are not single filter " - "expressions and therefore cannot be converted " - "to a single filter ({!r})".format(self._opsymbol) - ) return _datastore_query.make_filter(self._name, self._opsymbol, self._value) diff --git a/tests/system/test_query.py b/tests/system/test_query.py index df00a6b6..fb2e9bbb 100644 --- a/tests/system/test_query.py +++ b/tests/system/test_query.py @@ -865,6 +865,40 @@ def make_entities(): assert not more +@pytest.mark.usefixtures("client_context") +def test_fetch_page_in_query(dispose_of): + page_size = 5 + n_entities = page_size * 2 + + class SomeKind(ndb.Model): + foo = ndb.IntegerProperty() + + @ndb.toplevel + def make_entities(): + entities = [SomeKind(foo=n_entities) for i in range(n_entities)] + keys = yield [entity.put_async() for entity in entities] + raise ndb.Return(keys) + + for key in make_entities(): + dispose_of(key._key) + + query = SomeKind.query().filter(SomeKind.foo.IN([1, 2, n_entities], server_op=True)) + eventually(query.fetch, length_equals(n_entities)) + + results, cursor, more = query.fetch_page(page_size) + assert len(results) == page_size + assert more + + safe_cursor = cursor.urlsafe() + next_cursor = ndb.Cursor(urlsafe=safe_cursor) + results, cursor, more = query.fetch_page(page_size, start_cursor=next_cursor) + assert len(results) == page_size + + results, cursor, more = query.fetch_page(page_size, start_cursor=cursor) + assert not results + assert not more + + @pytest.mark.usefixtures("client_context") def test_polymodel_query(ds_entity): class Animal(ndb.PolyModel): diff --git a/tests/unit/test_model.py b/tests/unit/test_model.py index 3250d22d..82e4324c 100644 --- a/tests/unit/test_model.py +++ b/tests/unit/test_model.py @@ -549,7 +549,7 @@ def test__IN_wrong_container(): assert model.Property._FIND_METHODS_CACHE == {} @staticmethod - def test__IN(): + def test__IN_default(): prop = model.Property("name", indexed=True) or_node = prop._IN(["a", None, "xy"]) expected = query_module.DisjunctionNode( @@ -561,6 +561,33 @@ def test__IN(): # Also verify the alias assert or_node == prop.IN(["a", None, "xy"]) + @staticmethod + def test__IN_client(): + prop = model.Property("name", indexed=True) + or_node = prop._IN(["a", None, "xy"], server_op=False) + expected = query_module.DisjunctionNode( + query_module.FilterNode("name", "=", "a"), + query_module.FilterNode("name", "=", None), + query_module.FilterNode("name", "=", "xy"), + ) + assert or_node == expected + # Also verify the alias + assert or_node == prop.IN(["a", None, "xy"]) + + @staticmethod + def test_server__IN(): + prop = model.Property("name", indexed=True) + in_node = prop._IN(["a", None, "xy"], server_op=True) + assert in_node == prop.IN(["a", None, "xy"], server_op=True) + assert in_node != query_module.DisjunctionNode( + query_module.FilterNode("name", "=", "a"), + query_module.FilterNode("name", "=", None), + query_module.FilterNode("name", "=", "xy"), + ) + assert in_node == query_module.FilterNode( + "name", "in", ["a", None, "xy"], server_op=True + ) + @staticmethod def test___neg__(): prop = model.Property("name") diff --git a/tests/unit/test_query.py b/tests/unit/test_query.py index 589c9bcc..13da4740 100644 --- a/tests/unit/test_query.py +++ b/tests/unit/test_query.py @@ -701,13 +701,6 @@ def test__to_ne_filter_op(): filter_node = query_module.FilterNode("speed", "!=", 88) assert filter_node._to_filter(post=True) is None - @staticmethod - def test__to_filter_bad_op(): - filter_node = query_module.FilterNode("speed", ">=", 88) - filter_node._opsymbol = "in" - with pytest.raises(NotImplementedError): - filter_node._to_filter() - @staticmethod @mock.patch("google.cloud.ndb._datastore_query") def test__to_filter(_datastore_query):