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)