session: Create stream_recv() to read everything from a socket

Some session requests will return large binary data through the socket
until it's complete, then the socket is closed.  Provide a new method
stream_recv() which will support this behavior so callers don't need to
grovel around in our internal class methods and data.

Create a test for this.

Author: madscientist

pynuodb/session.py

@@ -26,8 +26,8 @@
     from urlparse import urlparse  # type: ignore
 
 try:
-    from typing import Dict, Iterable, Mapping  # pylint: disable=unused-import
-    from typing import Optional, Tuple, Union   # pylint: disable=unused-import
+    from typing import Dict, Generator, Iterable, Mapping  # pylint: disable=unused-import
+    from typing import Optional, Tuple, Union  # pylint: disable=unused-import
 except ImportError:
     pass
 
@@ -489,6 +489,31 @@ def __readFully(self, msgLength, timeout=None):
 
         return bytes(msg)
 
+    def stream_recv(self, blocksz=4096, timeout=None):
+        # type: (int, Optional[float]) -> Generator[bytes, None, None]
+        """Read data from the socket in blocksz increments.
+
+        Will yield bytes buffers of blocksz for as long as the sender is
+        sending.  After this function completes the socket has been closed.
+        Note it's best if blocksz is a multiple of 32, to ensure that block
+        ciphers will work.  This code doesn't manage block sizes or padding.
+
+        If timeout is not None, raises a socket.timeout exception on timeout.
+        The socket is still closed.
+        """
+        sock = self._sock
+        try:
+            sock.settimeout(timeout)
+            while True:
+                msg = sock.recv(blocksz)
+                if not msg:
+                    break
+                if self.__cipherIn:
+                    msg = self.__cipherIn.transform(msg)
+                yield msg
+        finally:
+            self.close()
+
     def close(self, force=False):
         # type: (bool) -> None
         """Close the current socket connection with the server."""

tests/nuodb_service_test.py

@@ -1,6 +1,9 @@
 #!/usr/bin/env python
 
+import os
 import unittest
+import tempfile
+import gzip
 import xml.etree.ElementTree as ET
 
 import pynuodb
@@ -64,6 +67,41 @@ def test_request_gc(self):
         self.assertEqual(len(info), 1)
         self.assertEqual(info[0].get('Action'), 'RequestGarbageCollection')
 
+    def test_stream_recv(self):
+        """Test the stream_recv() facility."""
+        ap_conn = get_ap_conn()
+        if ap_conn is None:
+            self.skipTest("No AP available")
+
+        procs = ap_conn.get_processes(db_name=DATABASE_NAME)
+        dbpasswd = ap_conn._get_db_password(DATABASE_NAME)
+
+        session = pynuodb.session.Session(
+            procs[0].address, service='Admin',
+            options={'verifyHostname': 'False'})
+        session.authorize('Cloud', dbpasswd)
+
+        session.send('''<Request Service="Admin">
+                          <Request Type="GetSysDepends"/>
+                        </Request>''')
+        resp = session.recv()
+        xml = ET.fromstring(resp)
+        self.assertIsNotNone(xml.find('Success'), "Failed: %s" % (resp))
+
+        deppath = os.path.join(tempfile.gettempdir(), 'deps.tar.gz')
+        with open(deppath, 'wb') as of:
+            for data in session.stream_recv():
+                of.write(data)
+
+        # The socket should be closed now: this will raise
+        self.assertRaises(pynuodb.session.SessionException,
+                          lambda: session._sock)
+
+        # Now make sure that what we read is uncompressable
+        with gzip.GzipFile(deppath, 'rb') as gz:
+            # We don't really care we just want to make sure it works
+            self.assertIsNotNone(gz.read(), "Failed to unzip %s" % (deppath))
+
 
 if __name__ == '__main__':
     unittest.main()