Merge pull request #1159 from Vic020/fix_connection_reading_socket_error

Connection Reading Socket Error Enhancement
redis · Apr 16, 2019 · 99c5899 · 99c5899
2 parents 2652c6b + 5d3d41a
commit 99c5899
Showing 1 changed file with 31 additions and 36 deletions.
diff --git a/redis/connection.py b/redis/connection.py
@@ -177,26 +177,19 @@ def _read_from_socket(self, length=None):
         buf.seek(self.bytes_written)
         marker = 0
 
-        try:
-            while True:
-                data = recv(self._sock, 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)
-                buf.write(data)
-                data_length = len(data)
-                self.bytes_written += data_length
-                marker += data_length
-
-                if length is not None and length > marker:
-                    continue
-                break
-        except socket.timeout:
-            raise TimeoutError("Timeout reading from socket")
-        except socket.error:
-            e = sys.exc_info()[1]
-            raise ConnectionError("Error while reading from socket: %s" %
-                                  (e.args,))
+        while True:
+            data = recv(self._sock, 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)
+            buf.write(data)
+            data_length = len(data)
+            self.bytes_written += data_length
+            marker += data_length
+
+            if length is not None and length > marker:
+                continue
+            break
 
     def read(self, length):
         length = length + 2  # make sure to read the \r\n terminator
@@ -391,22 +384,15 @@ def read_response(self):
         response = self._reader.gets()
         socket_read_size = self.socket_read_size
         while response is False:
-            try:
-                if HIREDIS_USE_BYTE_BUFFER:
-                    bufflen = recv_into(self._sock, self._buffer)
-                    if bufflen == 0:
-                        raise socket.error(SERVER_CLOSED_CONNECTION_ERROR)
-                else:
-                    buffer = recv(self._sock, socket_read_size)
-                    # an empty string indicates the server shutdown the socket
-                    if not isinstance(buffer, bytes) or len(buffer) == 0:
-                        raise socket.error(SERVER_CLOSED_CONNECTION_ERROR)
-            except socket.timeout:
-                raise TimeoutError("Timeout reading from socket")
-            except socket.error:
-                e = sys.exc_info()[1]
-                raise ConnectionError("Error while reading from socket: %s" %
-                                      (e.args,))
+            if HIREDIS_USE_BYTE_BUFFER:
+                bufflen = recv_into(self._sock, self._buffer)
+                if bufflen == 0:
+                    raise socket.error(SERVER_CLOSED_CONNECTION_ERROR)
+            else:
+                buffer = recv(self._sock, socket_read_size)
+                # an empty string indicates the server shutdown the socket
+                if not isinstance(buffer, bytes) or len(buffer) == 0:
+                    raise socket.error(SERVER_CLOSED_CONNECTION_ERROR)
             if HIREDIS_USE_BYTE_BUFFER:
                 self._reader.feed(self._buffer, 0, bufflen)
             else:
@@ -639,6 +625,15 @@ def read_response(self):
         "Read the response from a previously sent command"
         try:
             response = self._parser.read_response()
+        except socket.timeout:
+            self.disconnect()
+            raise TimeoutError("Timeout reading from %s:%s" %
+                               (self.host, self.port))
+        except socket.error:
+            self.disconnect()
+            e = sys.exc_info()[1]
+            raise ConnectionError("Error while reading from %s:%s : %s" %
+                                  (self.host, self.port, e.args))
         except:  # noqa: E722
             self.disconnect()
             raise