def _set_result_from_operation(self): """Set the result or exception from the operation if it is complete.""" # This must be done in a lock to prevent the polling thread # and main thread from both executing the completion logic # at the same time. with self._completion_lock: # If the operation isn't complete or if the result has already been # set, do not call set_result/set_exception again. # Note: self._result_set is set to True in set_result and # set_exception, in case those methods are invoked directly. if not self._operation.done or self._result_set: return if self._operation.HasField("response"): response = protobuf_helpers.from_any_pb( self._result_type, self._operation.response) self.set_result(response) elif self._operation.HasField("error"): exception = exceptions.from_grpc_status( status_code=self._operation.error.code, message=self._operation.error.message, errors=(self._operation.error, ), response=self._operation, ) self.set_exception(exception) else: # Some APIs set `done: true`, with an empty response. # Set the result to an empty message of the expected # result type. # https://google.aip.dev/151 self.set_result(self._result_type())
def _set_result_from_operation(self): """Set the result or exception from the operation if it is complete.""" # This must be done in a lock to prevent the polling thread # and main thread from both executing the completion logic # at the same time. with self._completion_lock: # If the operation isn't complete or if the result has already been # set, do not call set_result/set_exception again. # Note: self._result_set is set to True in set_result and # set_exception, in case those methods are invoked directly. if not self._operation.done or self._result_set: return if self._operation.HasField("response"): response = protobuf_helpers.from_any_pb( self._result_type, self._operation.response) self.set_result(response) elif self._operation.HasField("error"): exception = exceptions.from_grpc_status( status_code=self._operation.error.code, message=self._operation.error.message, errors=(self._operation.error, ), response=self._operation, ) self.set_exception(exception) else: exception = exceptions.GoogleAPICallError( "Unexpected state: Long-running operation had neither " "response nor error set.") self.set_exception(exception)
def test_from_grpc_status(): message = "message" exception = exceptions.from_grpc_status(grpc.StatusCode.OUT_OF_RANGE, message) assert isinstance(exception, exceptions.BadRequest) assert isinstance(exception, exceptions.OutOfRange) assert exception.code == http_client.BAD_REQUEST assert exception.grpc_status_code == grpc.StatusCode.OUT_OF_RANGE assert exception.message == message assert exception.errors == []
def test_from_grpc_status_as_int(): message = "message" exception = exceptions.from_grpc_status(11, message) assert isinstance(exception, exceptions.BadRequest) assert isinstance(exception, exceptions.OutOfRange) assert exception.code == http.client.BAD_REQUEST assert exception.grpc_status_code == grpc.StatusCode.OUT_OF_RANGE assert exception.message == message assert exception.errors == []
def _do_mutate_retryable_rows(self): """Mutate all the rows that are eligible for retry. A row is eligible for retry if it has not been tried or if it resulted in a transient error in a previous call. :rtype: list :return: The responses statuses, which is a list of :class:`~google.rpc.status_pb2.Status`. :raises: One of the following: * :exc:`~google.api_core.exceptions.ServiceUnavailable` if any row returned a transient error. This is "artificial" in the sense that we intentionally raise the error because it will be caught by the retry strategy. * :exc:`RuntimeError` if the number of responses doesn't match the number of rows that were retried """ retryable_rows = [] index_into_all_rows = [] for index, status in enumerate(self.responses_statuses): if self._is_retryable(status): retryable_rows.append(self.rows[index]) index_into_all_rows.append(index) if not retryable_rows: # All mutations are either successful or non-retryable now. return self.responses_statuses mutate_rows_request = _mutate_rows_request(self.table_name, retryable_rows) responses = self.client._data_stub.MutateRows(mutate_rows_request) num_responses = 0 num_retryable_responses = 0 for response in responses: for entry in response.entries: num_responses += 1 index = index_into_all_rows[entry.index] self.responses_statuses[index] = entry.status if self._is_retryable(entry.status): num_retryable_responses += 1 if entry.status.code == 0: self.rows[index].clear() if len(retryable_rows) != num_responses: raise RuntimeError('Unexpected the number of responses', num_responses, 'Expected', len(retryable_rows)) if num_retryable_responses: raise from_grpc_status(StatusCode.UNAVAILABLE, 'MutateRows retryable error.') return self.responses_statuses
def test_from_grpc_status_with_errors_and_response(): message = "message" response = mock.sentinel.response errors = ["1", "2"] exception = exceptions.from_grpc_status( grpc.StatusCode.OUT_OF_RANGE, message, errors=errors, response=response ) assert isinstance(exception, exceptions.OutOfRange) assert exception.message == message assert exception.errors == errors assert exception.response == response
def _do_mutate_retryable_rows(self): """Mutate all the rows that are eligible for retry. A row is eligible for retry if it has not been tried or if it resulted in a transient error in a previous call. :rtype: list :return: ``responses_statuses`` (`google.rpc.status_pb2.Status`) :raises: :exc:`~google.api_core.exceptions.ServiceUnavailable` if any row returned a transient error. An artificial exception to work with ``DEFAULT_RETRY``. """ retryable_rows = [] index_into_all_rows = [] for i, status in enumerate(self.responses_statuses): if self._is_retryable(status): retryable_rows.append(self.rows[i]) index_into_all_rows.append(i) if not retryable_rows: # All mutations are either successful or non-retryable now. return self.responses_statuses mutate_rows_request = _mutate_rows_request(self.table_name, retryable_rows) responses = self.client._data_stub.MutateRows(mutate_rows_request) num_responses = 0 num_retryable_responses = 0 for response in responses: for entry in response.entries: num_responses += 1 index = index_into_all_rows[entry.index] self.responses_statuses[index] = entry.status if self._is_retryable(entry.status): num_retryable_responses += 1 if entry.status.code == 0: self.rows[index].clear() assert len(retryable_rows) == num_responses if num_retryable_responses: raise from_grpc_status(StatusCode.UNAVAILABLE, 'MutateRows retryable error.') return self.responses_statuses
def _on_response(self, response: gapic_types.AppendRowsResponse): """Process a response from a consumer callback.""" # If the stream has closed, but somehow we still got a response message # back, discard it. The response futures queue has been drained, with # an exception reported. if self._closed: raise bqstorage_exceptions.StreamClosedError( f"Stream closed before receiving response: {response}") # Since we have 1 response per request, if we get here from a response # callback, the queue should never be empty. future: AppendRowsFuture = self._futures_queue.get_nowait() if response.error.code: exc = exceptions.from_grpc_status(response.error.code, response.error.message) future.set_exception(exc) else: future.set_result(response)
def test_from_grpc_status_unknown_code(): message = 'message' exception = exceptions.from_grpc_status(grpc.StatusCode.OK, message) assert exception.grpc_status_code == grpc.StatusCode.OK assert exception.message == message
def test_from_grpc_status_unknown_code(): message = "message" exception = exceptions.from_grpc_status(grpc.StatusCode.OK, message) assert exception.grpc_status_code == grpc.StatusCode.OK assert exception.message == message