def __init__(self, s: socket): self.s = s self.file = self.s.makefile("rwb") self.wrapper = IOWrapper(self.file, socket_=s) self.current_id = -1 self.handshake()
class TChannelConnection: file: IO wrapper: IOWrapper s: socket.socket @classmethod def open(cls, host: object, port: object, timeout: int = None) -> TChannelConnection: s: socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) s.settimeout(timeout) s.connect((host, port)) return cls(s) def __init__(self, s: socket): self.s = s self.file = self.s.makefile("rwb") self.wrapper = IOWrapper(self.file, socket_=s) self.current_id = -1 self.handshake() def set_next_timeout_cb(self, cb: Callable): self.wrapper.set_next_timeout_cb(cb) def new_id(self): self.current_id += 1 return self.current_id def handshake(self): req: InitReqFrame = InitReqFrame() req.id = self.new_id() req.headers.d["host_port"] = "0.0.0.0:0" req.headers.d["process_name"] = "python-process" self.write_frame(req) res = self.read_frame() if res.TYPE != 0x02: raise Exception("Unexpected response from server") def write_frame(self, frame: Frame): frame.write(self.wrapper) self.wrapper.flush() def read_frame(self): frame = Frame.read_frame(self.wrapper) if isinstance(frame, ErrorFrame): raise TChannelException(error_frame=frame) return frame def close(self): self.s.close() self.wrapper.close() def call_function(self, call: ThriftFunctionCall) -> ThriftFunctionResponse: frames = call.build_frames(self.new_id()) for frame in frames: self.write_frame(frame) response = ThriftFunctionResponse() while not response.is_complete(): frame = self.read_frame() if frame.TYPE not in (CallResFrame.TYPE, CallResContinueFrame.TYPE): raise Exception("Unexpected type: " + Frame.TYPE) assert isinstance(frame, FrameWithArgs) response.process_frame(frame) return response
def build_arg1(self) -> bytes: f = BytesIO() wrapper: IOWrapper = IOWrapper(f) wrapper.write_string(self.method_name) wrapper.flush() return f.getvalue()
def build_arg2(self) -> bytes: f = BytesIO() wrapper: IOWrapper = IOWrapper(f) h = KVHeaders(self.application_headers, 2) h.write_headers(wrapper) return f.getvalue()
def process_arg2(self, b): f: BytesIO = BytesIO(b) wrapper: IOWrapper = IOWrapper(f) h: KVHeaders = KVHeaders.read_kv_headers(wrapper, 2, "ThriftFunctionResponse") self.application_headers = h.d