def test_execute_sql_future(self):
stub = spanner_pb2_grpc.SpannerStub(self.channel)
session = stub.CreateSession(
spanner_pb2.CreateSessionRequest(database=_DATABASE))
self.assertEqual(1, len(self.channel._channel_refs))
self.assertEqual(1, self.channel._channel_refs[0]._affinity_ref)
self.assertEqual(0, self.channel._channel_refs[0]._active_stream_ref)
self.assertIsNotNone(session)
rendezvous = stub.ExecuteSql.future(
spanner_pb2.ExecuteSqlRequest(session=session.name, sql=_TEST_SQL))
self.assertEqual(1, len(self.channel._channel_refs))
self.assertEqual(1, self.channel._channel_refs[0]._affinity_ref)
self.assertEqual(1, self.channel._channel_refs[0]._active_stream_ref)
result_set = rendezvous.result()
self.assertEqual(1, len(self.channel._channel_refs))
self.assertEqual(1, self.channel._channel_refs[0]._affinity_ref)
self.assertEqual(0, self.channel._channel_refs[0]._active_stream_ref)
self.assertIsNotNone(result_set)
self.assertEqual(1, len(result_set.rows))
self.assertEqual(_TEST_COLUMN_DATA,
result_set.rows[0].values[0].string_value)
stub.DeleteSession(spanner_pb2.DeleteSessionRequest(name=session.name))
self.assertEqual(1, len(self.channel._channel_refs))
self.assertEqual(0, self.channel._channel_refs[0]._affinity_ref)
self.assertEqual(0, self.channel._channel_refs[0]._active_stream_ref)
def test_create_list_delete_session(self):
stub = spanner_pb2_grpc.SpannerStub(self.channel)
session = stub.CreateSession(
spanner_pb2.CreateSessionRequest(database=_DATABASE))
self.assertIsNotNone(session)
self.assertEqual(1, len(self.channel._channel_refs))
self.assertEqual(1, self.channel._channel_refs[0]._affinity_ref)
self.assertEqual(0, self.channel._channel_refs[0]._active_stream_ref)
sessions = stub.ListSessions(
spanner_pb2.ListSessionsRequest(database=_DATABASE))
self.assertIsNotNone(sessions.sessions)
self.assertIn(session.name, (s.name for s in sessions.sessions))
self.assertEqual(1, len(self.channel._channel_refs))
self.assertEqual(1, self.channel._channel_refs[0]._affinity_ref)
self.assertEqual(0, self.channel._channel_refs[0]._active_stream_ref)
stub.DeleteSession(spanner_pb2.DeleteSessionRequest(name=session.name))
self.assertEqual(1, len(self.channel._channel_refs))
self.assertEqual(0, self.channel._channel_refs[0]._affinity_ref)
self.assertEqual(0, self.channel._channel_refs[0]._active_stream_ref)
sessions = stub.ListSessions(
spanner_pb2.ListSessionsRequest(database=_DATABASE))
self.assertNotIn(session.name, (s.name for s in sessions.sessions))
self.assertEqual(1, len(self.channel._channel_refs))
self.assertEqual(0, self.channel._channel_refs[0]._affinity_ref)
self.assertEqual(0, self.channel._channel_refs[0]._active_stream_ref)
def test_create_session_reuse_channel(self):
stub = spanner_pb2_grpc.SpannerStub(self.channel)
for _ in range(_DEFAULT_MAX_CHANNELS_PER_TARGET * 2):
session = stub.CreateSession(
spanner_pb2.CreateSessionRequest(database=_DATABASE))
self.assertIsNotNone(session)
self.assertEqual(1, len(self.channel._channel_refs))
stub.DeleteSession(
spanner_pb2.DeleteSessionRequest(name=session.name))
def test_bound_after_unbind(self):
stub = spanner_pb2_grpc.SpannerStub(self.channel)
session = stub.CreateSession(
spanner_pb2.CreateSessionRequest(database=_DATABASE))
self.assertEqual(1, len(self.channel._channel_ref_by_affinity_key))
stub.DeleteSession(spanner_pb2.DeleteSessionRequest(name=session.name))
self.assertEqual(0, len(self.channel._channel_ref_by_affinity_key))
with self.assertRaises(Exception) as context:
stub.GetSession(spanner_pb2.GetSessionRequest(name=session.name))
self.assertEqual(grpc.StatusCode.NOT_FOUND, context.exception.code())
def test_bound_unbind_with_invalid_affinity_key(self):
stub = spanner_pb2_grpc.SpannerStub(self.channel)
with self.assertRaises(Exception) as context:
stub.GetSession(spanner_pb2.GetSessionRequest(name='random_name'))
self.assertEqual(grpc.StatusCode.INVALID_ARGUMENT,
context.exception.code())
with self.assertRaises(Exception) as context:
stub.DeleteSession(
spanner_pb2.DeleteSessionRequest(name='random_name'))
self.assertEqual(grpc.StatusCode.INVALID_ARGUMENT,
context.exception.code())
def _execute_probe(api, use_extension=False):
"""Execute a probe function given certain Cloud api and probe name.
Args:
api: the name of the api provider, e.g. "spanner", "firestore".
use_extension: option to use grpc-gcp extension when creating channel.
Raises:
NotImplementedError: An error occurred when api does not match any records.
"""
util = StackdriverUtil(api)
if api == 'spanner':
channel = _get_stub_channel(_SPANNER_TARGET, use_extension)
stub = spanner_pb2_grpc.SpannerStub(channel)
probe_functions = spanner_probes.PROBE_FUNCTIONS
elif api == 'firestore':
channel = _get_stub_channel(_FIRESTORE_TARGET)
stub = firestore_pb2_grpc.FirestoreStub(channel)
probe_functions = firestore_probes.PROBE_FUNCTIONS
else:
raise NotImplementedError('gRPC prober is not implemented for %s !' %
api)
total = len(probe_functions)
success = 0
metrics = {}
# Execute all probes for given api
for probe_name in probe_functions:
probe_function = probe_functions[probe_name]
try:
probe_function(stub, metrics)
success += 1
except Exception: # pylint: disable=broad-except
# report any kind of exception to Stackdriver
util.report_error(traceback.format_exc())
if success == total:
util.set_success(True)
# Summarize metrics
util.add_metrics_dict(metrics)
util.output_metrics()
# Fail this probe if any function fails
if success != total:
sys.exit(1)
def test_channel_connectivity_invalid_target(self):
config = config = grpc_gcp.api_config_from_text_pb(
pkg_resources.resource_string(__name__, 'spanner.grpc.config'))
http_request = Request()
credentials, _ = google.auth.default([_OAUTH_SCOPE], http_request)
invalid_channel = self._create_secure_gcp_channel(
credentials,
http_request,
'localhost:1234',
options=[(grpc_gcp.API_CONFIG_CHANNEL_ARG, config)])
callback = _Callback()
invalid_channel.subscribe(callback.update_first, try_to_connect=False)
stub = spanner_pb2_grpc.SpannerStub(invalid_channel)
with self.assertRaises(Exception) as context:
stub.CreateSession(
spanner_pb2.CreateSessionRequest(database=_DATABASE))
self.assertEqual(grpc.StatusCode.UNAVAILABLE, context.exception.code())
first_connectivities = callback.block_until_connectivities_satisfy(
lambda connectivities: len(connectivities) >= 3)
self.assertEqual(grpc.ChannelConnectivity.IDLE,
first_connectivities[0])
self.assertIn(grpc.ChannelConnectivity.CONNECTING,
first_connectivities)
self.assertIn(grpc.ChannelConnectivity.TRANSIENT_FAILURE,
first_connectivities)
invalid_channel.subscribe(callback.update_second, try_to_connect=True)
second_connectivities = callback.block_until_connectivities_satisfy(
lambda connectivities: len(connectivities) >= 3, False)
self.assertNotIn(grpc.ChannelConnectivity.IDLE, second_connectivities)
self.assertIn(grpc.ChannelConnectivity.CONNECTING,
second_connectivities)
self.assertIn(grpc.ChannelConnectivity.TRANSIENT_FAILURE,
second_connectivities)
self.assertEqual(2, len(invalid_channel._subscribers))
invalid_channel.unsubscribe(callback.update_first)
invalid_channel.unsubscribe(callback.update_second)
self.assertEqual(0, len(invalid_channel._subscribers))
def test_create_session_new_channel(self):
stub = spanner_pb2_grpc.SpannerStub(self.channel)
futures = []
for i in range(_DEFAULT_MAX_CHANNELS_PER_TARGET):
futures.append(
stub.CreateSession.future(
spanner_pb2.CreateSessionRequest(database=_DATABASE)))
self.assertEqual(i + 1, len(self.channel._channel_refs))
for future in futures:
stub.DeleteSession(
spanner_pb2.DeleteSessionRequest(name=future.result().name))
futures = []
for i in range(_DEFAULT_MAX_CHANNELS_PER_TARGET):
futures.append(
stub.CreateSession.future(
spanner_pb2.CreateSessionRequest(database=_DATABASE)))
self.assertEqual(_DEFAULT_MAX_CHANNELS_PER_TARGET,
len(self.channel._channel_refs))
for future in futures:
stub.DeleteSession(
spanner_pb2.DeleteSessionRequest(name=future.result().name))
def test_channel_connectivity_multiple_subchannels(self):
callback = _Callback()
self.channel.subscribe(callback.update_first, try_to_connect=False)
stub = spanner_pb2_grpc.SpannerStub(self.channel)
futures = []
for _ in range(2):
futures.append(
stub.CreateSession.future(
spanner_pb2.CreateSessionRequest(database=_DATABASE)))
connectivities = callback.block_until_connectivities_satisfy(
lambda connectivities: grpc.ChannelConnectivity.READY in
connectivities)
self.assertEqual(2, len(self.channel._channel_refs))
self.assertSequenceEqual((grpc.ChannelConnectivity.IDLE,
grpc.ChannelConnectivity.CONNECTING,
grpc.ChannelConnectivity.READY),
connectivities)
for future in futures:
stub.DeleteSession(
spanner_pb2.DeleteSessionRequest(name=future.result().name))
def test_channel_connectivity(self):
callback = _Callback()
self.channel.subscribe(callback.update_first, try_to_connect=False)
stub = spanner_pb2_grpc.SpannerStub(self.channel)
session = stub.CreateSession(
spanner_pb2.CreateSessionRequest(database=_DATABASE))
connectivities = callback.block_until_connectivities_satisfy(
lambda connectivities: grpc.ChannelConnectivity.READY in
connectivities)
self.assertEqual(3, len(connectivities))
self.assertSequenceEqual((grpc.ChannelConnectivity.IDLE,
grpc.ChannelConnectivity.CONNECTING,
grpc.ChannelConnectivity.READY),
connectivities)
stub.DeleteSession(spanner_pb2.DeleteSessionRequest(name=session.name))
self.channel.unsubscribe(callback.update_first)
session = stub.CreateSession(
spanner_pb2.CreateSessionRequest(database=_DATABASE))
self.assertEqual(3, len(connectivities))
stub.DeleteSession(spanner_pb2.DeleteSessionRequest(name=session.name))
def _create_stub(channel):
stub = spanner_pb2_grpc.SpannerStub(channel)
return stub
def test_concurrent_streams_watermark(self):
stub = spanner_pb2_grpc.SpannerStub(self.channel)
watermark = 2
self.channel._max_concurrent_streams_low_watermark = watermark
self.assertEqual(self.channel._max_concurrent_streams_low_watermark,
watermark)
session_list = []
rendezvous_list = []
# When active streams have not reached the concurrent_streams_watermark,
# gRPC calls should be reusing the same channel.
for i in range(watermark):
session = stub.CreateSession(
spanner_pb2.CreateSessionRequest(database=_DATABASE))
self.assertEqual(1, len(self.channel._channel_refs))
self.assertEqual(i + 1,
self.channel._channel_refs[0]._affinity_ref)
self.assertEqual(i,
self.channel._channel_refs[0]._active_stream_ref)
self.assertIsNotNone(session)
session_list.append(session)
rendezvous = stub.ExecuteStreamingSql(
spanner_pb2.ExecuteSqlRequest(session=session.name,
sql=_TEST_SQL))
self.assertEqual(1, len(self.channel._channel_refs))
self.assertEqual(i + 1,
self.channel._channel_refs[0]._affinity_ref)
self.assertEqual(i + 1,
self.channel._channel_refs[0]._active_stream_ref)
rendezvous_list.append(rendezvous)
# When active streams reach the concurrent_streams_watermark,
# channel pool will create a new channel.
another_session = stub.CreateSession(
spanner_pb2.CreateSessionRequest(database=_DATABASE))
self.assertEqual(2, len(self.channel._channel_refs))
self.assertEqual(2, self.channel._channel_refs[0]._affinity_ref)
self.assertEqual(2, self.channel._channel_refs[0]._active_stream_ref)
self.assertEqual(1, self.channel._channel_refs[1]._affinity_ref)
self.assertEqual(0, self.channel._channel_refs[1]._active_stream_ref)
self.assertIsNotNone(another_session)
session_list.append(another_session)
another_rendezvous = stub.ExecuteStreamingSql(
spanner_pb2.ExecuteSqlRequest(session=another_session.name,
sql=_TEST_SQL))
self.assertEqual(2, len(self.channel._channel_refs))
self.assertEqual(2, self.channel._channel_refs[0]._affinity_ref)
self.assertEqual(2, self.channel._channel_refs[0]._active_stream_ref)
self.assertEqual(1, self.channel._channel_refs[1]._affinity_ref)
self.assertEqual(1, self.channel._channel_refs[1]._active_stream_ref)
rendezvous_list.append(another_rendezvous)
# Iterate through the rendezous list to clean active streams.
for rendezvous in rendezvous_list:
for _ in rendezvous:
continue
# After cleaning, previously created channels will remain in the pool.
self.assertEqual(2, len(self.channel._channel_refs))
self.assertEqual(2, self.channel._channel_refs[0]._affinity_ref)
self.assertEqual(0, self.channel._channel_refs[0]._active_stream_ref)
self.assertEqual(1, self.channel._channel_refs[1]._affinity_ref)
self.assertEqual(0, self.channel._channel_refs[1]._active_stream_ref)
# Delete all sessions to clean affinity.
for session in session_list:
stub.DeleteSession(
spanner_pb2.DeleteSessionRequest(name=session.name))
self.assertEqual(2, len(self.channel._channel_refs))
self.assertEqual(0, self.channel._channel_refs[0]._affinity_ref)
self.assertEqual(0, self.channel._channel_refs[0]._active_stream_ref)
self.assertEqual(0, self.channel._channel_refs[1]._affinity_ref)
self.assertEqual(0, self.channel._channel_refs[1]._active_stream_ref)