def State( self, request_stream, # type: Iterable[beam_fn_api_pb2.StateRequest] context=None # type: Any ): # type: (...) -> Iterator[beam_fn_api_pb2.StateResponse] # Note that this eagerly mutates state, assuming any failures are fatal. # Thus it is safe to ignore instruction_id. for request in request_stream: request_type = request.WhichOneof('request') if request_type == 'get': data, continuation_token = self._state.get_raw( request.state_key, request.get.continuation_token) yield beam_fn_api_pb2.StateResponse( id=request.id, get=beam_fn_api_pb2.StateGetResponse( data=data, continuation_token=continuation_token)) elif request_type == 'append': self._state.append_raw(request.state_key, request.append.data) yield beam_fn_api_pb2.StateResponse( id=request.id, append=beam_fn_api_pb2.StateAppendResponse()) elif request_type == 'clear': self._state.clear(request.state_key) yield beam_fn_api_pb2.StateResponse( id=request.id, clear=beam_fn_api_pb2.StateClearResponse()) else: raise NotImplementedError('Unknown state request: %s' % request_type)
def State(self, request_stream, context=None): # Note that this eagerly mutates state, assuming any failures are fatal. # Thus it is safe to ignore instruction_reference. for request in request_stream: if request.get: yield beam_fn_api_pb2.StateResponse( id=request.id, get=beam_fn_api_pb2.StateGetResponse( data=self.blocking_get(request.state_key))) elif request.append: self.blocking_append(request.state_key, request.append.data) yield beam_fn_api_pb2.StateResponse( id=request.id, append=beam_fn_api_pb2.AppendResponse()) elif request.clear: self.blocking_clear(request.state_key) yield beam_fn_api_pb2.StateResponse( id=request.id, clear=beam_fn_api_pb2.ClearResponse())