[KIP-396] List offsets implementation (confluentinc#1576)

---------

Co-authored-by: Emanuele Sabellico <[email protected]>
Co-authored-by: Milind L <[email protected]>

Author: 3 people

docs/index.rst

@@ -25,6 +25,7 @@ Supporting classes
     - :ref:`Message <pythonclient_message>`
     - :ref:`TopicPartition <pythonclient_topicpartition>`
     - :ref:`ThrottleEvent <pythonclient_throttleevent>`
+    - :ref:`IsolationLevel <pythonclient_isolation_level>`
     - :ref:`TopicCollection <pythonclient_topic_collection>`
     - :ref:`TopicPartitionInfo <pythonclient_topic_partition_info>`
     - :ref:`Node <pythonclient_node>`
@@ -61,6 +62,8 @@ Supporting classes
        - :ref:`UserScramCredentialAlteration <pythonclient_user_scram_credential_alteration>`
        - :ref:`UserScramCredentialUpsertion <pythonclient_user_scram_credential_upsertion>`
        - :ref:`UserScramCredentialDeletion <pythonclient_user_scram_credential_deletion>`
+       - :ref:`OffsetSpec <pythonclient_offset_spec>`
+       - :ref:`ListOffsetsResultInfo <pythonclient_list_offsets_result_info>`
        - :ref:`TopicDescription <pythonclient_topic_description>`
        - :ref:`DescribeClusterResult <pythonclient_describe_cluster_result>`
        - :ref:`BrokerMetadata <pythonclient_broker_metadata>`
@@ -267,6 +270,24 @@ UserScramCredentialDeletion
 .. autoclass:: confluent_kafka.admin.UserScramCredentialDeletion
    :members:
 
+.. _pythonclient_offset_spec:
+
+**********
+OffsetSpec
+**********
+
+.. autoclass:: confluent_kafka.admin.OffsetSpec
+   :members:
+
+.. _pythonclient_list_offsets_result_info:
+
+*********************
+ListOffsetsResultInfo
+*********************
+
+.. autoclass:: confluent_kafka.admin.ListOffsetsResultInfo
+   :members:
+
 .. _pythonclient_topic_description:
 
 ****************
@@ -816,6 +837,14 @@ ThrottleEvent
 .. autoclass:: confluent_kafka.ThrottleEvent
    :members:
 
+.. _pythonclient_isolation_level:
+
+**************
+IsolationLevel
+**************
+
+.. autoclass:: confluent_kafka.IsolationLevel
+   :members:
 
 .. _avro_producer:
 

examples/README.md

@@ -12,7 +12,7 @@ The scripts in this directory provide various examples of using Confluent's Pyth
 * [protobuf_producer.py](protobuf_producer.py): Produce Protobuf serialized data using ProtobufSerializer.
 * [protobuf_consumer.py](protobuf_consumer.py): Read Protobuf serialized data using ProtobufDeserializer.
 * [sasl_producer.py](sasl_producer.py):  Demonstrates SASL Authentication.
-* [list_offsets.py](list_offsets.py): List committed offsets and consumer lag for group and topics.
+* [get_watermark_offsets.py](get_watermark_offsets.py): Consumer method for listing committed offsets and consumer lag for group and topics.
 * [oauth_producer.py](oauth_producer.py): Demonstrates OAuth Authentication (client credentials).
 
 Additional examples for [Confluent Cloud](https://www.confluent.io/confluent-cloud/):

examples/adminapi.py

@@ -18,13 +18,15 @@
 # Example use of AdminClient operations.
 
 from confluent_kafka import (KafkaException, ConsumerGroupTopicPartitions,
-                             TopicPartition, ConsumerGroupState, TopicCollection)
+                             TopicPartition, ConsumerGroupState, TopicCollection,
+                             IsolationLevel)
 from confluent_kafka.admin import (AdminClient, NewTopic, NewPartitions, ConfigResource,
                                    ConfigEntry, ConfigSource, AclBinding,
                                    AclBindingFilter, ResourceType, ResourcePatternType,
                                    AclOperation, AclPermissionType, AlterConfigOpType,
                                    ScramMechanism, ScramCredentialInfo,
-                                   UserScramCredentialUpsertion, UserScramCredentialDeletion)
+                                   UserScramCredentialUpsertion, UserScramCredentialDeletion,
+                                   OffsetSpec)
 import sys
 import threading
 import logging
@@ -765,6 +767,57 @@ def example_alter_user_scram_credentials(a, args):
             print("{}: Error: {}".format(username, e))
 
 
+def example_list_offsets(a, args):
+    topic_partition_offsets = {}
+    if len(args) == 0:
+        raise ValueError(
+            "Invalid number of arguments for list offsets, expected at least 1, got 0")
+    i = 1
+    partition_i = 1
+    isolation_level = IsolationLevel[args[0]]
+    while i < len(args):
+        if i + 3 > len(args):
+            raise ValueError(
+                f"Invalid number of arguments for list offsets, partition {partition_i}, expected 3," +
+                f" got {len(args) - i}")
+        topic = args[i]
+        partition = int(args[i+1])
+        topic_partition = TopicPartition(topic, partition)
+
+        if "EARLIEST" == args[i+2]:
+            offset_spec = OffsetSpec.earliest()
+
+        elif "LATEST" == args[i+2]:
+            offset_spec = OffsetSpec.latest()
+
+        elif "MAX_TIMESTAMP" == args[i+2]:
+            offset_spec = OffsetSpec.max_timestamp()
+
+        elif "TIMESTAMP" == args[i+2]:
+            if i + 4 > len(args):
+                raise ValueError(
+                    f"Invalid number of arguments for list offsets, partition {partition_i}, expected 4" +
+                    f", got {len(args) - i}")
+            offset_spec = OffsetSpec.for_timestamp(int(args[i+3]))
+            i += 1
+        else:
+            raise ValueError("Invalid OffsetSpec, must be EARLIEST, LATEST, MAX_TIMESTAMP or TIMESTAMP")
+        topic_partition_offsets[topic_partition] = offset_spec
+        i = i + 3
+        partition_i += 1
+
+    futmap = a.list_offsets(topic_partition_offsets, isolation_level=isolation_level, request_timeout=30)
+    for partition, fut in futmap.items():
+        try:
+            result = fut.result()
+            print("Topicname : {} Partition_Index : {} Offset : {} Timestamp : {}"
+                  .format(partition.topic, partition.partition, result.offset,
+                          result.timestamp))
+        except KafkaException as e:
+            print("Topicname : {} Partition_Index : {} Error : {}"
+                  .format(partition.topic, partition.partition, e))
+
+
 if __name__ == '__main__':
     if len(sys.argv) < 3:
         sys.stderr.write('Usage: %s <bootstrap-brokers> <operation> <args..>\n\n' % sys.argv[0])
@@ -801,6 +854,9 @@ def example_alter_user_scram_credentials(a, args):
                          '<iterations1> <password1> <salt1> ' +
                          '[UPSERT <user2> <mechanism2> <iterations2> ' +
                          ' <password2> <salt2> DELETE <user3> <mechanism3> ..]\n')
+        sys.stderr.write(' list_offsets <isolation_level> <topic1> <partition1> <offset_spec1> ' +
+                         '[<topic2> <partition2> <offset_spec2> ..]\n')
+
         sys.exit(1)
 
     broker = sys.argv[1]
@@ -829,7 +885,8 @@ def example_alter_user_scram_credentials(a, args):
               'list_consumer_group_offsets': example_list_consumer_group_offsets,
               'alter_consumer_group_offsets': example_alter_consumer_group_offsets,
               'describe_user_scram_credentials': example_describe_user_scram_credentials,
-              'alter_user_scram_credentials': example_alter_user_scram_credentials}
+              'alter_user_scram_credentials': example_alter_user_scram_credentials,
+              'list_offsets': example_list_offsets}
 
     if operation not in opsmap:
         sys.stderr.write('Unknown operation: %s\n' % operation)

examples/list_offsets.py examples/get_watermark_offsets.py

@@ -23,7 +23,8 @@
                      ConsumerGroupTopicPartitions,
                      ConsumerGroupState,
                      TopicCollection,
-                     TopicPartitionInfo)
+                     TopicPartitionInfo,
+                     IsolationLevel)
 
 from .cimpl import (Producer,
                     Consumer,
@@ -47,7 +48,8 @@
            'Producer', 'DeserializingConsumer',
            'SerializingProducer', 'TIMESTAMP_CREATE_TIME', 'TIMESTAMP_LOG_APPEND_TIME',
            'TIMESTAMP_NOT_AVAILABLE', 'TopicPartition', 'Node',
-           'ConsumerGroupTopicPartitions', 'ConsumerGroupState', 'Uuid']
+           'ConsumerGroupTopicPartitions', 'ConsumerGroupState', 'Uuid',
+           'IsolationLevel']
 
 __version__ = version()[0]
 

src/confluent_kafka/__init__.py

@@ -132,3 +132,20 @@ def __init__(self, id, leader, replicas, isr):
         self.leader = leader
         self.replicas = replicas
         self.isr = isr
+
+
+class IsolationLevel(Enum):
+    """
+    Enum for Kafka isolation levels.
+
+    Values:
+    -------
+    """
+
+    READ_UNCOMMITTED = cimpl.ISOLATION_LEVEL_READ_UNCOMMITTED  #: Receive all the offsets.
+    READ_COMMITTED = cimpl.ISOLATION_LEVEL_READ_COMMITTED  #: Skip offsets belonging to an aborted transaction.
+
+    def __lt__(self, other):
+        if self.__class__ != other.__class__:
+            return NotImplemented
+        return self.value < other.value

src/confluent_kafka/_model/__init__.py

@@ -46,10 +46,14 @@
                      ScramCredentialInfo,
                      ScramMechanism,
                      UserScramCredentialsDescription)
+
 from ._topic import (TopicDescription)  # noqa: F401
 
 from ._cluster import (DescribeClusterResult)  # noqa: F401
 
+from ._listoffsets import (OffsetSpec,  # noqa: F401
+                           ListOffsetsResultInfo)
+
 from .._model import TopicCollection as _TopicCollection
 
 from ..cimpl import (KafkaException,  # noqa: F401
@@ -71,11 +75,10 @@
                      RESOURCE_BROKER,
                      OFFSET_INVALID)
 
-from confluent_kafka import ConsumerGroupTopicPartitions \
-    as _ConsumerGroupTopicPartitions
-
-from confluent_kafka import ConsumerGroupState \
-    as _ConsumerGroupState
+from confluent_kafka import \
+    ConsumerGroupTopicPartitions as _ConsumerGroupTopicPartitions, \
+    ConsumerGroupState as _ConsumerGroupState, \
+    IsolationLevel as _IsolationLevel
 
 
 try:
@@ -295,6 +298,28 @@ def _make_user_scram_credentials_result(f, futmap):
             for _, fut in futmap.items():
                 fut.set_exception(e)
 
+    @staticmethod
+    def _make_futmap_result(f, futmap):
+        try:
+            results = f.result()
+            len_results = len(results)
+            len_futures = len(futmap)
+            if len(results) != len_futures:
+                raise RuntimeError(
+                    f"Results length {len_results} is different from future-map length {len_futures}")
+            for key, value in results.items():
+                fut = futmap.get(key, None)
+                if fut is None:
+                    raise RuntimeError(
+                        f"Key {key} not found in future-map: {futmap}")
+                if isinstance(value, KafkaError):
+                    fut.set_exception(KafkaException(value))
+                else:
+                    fut.set_result(value)
+        except Exception as e:
+            for _, fut in futmap.items():
+                fut.set_exception(e)
+
     @staticmethod
     def _create_future():
         f = concurrent.futures.Future()
@@ -384,7 +409,6 @@ def _check_list_consumer_group_offsets_request(request):
                     if topic_partition.partition < 0:
                         raise ValueError("Element of 'topic_partitions' must not have negative 'partition' value")
                     if topic_partition.offset != OFFSET_INVALID:
-                        print(topic_partition.offset)
                         raise ValueError("Element of 'topic_partitions' must not have 'offset' value")
 
     @staticmethod
@@ -479,6 +503,34 @@ def _check_alter_user_scram_credentials_request(alterations):
                                 "to be either a UserScramCredentialUpsertion or a " +
                                 "UserScramCredentialDeletion")
 
+    @staticmethod
+    def _check_list_offsets_request(topic_partition_offsets, kwargs):
+        if not isinstance(topic_partition_offsets, dict):
+            raise TypeError("Expected topic_partition_offsets to be " +
+                            "dict of [TopicPartitions,OffsetSpec] for list offsets request")
+
+        for topic_partition, offset_spec in topic_partition_offsets.items():
+            if topic_partition is None:
+                raise TypeError("partition cannot be None")
+            if not isinstance(topic_partition, _TopicPartition):
+                raise TypeError("partition must be a TopicPartition")
+            if topic_partition.topic is None:
+                raise TypeError("partition topic name cannot be None")
+            if not isinstance(topic_partition.topic, string_type):
+                raise TypeError("partition topic name must be string")
+            if not topic_partition.topic:
+                raise ValueError("partition topic name cannot be empty")
+            if topic_partition.partition < 0:
+                raise ValueError("partition index must be non-negative")
+            if offset_spec is None:
+                raise TypeError("OffsetSpec cannot be None")
+            if not isinstance(offset_spec, OffsetSpec):
+                raise TypeError("Value must be a OffsetSpec")
+
+        if 'isolation_level' in kwargs:
+            if not isinstance(kwargs['isolation_level'], _IsolationLevel):
+                raise TypeError("isolation_level argument should be an IsolationLevel")
+
     def create_topics(self, new_topics, **kwargs):
         """
         Create one or more new topics.
@@ -1103,5 +1155,45 @@ def alter_user_scram_credentials(self, alterations, **kwargs):
                                                  AdminClient._make_user_scram_credentials_result)
 
         super(AdminClient, self).alter_user_scram_credentials(alterations, f, **kwargs)
+        return futmap
+
+    def list_offsets(self, topic_partition_offsets, **kwargs):
+        """
+        Enables to find the beginning offset,
+        end offset as well as the offset matching a timestamp
+        or the offset with max timestamp in partitions.
+
+        :param dict([TopicPartition, OffsetSpec]) topic_partition_offsets: Dictionary of
+               TopicPartition objects associated with the corresponding OffsetSpec to query for.
+        :param IsolationLevel isolation_level: The isolation level to use when
+               querying.
+        :param float request_timeout: The overall request timeout in seconds,
+               including broker lookup, request transmission, operation time
+               on broker, and response. Default: `socket.timeout.ms*1000.0`
+
+        :returns: A dict of futures keyed by TopicPartition.
+                  The future result() method returns ListOffsetsResultInfo
+                  raises KafkaException
+
+        :rtype: dict[TopicPartition, future]
+
+        :raises TypeError: Invalid input type.
+        :raises ValueError: Invalid input value.
+        """
+        AdminClient._check_list_offsets_request(topic_partition_offsets, kwargs)
+
+        if 'isolation_level' in kwargs:
+            kwargs['isolation_level_value'] = kwargs['isolation_level'].value
+            del kwargs['isolation_level']
+
+        topic_partition_offsets_list = [
+            _TopicPartition(topic_partition.topic, int(topic_partition.partition),
+                            int(offset_spec._value))
+            for topic_partition, offset_spec in topic_partition_offsets.items()]
+
+        f, futmap = AdminClient._make_futures_v2(topic_partition_offsets_list,
+                                                 _TopicPartition,
+                                                 AdminClient._make_futmap_result)
 
+        super(AdminClient, self).list_offsets(topic_partition_offsets_list, f, **kwargs)
         return futmap