def test__next_page_w_skipped_lt_offset(self): from google.api_core import page_iterator from google.cloud.datastore_v1.types import datastore as datastore_pb2 from google.cloud.datastore_v1.types import entity as entity_pb2 from google.cloud.datastore_v1.types import query as query_pb2 from google.cloud.datastore.query import Query project = "prujekt" skipped_1 = 100 skipped_cursor_1 = b"DEADBEEF" skipped_2 = 50 skipped_cursor_2 = b"FACEDACE" more_enum = query_pb2.QueryResultBatch.MoreResultsType.NOT_FINISHED result_1 = _make_query_response([], b"", more_enum, skipped_1) result_1.batch.skipped_cursor = skipped_cursor_1 result_2 = _make_query_response([], b"", more_enum, skipped_2) result_2.batch.skipped_cursor = skipped_cursor_2 ds_api = _make_datastore_api(result_1, result_2) client = _Client(project, datastore_api=ds_api) query = Query(client) offset = 150 iterator = self._make_one(query, client, offset=offset) page = iterator._next_page() self.assertIsInstance(page, page_iterator.Page) self.assertIs(page._parent, iterator) partition_id = entity_pb2.PartitionId(project_id=project) read_options = datastore_pb2.ReadOptions() query_1 = query_pb2.Query(offset=offset) query_2 = query_pb2.Query(start_cursor=skipped_cursor_1, offset=(offset - skipped_1)) expected_calls = [ mock.call( request={ "project_id": project, "partition_id": partition_id, "read_options": read_options, "query": query, }) for query in [query_1, query_2] ] self.assertEqual(ds_api.run_query.call_args_list, expected_calls)
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.types import query as 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._pb.limit.value = limit - iterator.num_results self.assertEqual(pb, expected_pb)
def test__build_protobuf_empty(self): from google.cloud.datastore_v1.types import query as 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 _next_page_helper(self, txn_id=None, retry=None, timeout=None): from google.api_core import page_iterator from google.cloud.datastore_v1.types import datastore as datastore_pb2 from google.cloud.datastore_v1.types import entity as entity_pb2 from google.cloud.datastore_v1.types import query as query_pb2 from google.cloud.datastore.query import Query more_enum = query_pb2.QueryResultBatch.MoreResultsType.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( request={ "project_id": project, "partition_id": partition_id, "read_options": read_options, "query": empty_query, }, **kwargs, )
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.types import query as 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._pb.limit.value = limit - iterator.num_results self.assertEqual(pb, expected_pb)
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: projection = query_pb2.Projection() projection.property.name = projection_name pb.projection.append(projection) if query.kind: kind = query_pb2.KindExpression() kind.name = query.kind pb.kind.append(kind) composite_filter = pb.filter.composite_filter composite_filter.op = query_pb2.CompositeFilter.Operator.AND if query.ancestor: ancestor_pb = query.ancestor.to_protobuf() # Filter on __key__ HAS_ANCESTOR == ancestor. ancestor_filter = composite_filter.filters._pb.add().property_filter ancestor_filter.property.name = "__key__" ancestor_filter.op = query_pb2.PropertyFilter.Operator.HAS_ANCESTOR ancestor_filter.value.key_value.CopyFrom(ancestor_pb._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._pb.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._pb) else: helpers._set_protobuf_value(property_filter.value, value) if not composite_filter.filters: pb._pb.ClearField("filter") for prop in query.order: property_order = query_pb2.PropertyOrder() if prop.startswith("-"): property_order.property.name = prop[1:] property_order.direction = property_order.Direction.DESCENDING else: property_order.property.name = prop property_order.direction = property_order.Direction.ASCENDING pb.order.append(property_order) for distinct_on_name in query.distinct_on: ref = query_pb2.PropertyReference() ref.name = distinct_on_name pb.distinct_on.append(ref) return pb
def _next_page(self): """Get the next page in the iterator. :rtype: :class:`~google.cloud.iterator.Page` :returns: The next page in the iterator (or :data:`None` if there are no pages left). """ if not self._more_results: return None query_pb = self._build_protobuf() transaction = self.client.current_transaction if transaction is None: transaction_id = None else: transaction_id = transaction.id read_options = helpers.get_read_options(self._eventual, transaction_id) partition_id = entity_pb2.PartitionId( project_id=self._query.project, namespace_id=self._query.namespace) kwargs = {} if self._retry is not None: kwargs["retry"] = self._retry if self._timeout is not None: kwargs["timeout"] = self._timeout response_pb = self.client._datastore_api.run_query( request={ "project_id": self._query.project, "partition_id": partition_id, "read_options": read_options, "query": query_pb, }, **kwargs, ) while (response_pb.batch.more_results == _NOT_FINISHED and response_pb.batch.skipped_results < query_pb.offset): # We haven't finished processing. A likely reason is we haven't # skipped all of the results yet. Don't return any results. # Instead, rerun query, adjusting offsets. Datastore doesn't process # more than 1000 skipped results in a query. old_query_pb = query_pb query_pb = query_pb2.Query() query_pb._pb.CopyFrom(old_query_pb._pb) # copy for testability query_pb.start_cursor = response_pb.batch.skipped_cursor query_pb.offset -= response_pb.batch.skipped_results response_pb = self.client._datastore_api.run_query( request={ "project_id": self._query.project, "partition_id": partition_id, "read_options": read_options, "query": query_pb, }, **kwargs, ) entity_pbs = self._process_query_results(response_pb) return page_iterator.Page(self, entity_pbs, self.item_to_value)
def _make_query_pb(kind): from google.cloud.datastore_v1.types import query as query_pb2 return query_pb2.Query(kind=[query_pb2.KindExpression(name=kind)])