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) response_pb = self.client._datastore_api.run_query( self._query.project, partition_id, read_options, query=query_pb, ) entity_pbs = self._process_query_results(response_pb) return page_iterator.Page(self, entity_pbs, self.item_to_value)
def _datastore_run_query(query): """Run a query in Datastore. Args: query (query.QueryOptions): The query spec. Returns: tasklets.Future: """ query_pb = _query_to_protobuf(query) partition_id = entity_pb2.PartitionId(project_id=query.project, namespace_id=query.namespace) read_options = _datastore_api.get_read_options( query, default_read_consistency=_datastore_api.EVENTUAL) request = datastore_pb2.RunQueryRequest( project_id=query.project, partition_id=partition_id, query=query_pb, read_options=read_options, ) response = yield _datastore_api.make_call("RunQuery", request, timeout=query.timeout) log.debug(response) return response
def _next_page_helper(self, txn_id=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) iterator = self._make_one(query, client) 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)
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(self._query.project, partition_id, 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. 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( self._query.project, partition_id, 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 test_entity_key_only(): key_pb = entity_pb2.Key( partition_id=entity_pb2.PartitionId(project_id="testing"), path=[entity_pb2.Key.PathElement(kind="ThisKind", id=42)], ) result = _datastore_query._Result( _datastore_query.RESULT_TYPE_KEY_ONLY, mock.Mock( entity=mock.Mock(key=key_pb, spec=("key",)), cursor=b"123", spec=("entity", "cursor"), ), ) assert result.entity() == key_module.Key("ThisKind", 42)
def test_entity_full_entity(model): key_pb = entity_pb2.Key( partition_id=entity_pb2.PartitionId(project_id="testing"), path=[entity_pb2.Key.PathElement(kind="ThisKind", id=42)], ) entity = mock.Mock(key=key_pb) model._entity_from_protobuf.return_value = entity result = _datastore_query._Result( _datastore_query.RESULT_TYPE_FULL, mock.Mock(entity=entity, cursor=b"123", spec=("entity", "cursor")), ) assert result.entity() is entity model._entity_from_protobuf.assert_called_once_with(entity)
def test_entity_full_entity_no_cache(model): context = context_module.get_context() with context.new(cache_policy=False).use(): key_pb = entity_pb2.Key( partition_id=entity_pb2.PartitionId(project_id="testing"), path=[entity_pb2.Key.PathElement(kind="ThisKind", id=42)], ) entity = mock.Mock(key=key_pb) model._entity_from_protobuf.return_value = entity result = _datastore_query._Result( _datastore_query.RESULT_TYPE_FULL, mock.Mock(entity=entity, cursor=b"123", spec=("entity", "cursor")), ) assert result.entity() is entity
def test_run_query_w_eventual_no_transaction(self): 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 project = "PROJECT" kind = "Nonesuch" cursor = b"\x00" query_pb = self._make_query_pb(kind) partition_id = entity_pb2.PartitionId(project_id=project) read_options = datastore_pb2.ReadOptions( read_consistency=datastore_pb2.ReadOptions.EVENTUAL ) rsp_pb = datastore_pb2.RunQueryResponse( batch=query_pb2.QueryResultBatch( entity_result_type=query_pb2.EntityResult.FULL, end_cursor=cursor, more_results=query_pb2.QueryResultBatch.NO_MORE_RESULTS, ) ) # Create mock HTTP and client with response. http = _make_requests_session( [_make_response(content=rsp_pb.SerializeToString())] ) client_info = _make_client_info() client = mock.Mock( _http=http, _base_url="test.invalid", _client_info=client_info, spec=["_http", "_base_url", "_client_info"], ) # Make request. ds_api = self._make_one(client) response = ds_api.run_query(project, partition_id, read_options, query=query_pb) # Check the result and verify the callers. self.assertEqual(response, rsp_pb) uri = _build_expected_url(client._base_url, project, "runQuery") request = _verify_protobuf_call(http, uri, datastore_pb2.RunQueryRequest()) self.assertEqual(request.partition_id, partition_id) self.assertEqual(request.query, query_pb) self.assertEqual(request.read_options, read_options)
def test_entity_full_entity_cached(model): key = key_module.Key("ThisKind", 42) key_pb = entity_pb2.Key( partition_id=entity_pb2.PartitionId(project_id="testing"), path=[entity_pb2.Key.PathElement(kind="ThisKind", id=42)], ) entity = mock.Mock(key=key_pb) cached_entity = mock.Mock(key=key_pb, _key=key) context = context_module.get_context() context.cache[key] = cached_entity model._entity_from_protobuf.return_value = entity result = _datastore_query._Result( _datastore_query.RESULT_TYPE_FULL, mock.Mock(entity=entity, cursor=b"123", spec=("entity", "cursor")), ) assert result.entity() is not entity assert result.entity() is cached_entity
def test_it(_datastore_api): query = query_module.QueryOptions(project="testing", namespace="") query_pb = _datastore_query._query_to_protobuf(query) _datastore_api.make_call.return_value = utils.future_result("foo") read_options = datastore_pb2.ReadOptions() request = datastore_pb2.RunQueryRequest( project_id="testing", partition_id=entity_pb2.PartitionId(project_id="testing", namespace_id=""), query=query_pb, read_options=read_options, ) _datastore_api.get_read_options.return_value = read_options assert _datastore_query._datastore_run_query(query).result() == "foo" _datastore_api.make_call.assert_called_once_with("RunQuery", request, timeout=None) _datastore_api.get_read_options.assert_called_once_with( query, default_read_consistency=_datastore_api.EVENTUAL)
def test_run_query_w_namespace_nonempty_result(self): 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 project = "PROJECT" kind = "Kind" namespace = "NS" query_pb = self._make_query_pb(kind) partition_id = entity_pb2.PartitionId(project_id=project, namespace_id=namespace) read_options = datastore_pb2.ReadOptions() rsp_pb = datastore_pb2.RunQueryResponse( batch=query_pb2.QueryResultBatch( entity_result_type=query_pb2.EntityResult.FULL, entity_results=[ query_pb2.EntityResult(entity=entity_pb2.Entity()) ], more_results=query_pb2.QueryResultBatch.NO_MORE_RESULTS, )) # Create mock HTTP and client with response. http = _make_requests_session( [_make_response(content=rsp_pb.SerializeToString())]) client = mock.Mock(_http=http, _base_url="test.invalid", spec=["_http", "_base_url"]) # Make request. ds_api = self._make_one(client) response = ds_api.run_query(project, partition_id, read_options, query=query_pb) # Check the result and verify the callers. self.assertEqual(response, rsp_pb) uri = _build_expected_url(client._base_url, project, "runQuery") request = _verify_protobuf_call(http, uri, datastore_pb2.RunQueryRequest()) self.assertEqual(request.partition_id, partition_id) self.assertEqual(request.query, query_pb)
def test_run_query_wo_eventual_w_transaction(self): 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 project = 'PROJECT' kind = 'Nonesuch' cursor = b'\x00' transaction = b'TRANSACTION' query_pb = self._make_query_pb(kind) partition_id = entity_pb2.PartitionId(project_id=project) read_options = datastore_pb2.ReadOptions(transaction=transaction) rsp_pb = datastore_pb2.RunQueryResponse( batch=query_pb2.QueryResultBatch( entity_result_type=query_pb2.EntityResult.FULL, end_cursor=cursor, more_results=query_pb2.QueryResultBatch.NO_MORE_RESULTS, ) ) # Create mock HTTP and client with response. http = _make_requests_session( [_make_response(content=rsp_pb.SerializeToString())]) client = mock.Mock( _http=http, _base_url='test.invalid', spec=['_http', '_base_url']) # Make request. ds_api = self._make_one(client) response = ds_api.run_query( project, partition_id, read_options, query=query_pb) # Check the result and verify the callers. self.assertEqual(response, rsp_pb) uri = _build_expected_url(client._base_url, project, 'runQuery') request = _verify_protobuf_call( http, uri, datastore_pb2.RunQueryRequest()) self.assertEqual(request.partition_id, partition_id) self.assertEqual(request.query, query_pb) self.assertEqual(request.read_options, read_options)
def _run_query(project_id, namespace, query_pb): """Run a query in Datastore. Will potentially repeat the query to get all results. Args: project_id (str): The project/app id of the Datastore instance. namespace (str): The namespace to which to restrict results. query_pb (query_pb2.Query): The query protocol buffer representation. Returns: tasklets.Future: List[Tuple[query_pb2.EntityResult.ResultType, query_pb2.EntityResult]]: The raw query results. """ results = [] partition_id = entity_pb2.PartitionId(project_id=project_id, namespace_id=namespace) while True: # See what results we get from the backend request = datastore_pb2.RunQueryRequest(project_id=project_id, partition_id=partition_id, query=query_pb) response = yield _datastore_api.make_call("RunQuery", request) log.debug(response) batch = response.batch results.extend(((batch.entity_result_type, result) for result in batch.entity_results)) # Did we get all of them? if batch.more_results != MORE_RESULTS_TYPE_NOT_FINISHED: break # Still some results left to fetch. Update cursors and try again. query_pb.start_cursor = batch.end_cursor return results
def _datastore_run_query(query): """Run a query in Datastore. Args: query (query.QueryOptions): The query spec. Returns: tasklets.Future: """ query_pb = _query_to_protobuf(query) partition_id = entity_pb2.PartitionId(project_id=query.project, namespace_id=query.namespace) read_options = _datastore_api.get_read_options(query) request = datastore_pb2.RunQueryRequest( project_id=query.project, partition_id=partition_id, query=query_pb, read_options=read_options, ) response = yield _datastore_api.make_call("RunQuery", request, timeout=query.timeout) utils.logging_debug(log, response) raise tasklets.Return(response)