def __init__(self, host='localhost', port=6379, **kwargs): super(OneShotConnection, self).__init__(host=host, port=port, **kwargs) self.returnedResult = None # Get a basic parser that knows nothing about connections # and sockets: self._parser = BaseParser() # When more bytes have been pulled from the server # than were requested in one of the read_xxx() methods # below, the bytes are saved in the following: self.read_bytes = None # Similarly, readline() remembers lines self.lines_read = [] self.remnant = None self.connect()
class OneShotConnection(Connection): ''' Connection that is not backed by a buffer, and is not a thread. Used exclusively for messages to the Redis server that expect a single return data item. Methods get_int() and get_string() know about the Redis wire protocol, and extract the payload. No external parser is used. No callbacks are made to anywhere. ''' def __init__(self, host='localhost', port=6379, **kwargs): super(OneShotConnection, self).__init__(host=host, port=port, **kwargs) self.returnedResult = None # Get a basic parser that knows nothing about connections # and sockets: self._parser = BaseParser() # When more bytes have been pulled from the server # than were requested in one of the read_xxx() methods # below, the bytes are saved in the following: self.read_bytes = None # Similarly, readline() remembers lines self.lines_read = [] self.remnant = None self.connect() def read_int(self, block=True, timeout=None): ''' Reads a Redis int from the socket. Ensures that it is an int, i.e. of the form ":[0-9]*\r\n. :param block: whether to block, or return immediately if no data is available on the socket. :type block: bool :param timeout: if block==True, how long to block. None or 0: wait forever. :type timeout: float :return: integer that was returned by the Redis server :rtype: int :raise ResponseError if returned value is not an integer :raises TimeoutError: if no data arrives from server in time. ''' rawRes = self.readline(block=block, timeout=timeout) try: intRes = int(rawRes[1:].strip()) except (ValueError, IndexError): raise ResponseError("Server did not return an int; returned '%s'" % rawRes) return intRes def read(self, num_bytes, block=True, timeout=None): bytes_left_to_read = num_bytes while bytes_left_to_read > 0: if len(self.lines_read) > 0: # There already are lines that were read earlier. # read those first: result = '' for line in self.lines_read: if len(line) > bytes_left_to_read: result = line[:bytes_left_to_read] # Remove the read bytes from the line: self.lines_read = self.lines_read[bytes_left_to_read:] return result elif len(line) == bytes_left_to_read: return self.lines_read.pop(0) elif len(line) < bytes_left_to_read: # Need to use this line, plus more: line = self.lines_read.pop(0) result += line bytes_left_to_read -= len(line) continue if self.remnant is not None: if len(self.remnant) <= bytes_left_to_read: result += self.remnant if len(self.remnant) == bytes_left_to_read: self.remnant = None return result else: # Remnant > than what we still need len_remnant_part_used = bytes_left_to_read result += self.remnant[:bytes_left_to_read] self.remnant = self.remnant[bytes_left_to_read:] bytes_left_to_read -= len_remnant_part_used return result # Either no lines had previously been read, or # not enough data was in previously read lines # and self.remnant: next_line = self.readline(block=block, timeout=timeout) self.lines_read = [next_line] + self.lines_read def readline(self, block=True, timeout=None): ''' Return one line from the socket, without the closing CR/LF, or raise a TimeoutError exception. If socket yields more than one line, the remaining lines are stored in the self.lines_read array. If a partial line is at the end of the socket, that is stored in self.remnant, and prepended to subsequent readline() results. :param block: whether or not to wait for data :type block: boolean :param timeout: (fractional) seconds before timing out :type timeout: float :return: a string corresponding to one \r\n-delimited element of the wire protocol. :rtype: string :raise TimeoutError if no data from server in time. :raise socket.error if socket problem ''' if len(self.lines_read) > 0: return self.lines_read.pop(0) while True: try: # Wait for incoming messages: if block: if timeout is None: # Block forever: (readReady, writeReady, errReady) = select([self._sock],[],[]) #@UnusedVariable else: # Block, but with timeout: (readReady, writeReady, errReady) = select([self._sock],[],[], timeout) #@UnusedVariable else: # Just poll: (readReady, writeReady, errReady) = select([self._sock],[],[], 0) #@UnusedVariable return None # Something arrived on the socket. data = self._sock.recv(self.socket_read_size) # An empty string indicates the server shutdown the socket if isinstance(data, bytes) and len(data) == 0: raise socket.error(SERVER_CLOSED_CONNECTION_ERROR) # If over-read before, prepend that remnant: if self.remnant is not None: data = self.remnant + data self.remnant = None except socket.timeout: raise TimeoutError("Server did not respond in time when we expected a line of data.") except socket.error as e: if e.args[0] == errno.EAGAIN: time.sleep(0.3) continue else: raise if SYM_CRLF in data: lines = data.split(SYM_CRLF) # str.split() adds an empty str if the # str ends in the split symbol. If str does # not end in the split symbol, the resulting # array has the unfinished fragment at the end. for line in lines[:-1]: self.lines_read.append(line) # Final bytes may or may not have their # closing SYM_CRLF yet: if not data.endswith(SYM_CRLF): # Have a partial line at the end: self.remnant = lines[-1] # Return the first line: return self.lines_read.pop(0) def read_string(self, block=True, timeout=None): ''' Returns a Redis wire protocol encoded string from the socket. Expect incoming data to be a Redis 'simple string', or a lenth-specified string: Simple string form: +<str> Length specified string form: "*2\r\n$<strLen>\r\n<str> :param block: block for data to arrive from server, or not. :type block: bool :param timeout: if block == True, when to time out :type timeout: float :return: the string returned by the server :rtype: string :raise ResponseError: if response arrives, but is not in proper string format :raises TimeoutError: if no data arrives from server in time. ''' # Strings start with a +, followed by a string, followed # by a CR/LF, or with a length, like this: # '$9\r\nsubscribe\r\n$5\r\ntmp.0\r\n:1\r\n'\r\n # Get the length or simple-string rawRes = self.readline(block=block, timeout=timeout) # Is it a 'simple string', i.e. "+<str>"? if rawRes[0] == '+': return(rawRes[1:-len(SYM_CRLF)]) # Looking at '$<strLen>': try: str_len = int(rawRes[1:]) except ValueError: raise ResponseError("Expected integer string length, but received '%s'" % rawRes) res_str = self.readline(block=block, timeout=timeout) if len(res_str) != str_len: raise InvalidResponse("String length %d is not the length of '%s'" % (str_len, res_str)) return res_str def write_socket(self, msg): ''' Write an arbitrary string to this OneShotConnection's socket. :param msg: message to write :type msg: string :raises: socket.error if something bad happens in the socket system call ''' self._sock.sendall(msg) def parse_response(self, response=None, block=True, timeout=None): ''' Obtain a server stream of bytes directly from the socket, and parse it until one completely parsed Redis structure is obtained. Any additional bytes read from the buffer are safed in self.lines_read and self.remnant. :param response: a previously obtained line of Redis wire protocol. If none, readline() is called. :type response: string :param block: if True, wait for data from server :type block: boolean :param timeout: (fractional) seconds to wait for data from server :type timeout: float :return: an array of parsed Redis :rtype: [string] ''' # Parse a response return self._parser.parse_response(response, self, encoding=self.encoding, timeout=timeout, block=block) def pack_publish_command(self, channel, msg): ''' Given a message and a channel, return a string that is the corresponding wire message. This is an optimized special case for PUBLISH messages. :param channel: the channel to which to publish :type channel: string :param msg: the message to publish :type msg: string ''' wire_msg = '\r\n'.join(['*3', # 3 parts to follow '$7', # 7 letters 'PUBLISH', # <command> '$%d' % len(channel),# num topic-letters to follow channel, # <topic> '$%d' % len(msg), # num msg-letters to follow msg, # <msg> '' # forces a closing \r\n ]) return wire_msg def pack_subscription_command(self, command, channel): ''' Given a channel and a subscribe or unsubscribe name, return a string that is the corresponding wire message. This is an optimized special case for (un)subscribe messages. :param command: which command to construct; must be one of {SUBSCRIBE | UNSUBSCRIBE | PSUBSCRIBE | PUNSUBSCRIBE} :type command: string :param channel: the channel to which to subscribe :type channel: string ''' wire_msg = '\r\n'.join(['*2', # 2 parts to follow '$%d' % len(command), # number of letters in command command, # <command> '$%d' % len(channel), # num channel-letters to follow channel, # <topic> '' # forces a closing \r\n ]) return wire_msg