session.send(): Allow sending of bytes, str, or bytearray

In pynuodb 2.x we only accepted str here, then we changed it to accept
only bytes.  Instead allow users to send any of bytes, str, or
bytearray.

Add a test of the different types of send.

Note, regardless of the type of message given to send(), recv() still
always returns bytes.

Author: madscientist

pynuodb/session.py

@@ -26,7 +26,8 @@
     from urlparse import urlparse  # type: ignore
 
 try:
-    from typing import Dict, Iterable, Mapping, Optional, Tuple  # pylint: disable=unused-import
+    from typing import Dict, Iterable, Mapping  # pylint: disable=unused-import
+    from typing import Optional, Tuple, Union   # pylint: disable=unused-import
 except ImportError:
     pass
 
@@ -394,19 +395,31 @@ def __constructServiceMessage(self,
         return message
 
     def send(self, message):
-        # type: (bytes) -> None
-        """Send an encoded message to the server over the socket."""
+        # type: (Union[str, bytes, bytearray]) -> None
+        """Send an encoded message to the server over the socket.
+
+        The message to be sent is either already-encoded bytes or bytearray,
+        or it's a UTF-8 str.
+        """
         sock = self._sock
+
+        if isinstance(message, bytearray):
+            data = bytes(message)
+        elif isinstance(message, bytes) or isP2:
+            data = message  # type: ignore
+        elif isinstance(message, str):
+            data = message.encode('utf-8')
+
         if self.__cipherOut:
-            message = self.__cipherOut.transform(message)
+            data = self.__cipherOut.transform(data)
 
-        lenStr = struct.pack("!I", len(message))
+        lenStr = struct.pack("!I", len(data))
 
         try:
             # We should send this in two parts to avoid making a complete copy
             # of the message when we send it.  But, I think the server may be
             # unhappy if it receives the length then has to wait for the data.
-            sock.send(lenStr + message)
+            sock.send(lenStr + data)
         except Exception:
             self.close()
             raise

tests/nuodb_service_test.py

@@ -21,18 +21,24 @@ def test_query_memory(self):
         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='Query',
-            options={'verifyHostname': 'False'})
-        session.authorize('Cloud', dbpasswd)
+        def try_message(msg):
+            session = pynuodb.session.Session(
+                procs[0].address, service='Query',
+                options={'verifyHostname': 'False'})
+            session.authorize('Cloud', dbpasswd)
+            session.send(msg)
+            res = session.recv()
+            root = ET.fromstring(res)
+            self.assertEqual(root.tag, 'MemoryInfo')
+            info = root.findall('HeapInformation')
+            self.assertEqual(len(info), 1)
+            self.assertEqual(info[0].tag, 'HeapInformation')
 
-        session.send('<Request Service="Query" Type="Memory"/>'.encode('utf-8'))
-        res = session.recv()
-        root = ET.fromstring(res)
-        self.assertEqual(root.tag, 'MemoryInfo')
-        info = root.findall('HeapInformation')
-        self.assertEqual(len(info), 1)
-        self.assertEqual(info[0].tag, 'HeapInformation')
+        # Send with different types of buffers
+        msg = '<Request Service="Query" Type="Memory"/>'
+        try_message(msg)
+        try_message(msg.encode('utf-8'))
+        try_message(pynuodb.crypt.bytesToArray(msg.encode('utf-8')))
 
     def test_request_gc(self):
         """Test a request operation."""