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)])