def test__build_protobuf_all_values_except_offset(self):
# this test and the following (all_values_except_start_and_end_cursor)
# test mutually exclusive states; the offset is ignored
# if a start_cursor is supplied
from google.cloud.datastore_v1.proto import query_pb2
from google.cloud.datastore.query import Query
client = _Client(None)
query = Query(client)
limit = 15
start_bytes = b"i\xb7\x1d"
start_cursor = "abcd"
end_bytes = b"\xc3\x1c\xb3"
end_cursor = "wxyz"
iterator = self._make_one(query,
client,
limit=limit,
start_cursor=start_cursor,
end_cursor=end_cursor)
self.assertEqual(iterator.max_results, limit)
iterator.num_results = 4
iterator._skipped_results = 1
pb = iterator._build_protobuf()
expected_pb = query_pb2.Query(start_cursor=start_bytes,
end_cursor=end_bytes)
expected_pb.limit.value = limit - iterator.num_results
self.assertEqual(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_ancestor_with_composite_filter():
key = key_module.Key("Foo", 123)
foo = model.StringProperty("foo")
food = model.StringProperty("food")
query = query_module.QueryOptions(
ancestor=key,
filters=query_module.AND(foo == "bar", food == "barn"),
)
query_pb = _datastore_query._query_to_protobuf(query)
filter_pb1 = query_pb2.PropertyFilter(
property=query_pb2.PropertyReference(name="foo"),
op=query_pb2.PropertyFilter.EQUAL,
value=entity_pb2.Value(string_value="bar"),
)
filter_pb2 = query_pb2.PropertyFilter(
property=query_pb2.PropertyReference(name="food"),
op=query_pb2.PropertyFilter.EQUAL,
value=entity_pb2.Value(string_value="barn"),
)
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_pb1),
query_pb2.Filter(property_filter=filter_pb2),
query_pb2.Filter(property_filter=ancestor_pb),
],
)))
assert query_pb == expected_pb
def test__build_protobuf_all_values_except_start_and_end_cursor(self):
# this test and the previous (all_values_except_start_offset)
# test mutually exclusive states; the offset is ignored
# if a start_cursor is supplied
from google.cloud.datastore_v1.proto import query_pb2
from google.cloud.datastore.query import Query
client = _Client(None)
query = Query(client)
limit = 15
offset = 9
iterator = self._make_one(
query,
client,
limit=limit,
offset=offset,
)
self.assertEqual(iterator.max_results, limit)
iterator.num_results = 4
pb = iterator._build_protobuf()
expected_pb = query_pb2.Query(offset=offset -
iterator._skipped_results, )
expected_pb.limit.value = limit - iterator.num_results
self.assertEqual(pb, expected_pb)
def test__build_protobuf_empty(self):
from google.cloud.datastore_v1.proto import query_pb2
from google.cloud.datastore.query import Query
client = _Client(None)
query = Query(client)
iterator = self._make_one(query, client)
pb = iterator._build_protobuf()
expected_pb = query_pb2.Query()
self.assertEqual(pb, 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_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 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 _next_page_helper(self, txn_id=None, retry=None, timeout=None):
from google.api_core import page_iterator
from google.cloud.datastore_v1.proto import datastore_pb2
from google.cloud.datastore_v1.proto import entity_pb2
from google.cloud.datastore_v1.proto import query_pb2
from google.cloud.datastore.query import Query
more_enum = query_pb2.QueryResultBatch.NOT_FINISHED
result = _make_query_response([], b"", more_enum, 0)
project = "prujekt"
ds_api = _make_datastore_api(result)
if txn_id is None:
client = _Client(project, datastore_api=ds_api)
else:
transaction = mock.Mock(id=txn_id, spec=["id"])
client = _Client(project,
datastore_api=ds_api,
transaction=transaction)
query = Query(client)
kwargs = {}
if retry is not None:
kwargs["retry"] = retry
if timeout is not None:
kwargs["timeout"] = timeout
iterator = self._make_one(query, client, **kwargs)
page = iterator._next_page()
self.assertIsInstance(page, page_iterator.Page)
self.assertIs(page._parent, iterator)
partition_id = entity_pb2.PartitionId(project_id=project)
if txn_id is None:
read_options = datastore_pb2.ReadOptions()
else:
read_options = datastore_pb2.ReadOptions(transaction=txn_id)
empty_query = query_pb2.Query()
ds_api.run_query.assert_called_once_with(project,
partition_id,
read_options,
query=empty_query,
**kwargs)
def test_limit():
query = query_module.QueryOptions(limit=20)
expected_pb = query_pb2.Query()
expected_pb.limit.value = 20
assert _datastore_query._query_to_protobuf(query) == expected_pb
def test_offset():
query = query_module.QueryOptions(offset=20)
assert _datastore_query._query_to_protobuf(query) == query_pb2.Query(
offset=20
)
def test_kind():
query = query_module.QueryOptions(kind="Foo")
assert _datastore_query._query_to_protobuf(query) == query_pb2.Query(
kind=[query_pb2.KindExpression(name="Foo")]
)
def test_no_args():
query = query_module.QueryOptions()
assert _datastore_query._query_to_protobuf(query) == query_pb2.Query()
def _pb_from_query(query):
"""Convert a Query instance to the corresponding protobuf.
:type query: :class:`Query`
:param query: The source query.
:rtype: :class:`.query_pb2.Query`
:returns: A protobuf that can be sent to the protobuf API. N.b. that
it does not contain "in-flight" fields for ongoing query
executions (cursors, offset, limit).
"""
pb = query_pb2.Query()
for projection_name in query.projection:
pb.projection.add().property.name = projection_name
if query.kind:
pb.kind.add().name = query.kind
composite_filter = pb.filter.composite_filter
composite_filter.op = query_pb2.CompositeFilter.AND
if query.ancestor:
ancestor_pb = query.ancestor.to_protobuf()
# Filter on __key__ HAS_ANCESTOR == ancestor.
ancestor_filter = composite_filter.filters.add().property_filter
ancestor_filter.property.name = "__key__"
ancestor_filter.op = query_pb2.PropertyFilter.HAS_ANCESTOR
ancestor_filter.value.key_value.CopyFrom(ancestor_pb)
for property_name, operator, value in query.filters:
pb_op_enum = query.OPERATORS.get(operator)
# Add the specific filter
property_filter = composite_filter.filters.add().property_filter
property_filter.property.name = property_name
property_filter.op = pb_op_enum
# Set the value to filter on based on the type.
if property_name == "__key__":
key_pb = value.to_protobuf()
property_filter.value.key_value.CopyFrom(key_pb)
else:
helpers._set_protobuf_value(property_filter.value, value)
if not composite_filter.filters:
pb.ClearField("filter")
for prop in query.order:
property_order = pb.order.add()
if prop.startswith("-"):
property_order.property.name = prop[1:]
property_order.direction = property_order.DESCENDING
else:
property_order.property.name = prop
property_order.direction = property_order.ASCENDING
for distinct_on_name in query.distinct_on:
pb.distinct_on.add().name = distinct_on_name
return 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
def test_end_cursor():
query = query_module.QueryOptions(
end_cursor=_datastore_query.Cursor(b"abc"))
assert _datastore_query._query_to_protobuf(query) == query_pb2.Query(
end_cursor=b"abc")
def _make_query_pb(kind):
from google.cloud.datastore_v1.proto import query_pb2
return query_pb2.Query(kind=[query_pb2.KindExpression(name=kind)])
