PYTHON-1369 Extend driver vector support to arbitrary subtypes and fix handling of variable length types (OSS C* 5.0) (#1217)

Author: absurdfarce

cassandra/__init__.py

@@ -744,9 +744,3 @@ def __init__(self, msg, excs=[]):
         if excs:
             complete_msg += ("\nThe following exceptions were observed: \n - " + '\n - '.join(str(e) for e in excs))
         Exception.__init__(self, complete_msg)
-
-class VectorDeserializationFailure(DriverException):
-    """
-    The driver was unable to deserialize a given vector
-    """
-    pass

cassandra/cqltypes.py

@@ -48,8 +48,8 @@
                                int32_pack, int32_unpack, int64_pack, int64_unpack,
                                float_pack, float_unpack, double_pack, double_unpack,
                                varint_pack, varint_unpack, point_be, point_le,
-                               vints_pack, vints_unpack)
-from cassandra import util, VectorDeserializationFailure
+                               vints_pack, vints_unpack, uvint_unpack, uvint_pack)
+from cassandra import util
 
 _little_endian_flag = 1  # we always serialize LE
 import ipaddress
@@ -392,6 +392,9 @@ def cass_parameterized_type(cls, full=False):
         """
         return cls.cass_parameterized_type_with(cls.subtypes, full=full)
 
+    @classmethod
+    def serial_size(cls):
+        return None
 
 # it's initially named with a _ to avoid registering it as a real type, but
 # client programs may want to use the name still for isinstance(), etc
@@ -457,10 +460,12 @@ def serialize(uuid, protocol_version):
         except AttributeError:
             raise TypeError("Got a non-UUID object for a UUID value")
 
+    @classmethod
+    def serial_size(cls):
+        return 16
 
 class BooleanType(_CassandraType):
     typename = 'boolean'
-    serial_size = 1
 
     @staticmethod
     def deserialize(byts, protocol_version):
@@ -470,6 +475,10 @@ def deserialize(byts, protocol_version):
     def serialize(truth, protocol_version):
         return int8_pack(truth)
 
+    @classmethod
+    def serial_size(cls):
+        return 1
+
 class ByteType(_CassandraType):
     typename = 'tinyint'
 
@@ -500,7 +509,6 @@ def serialize(var, protocol_version):
 
 class FloatType(_CassandraType):
     typename = 'float'
-    serial_size = 4
 
     @staticmethod
     def deserialize(byts, protocol_version):
@@ -510,10 +518,12 @@ def deserialize(byts, protocol_version):
     def serialize(byts, protocol_version):
         return float_pack(byts)
 
+    @classmethod
+    def serial_size(cls):
+        return 4
 
 class DoubleType(_CassandraType):
     typename = 'double'
-    serial_size = 8
 
     @staticmethod
     def deserialize(byts, protocol_version):
@@ -523,10 +533,12 @@ def deserialize(byts, protocol_version):
     def serialize(byts, protocol_version):
         return double_pack(byts)
 
+    @classmethod
+    def serial_size(cls):
+        return 8
 
 class LongType(_CassandraType):
     typename = 'bigint'
-    serial_size = 8
 
     @staticmethod
     def deserialize(byts, protocol_version):
@@ -536,10 +548,12 @@ def deserialize(byts, protocol_version):
     def serialize(byts, protocol_version):
         return int64_pack(byts)
 
+    @classmethod
+    def serial_size(cls):
+        return 8
 
 class Int32Type(_CassandraType):
     typename = 'int'
-    serial_size = 4
 
     @staticmethod
     def deserialize(byts, protocol_version):
@@ -549,6 +563,9 @@ def deserialize(byts, protocol_version):
     def serialize(byts, protocol_version):
         return int32_pack(byts)
 
+    @classmethod
+    def serial_size(cls):
+        return 4
 
 class IntegerType(_CassandraType):
     typename = 'varint'
@@ -645,14 +662,16 @@ def serialize(v, protocol_version):
 
         return int64_pack(int(timestamp))
 
+    @classmethod
+    def serial_size(cls):
+        return 8
 
 class TimestampType(DateType):
     pass
 
 
 class TimeUUIDType(DateType):
     typename = 'timeuuid'
-    serial_size = 16
 
     def my_timestamp(self):
         return util.unix_time_from_uuid1(self.val)
@@ -668,6 +687,9 @@ def serialize(timeuuid, protocol_version):
         except AttributeError:
             raise TypeError("Got a non-UUID object for a UUID value")
 
+    @classmethod
+    def serial_size(cls):
+        return 16
 
 class SimpleDateType(_CassandraType):
     typename = 'date'
@@ -699,7 +721,6 @@ def serialize(val, protocol_version):
 
 class ShortType(_CassandraType):
     typename = 'smallint'
-    serial_size = 2
 
     @staticmethod
     def deserialize(byts, protocol_version):
@@ -709,10 +730,14 @@ def deserialize(byts, protocol_version):
     def serialize(byts, protocol_version):
         return int16_pack(byts)
 
-
 class TimeType(_CassandraType):
     typename = 'time'
-    serial_size = 8
+    # Time should be a fixed size 8 byte type but Cassandra 5.0 code marks it as
+    # variable size... and we have to match what the server expects since the server
+    # uses that specification to encode data of that type.
+    #@classmethod
+    #def serial_size(cls):
+    #    return 8
 
     @staticmethod
     def deserialize(byts, protocol_version):
@@ -1409,6 +1434,11 @@ class VectorType(_CassandraType):
     vector_size = 0
     subtype = None
 
+    @classmethod
+    def serial_size(cls):
+        serialized_size = cls.subtype.serial_size()
+        return cls.vector_size * serialized_size if serialized_size is not None else None
+
     @classmethod
     def apply_parameters(cls, params, names):
         assert len(params) == 2
@@ -1418,19 +1448,50 @@ def apply_parameters(cls, params, names):
 
     @classmethod
     def deserialize(cls, byts, protocol_version):
-        serialized_size = getattr(cls.subtype, "serial_size", None)
-        if not serialized_size:
-            raise VectorDeserializationFailure("Cannot determine serialized size for vector with subtype %s" % cls.subtype.__name__)
-        indexes = (serialized_size * x for x in range(0, cls.vector_size))
-        return [cls.subtype.deserialize(byts[idx:idx + serialized_size], protocol_version) for idx in indexes]
+        serialized_size = cls.subtype.serial_size()
+        if serialized_size is not None:
+            expected_byte_size = serialized_size * cls.vector_size
+            if len(byts) != expected_byte_size:
+                raise ValueError(
+                    "Expected vector of type {0} and dimension {1} to have serialized size {2}; observed serialized size of {3} instead"\
+                    .format(cls.subtype.typename, cls.vector_size, expected_byte_size, len(byts)))
+            indexes = (serialized_size * x for x in range(0, cls.vector_size))
+            return [cls.subtype.deserialize(byts[idx:idx + serialized_size], protocol_version) for idx in indexes]
+
+        idx = 0
+        rv = []
+        while (len(rv) < cls.vector_size):
+            try:
+                size, bytes_read = uvint_unpack(byts[idx:])
+                idx += bytes_read
+                rv.append(cls.subtype.deserialize(byts[idx:idx + size], protocol_version))
+                idx += size
+            except:
+                raise ValueError("Error reading additional data during vector deserialization after successfully adding {} elements"\
+                .format(len(rv)))
+
+        # If we have any additional data in the serialized vector treat that as an error as well
+        if idx < len(byts):
+            raise ValueError("Additional bytes remaining after vector deserialization completed")
+        return rv
 
     @classmethod
     def serialize(cls, v, protocol_version):
+        v_length = len(v)
+        if cls.vector_size != v_length:
+            raise ValueError(
+                "Expected sequence of size {0} for vector of type {1} and dimension {0}, observed sequence of length {2}"\
+                .format(cls.vector_size, cls.subtype.typename, v_length))
+
+        serialized_size = cls.subtype.serial_size()
         buf = io.BytesIO()
         for item in v:
-            buf.write(cls.subtype.serialize(item, protocol_version))
+            item_bytes = cls.subtype.serialize(item, protocol_version)
+            if serialized_size is None:
+                buf.write(uvint_pack(len(item_bytes)))
+            buf.write(item_bytes)
         return buf.getvalue()
 
     @classmethod
     def cql_parameterized_type(cls):
-        return "%s<%s, %s>" % (cls.typename, cls.subtype.typename, cls.vector_size)
+        return "%s<%s, %s>" % (cls.typename, cls.subtype.cql_parameterized_type(), cls.vector_size)

cassandra/encoder.py

@@ -21,6 +21,7 @@
 log = logging.getLogger(__name__)
 
 from binascii import hexlify
+from decimal import Decimal
 import calendar
 import datetime
 import math
@@ -59,6 +60,7 @@ class Encoder(object):
     def __init__(self):
         self.mapping = {
             float: self.cql_encode_float,
+            Decimal: self.cql_encode_decimal,
             bytearray: self.cql_encode_bytes,
             str: self.cql_encode_str,
             int: self.cql_encode_object,
@@ -217,3 +219,6 @@ def cql_encode_ipaddress(self, val):
         is suitable for ``inet`` type columns.
         """
         return "'%s'" % val.compressed
+
+    def cql_encode_decimal(self, val):
+        return self.cql_encode_float(float(val))

cassandra/marshal.py

@@ -111,7 +111,6 @@ def vints_unpack(term):  # noqa
 
     return tuple(values)
 
-
 def vints_pack(values):
     revbytes = bytearray()
     values = [int(v) for v in values[::-1]]
@@ -143,3 +142,48 @@ def vints_pack(values):
 
     revbytes.reverse()
     return bytes(revbytes)
+
+def uvint_unpack(bytes):
+    first_byte = bytes[0]
+
+    if (first_byte & 128) == 0:
+        return (first_byte,1)
+
+    num_extra_bytes = 8 - (~first_byte & 0xff).bit_length()
+    rv = first_byte & (0xff >> num_extra_bytes)
+    for idx in range(1,num_extra_bytes + 1):
+        new_byte = bytes[idx]
+        rv <<= 8
+        rv |= new_byte & 0xff
+
+    return (rv, num_extra_bytes + 1)
+
+def uvint_pack(val):
+    rv = bytearray()
+    if val < 128:
+        rv.append(val)
+    else:
+        v = val
+        num_extra_bytes = 0
+        num_bits = v.bit_length()
+        # We need to reserve (num_extra_bytes+1) bits in the first byte
+        # ie. with 1 extra byte, the first byte needs to be something like '10XXXXXX' # 2 bits reserved
+        # ie. with 8 extra bytes, the first byte needs to be '11111111'  # 8 bits reserved
+        reserved_bits = num_extra_bytes + 1
+        while num_bits > (8-(reserved_bits)):
+            num_extra_bytes += 1
+            num_bits -= 8
+            reserved_bits = min(num_extra_bytes + 1, 8)
+            rv.append(v & 0xff)
+            v >>= 8
+
+        if num_extra_bytes > 8:
+            raise ValueError('Value %d is too big and cannot be encoded as vint' % val)
+
+        # We can now store the last bits in the first byte
+        n = 8 - num_extra_bytes
+        v |= (0xff >> n << n)
+        rv.append(abs(v))
+
+    rv.reverse()
+    return bytes(rv)

tests/integration/__init__.py

@@ -330,9 +330,10 @@ def _id_and_mark(f):
 greaterthanorequalcass36 = unittest.skipUnless(CASSANDRA_VERSION >= Version('3.6'), 'Cassandra version 3.6 or greater required')
 greaterthanorequalcass3_10 = unittest.skipUnless(CASSANDRA_VERSION >= Version('3.10'), 'Cassandra version 3.10 or greater required')
 greaterthanorequalcass3_11 = unittest.skipUnless(CASSANDRA_VERSION >= Version('3.11'), 'Cassandra version 3.11 or greater required')
-greaterthanorequalcass40 = unittest.skipUnless(CASSANDRA_VERSION >= Version('4.0-a'), 'Cassandra version 4.0 or greater required')
-lessthanorequalcass40 = unittest.skipUnless(CASSANDRA_VERSION <= Version('4.0-a'), 'Cassandra version less or equal to 4.0 required')
-lessthancass40 = unittest.skipUnless(CASSANDRA_VERSION < Version('4.0-a'), 'Cassandra version less than 4.0 required')
+greaterthanorequalcass40 = unittest.skipUnless(CASSANDRA_VERSION >= Version('4.0'), 'Cassandra version 4.0 or greater required')
+greaterthanorequalcass50 = unittest.skipUnless(CASSANDRA_VERSION >= Version('5.0-beta'), 'Cassandra version 5.0 or greater required')
+lessthanorequalcass40 = unittest.skipUnless(CASSANDRA_VERSION <= Version('4.0'), 'Cassandra version less or equal to 4.0 required')
+lessthancass40 = unittest.skipUnless(CASSANDRA_VERSION < Version('4.0'), 'Cassandra version less than 4.0 required')
 lessthancass30 = unittest.skipUnless(CASSANDRA_VERSION < Version('3.0'), 'Cassandra version less then 3.0 required')
 greaterthanorequaldse68 = unittest.skipUnless(DSE_VERSION and DSE_VERSION >= Version('6.8'), "DSE 6.8 or greater required for this test")
 greaterthanorequaldse67 = unittest.skipUnless(DSE_VERSION and DSE_VERSION >= Version('6.7'), "DSE 6.7 or greater required for this test")