def test_ancestor_with_property_filter():
key = key_module.Key("Foo", 123)
foo = model.StringProperty("foo")
query = query_module.QueryOptions(ancestor=key, filters=foo == "bar")
query_pb = _datastore_query._query_to_protobuf(query)
filter_pb = query_pb2.PropertyFilter(
property=query_pb2.PropertyReference(name="foo"),
op=query_pb2.PropertyFilter.EQUAL,
value=entity_pb2.Value(string_value="bar"),
)
ancestor_pb = query_pb2.PropertyFilter(
property=query_pb2.PropertyReference(name="__key__"),
op=query_pb2.PropertyFilter.HAS_ANCESTOR,
)
ancestor_pb.value.key_value.CopyFrom(key._key.to_protobuf())
expected_pb = query_pb2.Query(
filter=query_pb2.Filter(
composite_filter=query_pb2.CompositeFilter(
op=query_pb2.CompositeFilter.AND,
filters=[
query_pb2.Filter(property_filter=filter_pb),
query_pb2.Filter(property_filter=ancestor_pb),
],
)
)
)
assert query_pb == expected_pb
def test_distinct_on():
query = query_module.QueryOptions(distinct_on=("a", "b"))
expected_pb = query_pb2.Query(distinct_on=[
query_pb2.PropertyReference(name="a"),
query_pb2.PropertyReference(name="b"),
])
assert _datastore_query._query_to_protobuf(query) == expected_pb
def test_make_filter():
expected = query_pb2.PropertyFilter(
property=query_pb2.PropertyReference(name="harry"),
op=query_pb2.PropertyFilter.EQUAL,
value=entity_pb2.Value(string_value="Harold"),
)
assert _datastore_query.make_filter("harry", "=", u"Harold") == expected
def test_order_by():
query = query_module.QueryOptions(order_by=[
query_module.PropertyOrder("a"),
query_module.PropertyOrder("b", reverse=True),
])
expected_pb = query_pb2.Query(order=[
query_pb2.PropertyOrder(
property=query_pb2.PropertyReference(name="a"),
direction=query_pb2.PropertyOrder.ASCENDING,
),
query_pb2.PropertyOrder(
property=query_pb2.PropertyReference(name="b"),
direction=query_pb2.PropertyOrder.DESCENDING,
),
])
assert _datastore_query._query_to_protobuf(query) == expected_pb
def test_ancestor():
key = key_module.Key("Foo", 123)
query = query_module.QueryOptions(ancestor=key)
expected_pb = query_pb2.Query(filter=query_pb2.Filter(
property_filter=query_pb2.PropertyFilter(
property=query_pb2.PropertyReference(name="__key__"),
op=query_pb2.PropertyFilter.HAS_ANCESTOR,
)))
expected_pb.filter.property_filter.value.key_value.CopyFrom(
key._key.to_protobuf())
assert _datastore_query._query_to_protobuf(query) == expected_pb
def test_make_composite_and_filter():
filters = [
query_pb2.PropertyFilter(
property=query_pb2.PropertyReference(name="harry"),
op=query_pb2.PropertyFilter.EQUAL,
value=entity_pb2.Value(string_value="Harold"),
),
query_pb2.PropertyFilter(
property=query_pb2.PropertyReference(name="josie"),
op=query_pb2.PropertyFilter.EQUAL,
value=entity_pb2.Value(string_value="Josephine"),
),
]
expected = query_pb2.CompositeFilter(
op=query_pb2.CompositeFilter.AND,
filters=[
query_pb2.Filter(property_filter=sub_filter)
for sub_filter in filters
],
)
assert _datastore_query.make_composite_and_filter(filters) == expected
def test_filter_pb():
foo = model.StringProperty("foo")
query = query_module.QueryOptions(kind="Foo", filters=(foo == "bar"))
query_pb = _datastore_query._query_to_protobuf(query)
filter_pb = query_pb2.PropertyFilter(
property=query_pb2.PropertyReference(name="foo"),
op=query_pb2.PropertyFilter.EQUAL,
value=entity_pb2.Value(string_value="bar"),
)
expected_pb = query_pb2.Query(
kind=[query_pb2.KindExpression(name="Foo")],
filter=query_pb2.Filter(property_filter=filter_pb),
)
assert query_pb == expected_pb
def make_filter(name, op, value):
"""Make a property filter protocol buffer.
Args:
name (str): The name of the property to filter by.
op (str): The operator to apply in the filter. Must be one of "=", "<",
"<=", ">", or ">=".
value (Any): The value for comparison.
Returns:
query_pb2.PropertyFilter: The filter protocol buffer.
"""
filter_pb = query_pb2.PropertyFilter(
property=query_pb2.PropertyReference(name=name),
op=FILTER_OPERATORS[op],
)
helpers._set_protobuf_value(filter_pb.value, value)
return filter_pb
def _query_to_protobuf(query):
"""Convert an NDB query to a Datastore protocol buffer.
Args:
query (query.QueryOptions): The query spec.
Returns:
query_pb2.Query: The protocol buffer representation of the query.
"""
query_args = {}
if query.kind:
query_args["kind"] = [query_pb2.KindExpression(name=query.kind)]
if query.projection:
query_args["projection"] = [
query_pb2.Projection(property=query_pb2.PropertyReference(
name=name)) for name in query.projection
]
if query.distinct_on:
query_args["distinct_on"] = [
query_pb2.PropertyReference(name=name)
for name in query.distinct_on
]
if query.order_by:
query_args["order"] = [
query_pb2.PropertyOrder(
property=query_pb2.PropertyReference(name=order.name),
direction=DOWN if order.reverse else UP,
) for order in query.order_by
]
filter_pb = query.filters._to_filter() if query.filters else None
if query.ancestor:
ancestor_pb = query.ancestor._key.to_protobuf()
ancestor_filter_pb = query_pb2.PropertyFilter(
property=query_pb2.PropertyReference(name="__key__"),
op=query_pb2.PropertyFilter.HAS_ANCESTOR,
)
ancestor_filter_pb.value.key_value.CopyFrom(ancestor_pb)
if filter_pb is None:
filter_pb = ancestor_filter_pb
elif isinstance(filter_pb, query_pb2.CompositeFilter):
filter_pb.filters.add(property_filter=ancestor_filter_pb)
else:
filter_pb = query_pb2.CompositeFilter(
op=query_pb2.CompositeFilter.AND,
filters=[
_filter_pb(filter_pb),
_filter_pb(ancestor_filter_pb),
],
)
if filter_pb is not None:
query_args["filter"] = _filter_pb(filter_pb)
if query.start_cursor:
query_args["start_cursor"] = query.start_cursor.cursor
if query.end_cursor:
query_args["end_cursor"] = query.end_cursor.cursor
query_pb = query_pb2.Query(**query_args)
if query.offset:
query_pb.offset = query.offset
if query.limit:
query_pb.limit.value = query.limit
return query_pb
