def _handleResponse(self, token, data): try: cursor = self._cursor_cache.get(token) if cursor is not None: cursor._extend(data) elif token in self._user_queries: query, deferred = self._user_queries[token] res = Response(token, data, self._parent._get_json_decoder(query)) if res.type == pResponse.SUCCESS_ATOM: deferred.callback(maybe_profile(res.data[0], res)) elif res.type in (pResponse.SUCCESS_SEQUENCE, pResponse.SUCCESS_PARTIAL): cursor = TwistedCursor(self, query, res) deferred.callback(maybe_profile(cursor, res)) elif res.type == pResponse.WAIT_COMPLETE: deferred.callback(None) elif res.type == pResponse.SERVER_INFO: deferred.callback(res.data[0]) else: deferred.errback(res.make_error(query)) del self._user_queries[token] elif not self._closing: raise ReqlDriverError("Unexpected response received.") except Exception as e: if not self._closing: self.close(exception=e)
def _reader(self): try: while True: buf = yield from self._streamreader.readexactly(12) (token, length,) = struct.unpack("<qL", buf) buf = yield from self._streamreader.readexactly(length) cursor = self._cursor_cache.get(token) if cursor is not None: cursor._extend(buf) elif token in self._user_queries: # Do not pop the query from the dict until later, so # we don't lose track of it in case of an exception query, future = self._user_queries[token] res = Response(token, buf, self._parent._get_json_decoder(query)) if res.type == pResponse.SUCCESS_ATOM: future.set_result(maybe_profile(res.data[0], res)) elif res.type in ( pResponse.SUCCESS_SEQUENCE, pResponse.SUCCESS_PARTIAL, ): cursor = AsyncioCursor(self, query, res) future.set_result(maybe_profile(cursor, res)) elif res.type == pResponse.WAIT_COMPLETE: future.set_result(None) elif res.type == pResponse.SERVER_INFO: future.set_result(res.data[0]) else: future.set_exception(res.make_error(query)) del self._user_queries[token] elif not self._closing: raise ReqlDriverError("Unexpected response received.") except Exception as ex: if not self._closing: yield from self.close(exception=ex)
def data_received(self, data): if self._response_buf: data = self._response_buf + data try: value = None if self.connection.requests_processed == 1: if b'\0' in data: idx = data.index(b'\0') msg, data = data[:idx], data[idx+1:] if msg != b'SUCCESS': raise RqlDriverError( 'Server dropped connection: "%s"' % msg.decode('utf-8').strip()) elif len(data) >= 12: buf, data = data[:12], data[12:] token, size = struct.unpack("<qL", buf) if len(data) >= size: buf, data = data[:size], data[size:] response = Response(token, buf) if data: raise RqlDriverError('protocol error') elif response.token != self.request.token: raise RqlDriverError('Unexpected response token') value = self._check_response(response) else: data = buf + data except Exception as exc: self.finished(exc=exc) else: self._response_buf = data if not data: self.finished(value)
def _extend(self, res_buf): ''' Override so that we can make this async, and also to wake up blocked tasks. ''' self.outstanding_requests -= 1 self._maybe_fetch_batch() res = Response(self.query.token, res_buf, self._json_decoder) self._extend_internal(res) self._new_response.set() self._new_response = trio.Event()