Add support for byte and smallint types

cassandra/cqltypes.py

@@ -44,7 +44,7 @@ from uuid import UUID
 import warnings
 
 
-from cassandra.marshal import (int8_pack, int8_unpack,
+from cassandra.marshal import (int8_pack, int8_unpack, int16_pack, int16_unpack,
                                uint16_pack, uint16_unpack, uint32_pack, uint32_unpack,
                                int32_pack, int32_unpack, int64_pack, int64_unpack,
                                float_pack, float_unpack, double_pack, double_unpack,
@@ -424,6 +424,17 @@ class BooleanType(_CassandraType):
     def serialize(truth, protocol_version):
         return int8_pack(truth)
 
+class TinyIntType(_CassandraType):
+    typename = 'tinyint'
+
+    @staticmethod
+    def deserialize(byts, protocol_version):
+        return int8_unpack(byts)
+
+    @staticmethod
+    def serialize(byts, protocol_version):
+        return int8_pack(byts)
+
 
 if six.PY2:
     class AsciiType(_CassandraType):
@@ -650,6 +661,18 @@ class SimpleDateType(_CassandraType):
         return util.Date(days)
 
 
+class SmallIntType(_CassandraType):
+    typename = 'smallint'
+
+    @staticmethod
+    def deserialize(byts, protocol_version):
+        return int16_unpack(byts)
+
+    @staticmethod
+    def serialize(byts, protocol_version):
+        return int16_pack(byts)
+
+
 class TimeType(_CassandraType):
     typename = 'time'
 

cassandra/protocol.py

@@ -34,7 +34,7 @@ from cassandra.cqltypes import (AsciiType, BytesType, BooleanType,
                                 LongType, MapType, SetType, TimeUUIDType,
                                 UTF8Type, UUIDType, UserType,
                                 TupleType, lookup_casstype, SimpleDateType,
-                                TimeType)
+                                TimeType, TinyIntType, SmallIntType)
 from cassandra.policies import WriteType
 
 log = logging.getLogger(__name__)
@@ -533,6 +533,8 @@ class ResultMessage(_MessageType):
         0x0010: InetAddressType,
         0x0011: SimpleDateType,
         0x0012: TimeType,
+        0x0013: SmallIntType,
+        0x0014: TinyIntType,
         0x0020: ListType,
         0x0021: MapType,
         0x0022: SetType,

tests/integration/datatype_utils.py

@@ -60,6 +60,8 @@ def update_datatypes():
     if _cass_version >= (3, 0, 0):
         PRIMITIVE_DATATYPES.append('date')
         PRIMITIVE_DATATYPES.append('time')
+        PRIMITIVE_DATATYPES.append('tinyint')
+        PRIMITIVE_DATATYPES.append('smallint')
 
 
 def get_sample_data():
@@ -117,6 +119,12 @@ def get_sample_data():
         elif datatype == 'time':
             sample_data[datatype] = time(16, 47, 25, 7)
 
+        elif datatype == 'tinyint':
+            sample_data[datatype] = 123
+
+        elif datatype == 'smallint':
+            sample_data[datatype] = 32523
+
         else:
             raise Exception("Missing handling of {0}".format(datatype))
 

tests/unit/test_marshalling.py

@@ -81,7 +81,11 @@ marshalled_value_pairs = (
     (b'\x00\x01\x00\x10\xafYC\xa3\xea<\x11\xe1\xabc\xc4,\x03"y\xf0', 'ListType(TimeUUIDType)', [UUID(bytes=b'\xafYC\xa3\xea<\x11\xe1\xabc\xc4,\x03"y\xf0')]),
     (b'\x80\x00\x00\x01', 'SimpleDateType', Date(1)),
     (b'\x7f\xff\xff\xff', 'SimpleDateType', Date('1969-12-31')),
-    (b'\x00\x00\x00\x00\x00\x00\x00\x01', 'TimeType', Time(1))
+    (b'\x00\x00\x00\x00\x00\x00\x00\x01', 'TimeType', Time(1)),
+    (b'\x7f', 'TinyIntType', 127),
+    (b'\xff\xff\xff\x80', 'TinyIntType', -128),
+    (b'\xff\xff\x80\x00', 'SmallIntType', 32767),
+    (b'\xff\xff\x80\x00', 'SmallIntType', -32768)
 )
 
 ordered_map_value = OrderedMapSerializedKey(UTF8Type, 2)

tests/unit/test_types.py

@@ -26,7 +26,7 @@ import cassandra
 from cassandra.cqltypes import (BooleanType, lookup_casstype_simple, lookup_casstype,
                                 LongType, DecimalType, SetType, cql_typename,
                                 CassandraType, UTF8Type, parse_casstype_args,
-                                SimpleDateType, TimeType,
+                                SimpleDateType, TimeType, TinyIntType, SmallIntType,
                                 EmptyValue, _CassandraType, DateType, int64_pack)
 from cassandra.encoder import cql_quote
 from cassandra.protocol import (write_string, read_longstring, write_stringmap,
@@ -55,6 +55,8 @@ class TypeTests(unittest.TestCase):
         self.assertEqual(lookup_casstype_simple('UTF8Type'), cassandra.cqltypes.UTF8Type)
         self.assertEqual(lookup_casstype_simple('DateType'), cassandra.cqltypes.DateType)
         self.assertEqual(lookup_casstype_simple('SimpleDateType'), cassandra.cqltypes.SimpleDateType)
+        self.assertEqual(lookup_casstype_simple('TinyIntType'), cassandra.cqltypes.TinyIntType)
+        self.assertEqual(lookup_casstype_simple('SmallIntType'), cassandra.cqltypes.SmallIntType)
         self.assertEqual(lookup_casstype_simple('TimeUUIDType'), cassandra.cqltypes.TimeUUIDType)
         self.assertEqual(lookup_casstype_simple('TimeType'), cassandra.cqltypes.TimeType)
         self.assertEqual(lookup_casstype_simple('UUIDType'), cassandra.cqltypes.UUIDType)
@@ -87,6 +89,8 @@ class TypeTests(unittest.TestCase):
         self.assertEqual(lookup_casstype('UTF8Type'), cassandra.cqltypes.UTF8Type)
         self.assertEqual(lookup_casstype('DateType'), cassandra.cqltypes.DateType)
         self.assertEqual(lookup_casstype('TimeType'), cassandra.cqltypes.TimeType)
+        self.assertEqual(lookup_casstype('TinyIntType'), cassandra.cqltypes.TinyIntType)
+        self.assertEqual(lookup_casstype('SmallIntType'), cassandra.cqltypes.SmallIntType)
         self.assertEqual(lookup_casstype('TimeUUIDType'), cassandra.cqltypes.TimeUUIDType)
         self.assertEqual(lookup_casstype('UUIDType'), cassandra.cqltypes.UUIDType)
         self.assertEqual(lookup_casstype('IntegerType'), cassandra.cqltypes.IntegerType)