From bca829b61e43a038dcb6ade3ff8b780043fe785e Mon Sep 17 00:00:00 2001 From: Han Weng Date: Wed, 16 Mar 2022 15:56:37 -0400 Subject: [PATCH] feat: support IN/NOT_IN/NOT_EQUAL operators --- google/cloud/datastore/query.py | 11 +++++++---- tests/unit/test_query.py | 8 +++++++- 2 files changed, 14 insertions(+), 5 deletions(-) diff --git a/google/cloud/datastore/query.py b/google/cloud/datastore/query.py index 5e4f4937..57c19205 100644 --- a/google/cloud/datastore/query.py +++ b/google/cloud/datastore/query.py @@ -86,6 +86,9 @@ class Query(object): "<": query_pb2.PropertyFilter.Operator.LESS_THAN, ">": query_pb2.PropertyFilter.Operator.GREATER_THAN, "=": query_pb2.PropertyFilter.Operator.EQUAL, + "!=": query_pb2.PropertyFilter.Operator.NOT_EQUAL, + "IN": query_pb2.PropertyFilter.Operator.IN, + "NOT_IN": query_pb2.PropertyFilter.Operator.NOT_IN, } """Mapping of operator strings and their protobuf equivalents.""" @@ -215,7 +218,7 @@ def add_filter(self, property_name, operator, value): where property is a property stored on the entity in the datastore and operator is one of ``OPERATORS`` - (ie, ``=``, ``<``, ``<=``, ``>``, ``>=``): + (ie, ``=``, ``<``, ``<=``, ``>``, ``>=``, ``!=``, ``IN``, ``NOT_IN``): .. testsetup:: query-filter @@ -235,7 +238,7 @@ def add_filter(self, property_name, operator, value): :param property_name: A property name. :type operator: str - :param operator: One of ``=``, ``<``, ``<=``, ``>``, ``>=``. + :param operator: One of ``=``, ``<``, ``<=``, ``>``, ``>=``, ``!=``, ``IN``, ``NOT_IN``. :type value: :class:`int`, :class:`str`, :class:`bool`, :class:`float`, :class:`NoneType`, @@ -252,7 +255,7 @@ def add_filter(self, property_name, operator, value): """ if self.OPERATORS.get(operator) is None: error_message = 'Invalid expression: "%s"' % (operator,) - choices_message = "Please use one of: =, <, <=, >, >=." + choices_message = "Please use one of: =, <, <=, >, >=, !=, IN, NOT_IN." raise ValueError(error_message, choices_message) if property_name == "__key__" and not isinstance(value, Key): @@ -293,7 +296,7 @@ def key_filter(self, key, operator="="): :param key: The key to filter on. :type operator: str - :param operator: (Optional) One of ``=``, ``<``, ``<=``, ``>``, ``>=``. + :param operator: (Optional) One of ``=``, ``<``, ``<=``, ``>``, ``>=``, ``!=``, ``IN``, ``NOT_IN``. Defaults to ``=``. """ self.add_filter("__key__", operator, key) diff --git a/tests/unit/test_query.py b/tests/unit/test_query.py index 3cbd95b8..1f250f46 100644 --- a/tests/unit/test_query.py +++ b/tests/unit/test_query.py @@ -175,12 +175,18 @@ def test_query_add_filter_w_all_operators(): query.add_filter("lt_prop", "<", "val3") query.add_filter("gt_prop", ">", "val4") query.add_filter("eq_prop", "=", "val5") - assert len(query.filters) == 5 + query.add_filter("in_prop", "IN", ["val6"]) + query.add_filter("neq_prop", "!=", "val9") + query.add_filter("not_in_prop", "NOT_IN", ["val13"]) + assert len(query.filters) == 8 assert query.filters[0] == ("leq_prop", "<=", "val1") assert query.filters[1] == ("geq_prop", ">=", "val2") assert query.filters[2] == ("lt_prop", "<", "val3") assert query.filters[3] == ("gt_prop", ">", "val4") assert query.filters[4] == ("eq_prop", "=", "val5") + assert query.filters[5] == ("in_prop", "IN", ["val6"]) + assert query.filters[6] == ("neq_prop", "!=", "val9") + assert query.filters[7] == ("not_in_prop", "NOT_IN", ["val13"]) def test_query_add_filter_w_known_operator_and_entity():