PYTHON-1358 Add a TimeUUID class

cassandra/cqltypes.py

@@ -677,7 +677,7 @@ def my_timestamp(self):
 
     @staticmethod
     def deserialize(byts, protocol_version):
-        return UUID(bytes=byts)
+        return util.TimeUUID(bytes=byts)
 
     @staticmethod
     def serialize(timeuuid, protocol_version):

cassandra/util.py

@@ -116,6 +116,35 @@ def max_uuid_from_time(timestamp):
     """
     return uuid_from_time(timestamp, 0x7f7f7f7f7f7f, 0x3f7f)  # Max signed bytes (0x7f = 127)
 
+class TimeUUID(uuid.UUID):
+
+    @property
+    def lsb(self) -> int:
+        return int.from_bytes(self.bytes[8:16], 'big', signed=True)
+
+    def __eq__(self, other):
+        if isinstance(other, TimeUUID):
+            return self.time == other.time and self.lsb == other.lsb 
+        if isinstance(other, uuid.UUID):
+            return self.time == other.time and self.bytes[8:16] == other.bytes[8:16] 
+        return NotImplemented
+
+    def __lt__(self, other):
+        if isinstance(other, TimeUUID):
+            if self.time != other.time:
+                return self.time < other.time
+            else:
+                return self.lsb < other.lsb
+        return NotImplemented
+
+    def __gt__(self, other):
+        return not self <= other
+
+    def __le__(self, other):
+        return self < other or self == other
+
+    def __ge__(self, other):
+        return not self < other
 
 def uuid_from_time(time_arg, node=None, clock_seq=None):
     """
@@ -166,7 +195,7 @@ def uuid_from_time(time_arg, node=None, clock_seq=None):
     if node is None:
         node = random.getrandbits(48)
 
-    return uuid.UUID(fields=(time_low, time_mid, time_hi_version,
+    return TimeUUID(fields=(time_low, time_mid, time_hi_version,
                              clock_seq_hi_variant, clock_seq_low, node), version=1)
 
 LOWEST_TIME_UUID = uuid.UUID('00000000-0000-1000-8080-808080808080')

tests/unit/test_types.py

@@ -946,3 +946,38 @@ def test_token_order(self):
         tokens_equal = [Token(1), Token(1)]
         check_sequence_consistency(self, tokens)
         check_sequence_consistency(self, tokens_equal, equal=True)
+
+    def test_timeuuid_comparison(self):
+        """
+        Test TimeUUID is compared just as Cassandra does
+        If first compares the timestamps,
+        then the last 64 bits interpreted as a signed int
+
+        @jira_ticket PYTHON-1358
+        @expected_result The TimeUUID from 2023-06-19 23:49:56.990000+0000
+        should be smaller than the one from 2023-06-19 23:49:59.961000+0000
+
+        @test_category data_types
+        """
+        time1 = datetime.datetime(2023, 6, 19, 23, 49, 56, 990000)
+        time2 = datetime.datetime(2023, 6, 19, 23, 49, 59, 961000)
+        timeuuid1 = util.uuid_from_time(time1)
+        timeuuid2 = util.uuid_from_time(time2)
+        self.assertTrue(timeuuid1 < timeuuid2)
+
+    def test_timeuuid_with_same_time(self):
+        """
+        Test TimeUUID is compared just as Cassandra does
+        If first compares the timestamps,
+        then the last 64 bits interpreted as a signed int
+
+        @jira_ticket PYTHON-1358
+        @expected_result The min TimeUUID should be smaller than 
+        the max TimeUUID with the same timestamp
+
+        @test_category data_types
+        """
+        timestamp = 1693776594
+        min_timeuuid = util.min_uuid_from_time(timestamp)
+        max_timeuuid = util.max_uuid_from_time(timestamp)
+        self.assertTrue(min_timeuuid < max_timeuuid)