diff --git a/cassandra/cqltypes.py b/cassandra/cqltypes.py index 5b95d4c0..48cbc198 100644 --- a/cassandra/cqltypes.py +++ b/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' diff --git a/cassandra/protocol.py b/cassandra/protocol.py index bd929dad..031501be 100644 --- a/cassandra/protocol.py +++ b/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, diff --git a/tests/integration/datatype_utils.py b/tests/integration/datatype_utils.py index fb437d75..a7d8aeb1 100644 --- a/tests/integration/datatype_utils.py +++ b/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)) diff --git a/tests/unit/test_marshalling.py b/tests/unit/test_marshalling.py index bb9a4068..eeefb0f1 100644 --- a/tests/unit/test_marshalling.py +++ b/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) diff --git a/tests/unit/test_types.py b/tests/unit/test_types.py index bc0eac2d..43863803 100644 --- a/tests/unit/test_types.py +++ b/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)