def test_partition_query(self):
# Setup Expected Response
expected_response = {}
expected_response = spanner_pb2.PartitionResponse(**expected_response)
# Mock the API response
channel = ChannelStub(responses=[expected_response])
patch = mock.patch("google.api_core.grpc_helpers.create_channel")
with patch as create_channel:
create_channel.return_value = channel
client = spanner_v1.SpannerClient()
# Setup Request
session = client.session_path("[PROJECT]", "[INSTANCE]", "[DATABASE]",
"[SESSION]")
sql = "sql114126"
response = client.partition_query(session, sql)
assert expected_response == response
assert len(channel.requests) == 1
expected_request = spanner_pb2.PartitionQueryRequest(session=session,
sql=sql)
actual_request = channel.requests[0][1]
assert expected_request == actual_request
def _partition(stub):
"""Probe to test PartitionQuery and PartitionRead grpc call from Spanner stub.
Args:
stub: An object of SpannerStub.
"""
_partition_tracer = initialize_tracer()
with _partition_tracer.span(name='_partition') as root_span:
root_span.add_annotation('endpoint info available',
endpoint=_SPANNER_TARGET)
session = None
try:
with _partition_tracer.span(name='stub.CreateSession'):
session = stub.CreateSession(
spanner_pb2.CreateSessionRequest(database=_DATABASE))
txn_options = transaction_pb2.TransactionOptions(
read_only=transaction_pb2.TransactionOptions.ReadOnly())
txn_selector = transaction_pb2.TransactionSelector(
begin=txn_options)
# Probing PartitionQuery call
ptn_query_request = spanner_pb2.PartitionQueryRequest(
session=session.name,
sql='select * FROM users',
transaction=txn_selector,
)
with _partition_tracer.span(name='stub.PartitionQuery'):
stub.PartitionQuery(ptn_query_request)
# Probing PartitionRead call
ptn_read_request = spanner_pb2.PartitionReadRequest(
session=session.name,
table='users',
transaction=txn_selector,
key_set=keys_pb2.KeySet(all=True),
columns=['username', 'firstname', 'lastname'])
with _partition_tracer.span(name='stub.PartitionRead'):
stub.PartitionRead(ptn_read_request)
finally:
if session is not None:
with _partition_tracer.span(name='stub.DeleteSession'):
stub.DeleteSession(
spanner_pb2.DeleteSessionRequest(name=session.name))
def test_partition_query(self):
# Setup Expected Response
expected_response = {}
expected_response = spanner_pb2.PartitionResponse(**expected_response)
# Mock the API response
channel = ChannelStub(responses=[expected_response])
client = spanner_v1.SpannerClient(channel=channel)
# Setup Request
session = client.session_path('[PROJECT]', '[INSTANCE]', '[DATABASE]',
'[SESSION]')
sql = 'sql114126'
response = client.partition_query(session, sql)
assert expected_response == response
assert len(channel.requests) == 1
expected_request = spanner_pb2.PartitionQueryRequest(session=session,
sql=sql)
actual_request = channel.requests[0][1]
assert expected_request == actual_request