Skip to content

Commit

Permalink
feat: support IN/NOT_IN/NOT_EQUAL operators
Browse files Browse the repository at this point in the history
  • Loading branch information
Vincent-Weng committed Jun 8, 2022
1 parent db044b1 commit bca829b
Show file tree
Hide file tree
Showing 2 changed files with 14 additions and 5 deletions.
11 changes: 7 additions & 4 deletions google/cloud/datastore/query.py
Expand Up @@ -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."""

Expand Down Expand Up @@ -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
Expand All @@ -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`,
Expand All @@ -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):
Expand Down Expand Up @@ -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)
Expand Down
8 changes: 7 additions & 1 deletion tests/unit/test_query.py
Expand Up @@ -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():
Expand Down

0 comments on commit bca829b

Please sign in to comment.