def test_fake_server(self):
"""Fake server test."""
channel = grpc.insecure_channel(BigqueryOpsTest.server.endpoint())
stub = storage_pb2_grpc.BigQueryStorageStub(channel)
create_read_session_request = storage_pb2.CreateReadSessionRequest()
create_read_session_request.table_reference.project_id = self.GCP_PROJECT_ID
create_read_session_request.table_reference.dataset_id = self.DATASET_ID
create_read_session_request.table_reference.table_id = self.TABLE_ID
create_read_session_request.requested_streams = 2
read_session_response = stub.CreateReadSession(create_read_session_request)
self.assertEqual(2, len(read_session_response.streams))
read_rows_request = storage_pb2.ReadRowsRequest()
read_rows_request.read_position.stream.name = read_session_response.streams[
0].name
read_rows_response = stub.ReadRows(read_rows_request)
row = read_rows_response.next()
self.assertEqual(
self._serialize_to_avro(self.STREAM_1_ROWS, self.AVRO_SCHEMA),
row.avro_rows.serialized_binary_rows)
self.assertEqual(len(self.STREAM_1_ROWS), row.avro_rows.row_count)
read_rows_request = storage_pb2.ReadRowsRequest()
read_rows_request.read_position.stream.name = read_session_response.streams[
1].name
read_rows_response = stub.ReadRows(read_rows_request)
row = read_rows_response.next()
self.assertEqual(
self._serialize_to_avro(self.STREAM_2_ROWS, self.AVRO_SCHEMA),
row.avro_rows.serialized_binary_rows)
self.assertEqual(len(self.STREAM_2_ROWS), row.avro_rows.row_count)
def test_read_rows(self):
# Setup Expected Response
expected_response = {}
expected_response = storage_pb2.ReadRowsResponse(**expected_response)
# Mock the API response
channel = ChannelStub(responses=[iter([expected_response])])
patch = mock.patch("google.api_core.grpc_helpers.create_channel")
with patch as create_channel:
create_channel.return_value = channel
client = big_query_storage_client.BigQueryStorageClient()
# Setup Request
read_position = {}
response = client.read_rows(read_position)
resources = list(response)
assert len(resources) == 1
assert expected_response == resources[0]
assert len(channel.requests) == 1
expected_request = storage_pb2.ReadRowsRequest(
read_position=read_position)
actual_request = channel.requests[0][1]
assert expected_request == actual_request
def read_rows(
self,
read_position,
retry=google.api_core.gapic_v1.method.DEFAULT,
timeout=google.api_core.gapic_v1.method.DEFAULT,
metadata=None,
):
"""
Reads rows from the table in the format prescribed by the read session.
Each response contains one or more table rows, up to a maximum of 10 MiB
per response; read requests which attempt to read individual rows larger
than this will fail.
Each request also returns a set of stream statistics reflecting the
estimated total number of rows in the read stream. This number is computed
based on the total table size and the number of active streams in the read
session, and may change as other streams continue to read data.
Example:
>>> from google.cloud import bigquery_storage_v1beta1
>>>
>>> client = bigquery_storage_v1beta1.BigQueryStorageClient()
>>>
>>> # TODO: Initialize `read_position`:
>>> read_position = {}
>>>
>>> for element in client.read_rows(read_position):
... # process element
... pass
Args:
read_position (Union[dict, ~google.cloud.bigquery_storage_v1beta1.types.StreamPosition]): Required. Identifier of the position in the stream to start reading from.
The offset requested must be less than the last row read from ReadRows.
Requesting a larger offset is undefined.
If a dict is provided, it must be of the same form as the protobuf
message :class:`~google.cloud.bigquery_storage_v1beta1.types.StreamPosition`
retry (Optional[google.api_core.retry.Retry]): A retry object used
to retry requests. If ``None`` is specified, requests will
be retried using a default configuration.
timeout (Optional[float]): The amount of time, in seconds, to wait
for the request to complete. Note that if ``retry`` is
specified, the timeout applies to each individual attempt.
metadata (Optional[Sequence[Tuple[str, str]]]): Additional metadata
that is provided to the method.
Returns:
Iterable[~google.cloud.bigquery_storage_v1beta1.types.ReadRowsResponse].
Raises:
google.api_core.exceptions.GoogleAPICallError: If the request
failed for any reason.
google.api_core.exceptions.RetryError: If the request failed due
to a retryable error and retry attempts failed.
ValueError: If the parameters are invalid.
"""
# Wrap the transport method to add retry and timeout logic.
if "read_rows" not in self._inner_api_calls:
self._inner_api_calls[
"read_rows"
] = google.api_core.gapic_v1.method.wrap_method(
self.transport.read_rows,
default_retry=self._method_configs["ReadRows"].retry,
default_timeout=self._method_configs["ReadRows"].timeout,
client_info=self._client_info,
)
request = storage_pb2.ReadRowsRequest(read_position=read_position)
if metadata is None:
metadata = []
metadata = list(metadata)
try:
routing_header = [("read_position.stream.name", read_position.stream.name)]
except AttributeError:
pass
else:
routing_metadata = google.api_core.gapic_v1.routing_header.to_grpc_metadata(
routing_header
)
metadata.append(routing_metadata) # pragma: no cover
return self._inner_api_calls["read_rows"](
request, retry=retry, timeout=timeout, metadata=metadata
)
