def fetch(self, request_id: UUID) -> Union[transaction_proto.Transaction.Res, transaction_proto.Transaction.ResPart]:
# Keep taking responses until we get one that matches the request ID
while True:
try:
return self._response_collector.get(request_id).get(block=False)
except Empty:
pass
try:
if not self._is_open.get():
raise TypeDBClientException.of(TRANSACTION_CLOSED)
server_msg = next(self._response_iterator)
except RpcError as e:
error = TypeDBClientException.of_rpc(e)
self.close(error)
raise error
except StopIteration:
self.close()
raise TypeDBClientException.of(TRANSACTION_CLOSED)
server_case = server_msg.WhichOneof("server")
if server_case == "res":
self._collect(server_msg.res)
elif server_case == "res_part":
self._collect(server_msg.res_part)
else:
raise TypeDBClientException.of(ILLEGAL_ARGUMENT)
def close(self, error: TypeDBClientException = None):
if self._is_open.compare_and_set(True, False):
self._error = error
self._response_collector.close(error)
try:
self._dispatcher.close()
except RpcError as e:
raise TypeDBClientException.of_rpc(e)
def get(self, block: bool) -> R:
response = self._response_queue.get(block=block)
if response.is_value():
return response.value
elif response.is_done() and response.error is None:
raise TypeDBClientException.of(TRANSACTION_CLOSED)
elif response.is_done() and response.error is not None:
raise TypeDBClientException.of_rpc(response.error)
else:
raise TypeDBClientException.of(ILLEGAL_STATE)
def __init__(self, channel: Channel, credential: TypeDBCredential):
super(_ClusterServerStub, self).__init__()
self._credential = credential
self._channel = channel
self._stub = core_service_proto.TypeDBStub(channel)
self._cluster_stub = cluster_service_proto.TypeDBClusterStub(channel)
self._token = None
try:
res = self._cluster_stub.user_token(
cluster_user_token_req(self._credential.username()))
self._token = res.token
except RpcError as e:
e2 = TypeDBClientException.of_rpc(e)
if e2.error_message is not None and e2.error_message is not UNABLE_TO_CONNECT:
raise e2
def may_renew_token(self, function: Callable[[], T]) -> T:
try:
return self.resilient_call(function)
except TypeDBClientException as e:
if e.error_message is not None and e.error_message is CLUSTER_TOKEN_CREDENTIAL_INVALID:
self._token = None
res = self._cluster_stub.user_token(
cluster_user_token_req(self._credential.username()))
self._token = res.token
try:
return self.resilient_call(function)
except RpcError as e2:
raise TypeDBClientException.of_rpc(e2)
else:
raise e
def __init__(self,
session: "_TypeDBSessionImpl",
transaction_type: TransactionType,
options: TypeDBOptions = None):
if not options:
options = TypeDBOptions.core()
self._transaction_type = transaction_type
self._options = options
self._concept_manager = _ConceptManager(self)
self._query_manager = _QueryManager(self)
self._logic_manager = _LogicManager(self)
try:
self._channel, stub = session.client().new_channel_and_stub()
self._bidirectional_stream = BidirectionalStream(
stub, session.transmitter())
req = transaction_open_req(session.session_id(),
transaction_type.proto(),
options.proto(),
session.network_latency_millis())
self.execute(request=req, batch=False)
except RpcError as e:
raise TypeDBClientException.of_rpc(e)
def resilient_call(function: Callable[[], T]) -> T:
try:
# TODO actually implement forced gRPC to reconnected rapidly, which provides resilience
return function()
except RpcError as e:
raise TypeDBClientException.of_rpc(e)