Restore socket timeout after partial-read recovery in _recv_exact_py

When a TimeoutError occurs mid-message, the function removes the socket
timeout to avoid protocol desync. Previously the timeout was never
restored, causing the caller's shutdown-polling loop to block
indefinitely on subsequent reads. Now the original timeout is saved
and restored on both EOF and successful completion.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: kesmit13

singlestoredb/functions/ext/collocated/connection.py

@@ -394,6 +394,7 @@ def _recv_exact_py(sock: socket.socket, n: int) -> bytes | None:
     buf = bytearray(n)
     view = memoryview(buf)
     pos = 0
+    orig_timeout = sock.gettimeout()
     while pos < n:
         try:
             nbytes = sock.recv_into(view[pos:])
@@ -405,8 +406,12 @@ def _recv_exact_py(sock: socket.socket, n: int) -> bytes | None:
             sock.settimeout(None)
             continue
         if nbytes == 0:
+            if orig_timeout is not None:
+                sock.settimeout(orig_timeout)
             return None
         pos += nbytes
+    if orig_timeout is not None:
+        sock.settimeout(orig_timeout)
     return bytes(buf)