Repository: Parsely/pykafka
Pull request #750 (2017-12-06 00:00:59)

Topics APIs
This pull request addresses #722 by implementing protocol support for the create/delete topics APIs. It also adds an interface to the `cli` module that uses these APIs.


Patch:
diff --git a/pykafka/broker.py b/pykafka/broker.py
index 0ee8f41..8becad8 100644
--- a/pykafka/broker.py
+++ b/pykafka/broker.py
@@ -30,7 +30,8 @@ from .protocol import (
     SyncGroupRequest, SyncGroupResponse, HeartbeatRequest, HeartbeatResponse,
     LeaveGroupRequest, LeaveGroupResponse, ListGroupsRequest, ListGroupsResponse,
     DescribeGroupsRequest, DescribeGroupsResponse, ApiVersionsRequest,
-    ApiVersionsResponse)
+    ApiVersionsResponse, CreateTopicsRequest, CreateTopicsResponse, DeleteTopicsRequest,
+    DeleteTopicsResponse)
 from .utils.compat import range, iteritems, get_bytes
 
 log = logging.getLogger(__name__)
@@ -552,8 +553,38 @@ class Broker(object):
         future = self._req_handler.request(DescribeGroupsRequest(group_ids))
         return future.get(DescribeGroupsResponse)
 
+    ############################
+    # Create/Delete Topics API #
+    ############################
+    @_check_handler
+    def create_topics(self, topic_reqs, timeout=0):
+        """Create topics via the Topic Creation API
+
+        :param topic_reqs: The topic creation requests to issue
+        :type topics: Iterable of :class:`pykafka.protocol.CreateTopicRequest`
+        :param timeout: The time in ms to wait for a topic to be completely created.
+            Values <= 0 will trigger topic creation and return immediately.
+        :type timeout: int
+        """
+        future = self._req_handler.request(CreateTopicsRequest(topic_reqs,
+                                                               timeout=timeout))
+        return future.get(CreateTopicsResponse)
+
+    @_check_handler
+    def delete_topics(self, topics, timeout=0):
+        """Delete topics viaa the Topic Deletion API
+
+        :param topics: The names of the topics to delete
+        :type topics: Iterable of str
+        :param timeout: The time in ms to wait for a topic to be completely deleted.
+            Values <= 0 will trigger topic deletion and return immediately.
+        :type timeout: int
+        """
+        future = self._req_handler.request(DeleteTopicsRequest(topics, timeout=timeout))
+        return future.get(DeleteTopicsResponse)
+
     @_check_handler
     def fetch_api_versions(self):
-        """Send an ApiVersionsRequest"""
+        """Fetch supported API versions from this broker"""
         future = self._req_handler.request(ApiVersionsRequest())
         return future.get(ApiVersionsResponse.get_version_impl(self._api_versions))
diff --git a/pykafka/cli/kafka_tools.py b/pykafka/cli/kafka_tools.py
index 35c6c3c..a83ec7f 100644
--- a/pykafka/cli/kafka_tools.py
+++ b/pykafka/cli/kafka_tools.py
@@ -5,12 +5,13 @@ import calendar
 import datetime as dt
 import sys
 import time
+from pkg_resources import parse_version
 
 import tabulate
 
 import pykafka
 from pykafka.common import OffsetType
-from pykafka.protocol import PartitionOffsetCommitRequest
+from pykafka.protocol import PartitionOffsetCommitRequest, CreateTopicRequest
 from pykafka.utils.compat import PY3, iteritems
 
 #
@@ -261,6 +262,24 @@ def reset_offsets(client, args):
     )
 
 
+def create_topic(client, args):
+    if parse_version(args.broker_version) < parse_version('0.10.0'):
+        raise ValueError("The topic creation API is not usable on brokers older than "
+                         "0.10.0. Use --broker_version to specify the version")
+    topic_req = CreateTopicRequest(args.topic, args.num_partitions,
+                                   args.replication_factor,
+                                   args.replica_assignment or [],
+                                   args.config_entries or [])
+    client.cluster.controller_broker.create_topics([topic_req], args.timeout)
+
+
+def delete_topic(client, args):
+    if parse_version(args.broker_version) < parse_version('0.10.0'):
+        raise ValueError("The topic deletoin API is not usable on brokers older than "
+                         "0.10.0. Use --broker_version to specify the version")
+    client.cluster.controller_broker.delete_topics([args.topic], args.timeout)
+
+
 def _encode_utf8(string):
     """Converts argument to UTF-8-encoded bytes.
 
@@ -286,6 +305,13 @@ def _add_limit(parser):
                         type=int, default=10)
 
 
+def _add_timeout(parser):
+    parser.add_argument('-t', '--timeout',
+                        help='Time in ms to wait for the operation to complete'
+                             '(default: %(default)s)',
+                        type=int, default=5000)
+
+
 def _add_offset(parser):
     """Add offset to arg parser."""
     parser.add_argument('offset',
@@ -320,6 +346,11 @@ def _get_arg_parser():
                         dest='host',
                         help='host:port of any Kafka broker. '
                              '[default: localhost:9092]')
+    output.add_argument('-o', '--broker_version',
+                        required=False,
+                        default='0.9.0',
+                        dest='broker_version',
+                        help="The version string of the broker with which to communicate")
 
     subparsers = output.add_subparsers(help='Commands', dest='command')
 
@@ -383,6 +414,41 @@ def _get_arg_parser():
     _add_consumer_group(parser)
     _add_offset(parser)
 
+    parser = subparsers.add_parser(
+        'create_topic',
+        help='Create a topic'
+    )
+    parser.set_defaults(func=create_topic)
+    _add_topic(parser)
+    _add_timeout(parser)
+    parser.add_argument('-p', '--num_partitions',
+                        help='Number of partitions to be created. -1 indicates unset. '
+                             '(default: %(default)s)',
+                        type=int, default=1)
+    parser.add_argument('-r', '--replication_factor',
+                        help='Replication factor for the topic. -1 indicates unset. '
+                             '(default: %(default)s)',
+                        type=int, default=1)
+    parser.add_argument('-a', '--replica_assignment',
+                        help='Replica assignment among kafka brokers for this topic '
+                             'partitions. If this is set num_partitions and '
+                             'replication_factor must be unset. Represent as a JSON '
+                             'object with partition IDs as keys as lists of node IDs '
+                             'as values',
+                        type=_encode_utf8)
+    parser.add_argument('-c', '--config_entries',
+                        help='Topic level configuration for topic to be set. Represent '
+                             'as a JSON object.',
+                        type=_encode_utf8)
+
+    parser = subparsers.add_parser(
+        'delete_topic',
+        help='Delete a topic'
+    )
+    parser.set_defaults(func=delete_topic)
+    _add_topic(parser)
+    _add_timeout(parser)
+
     return output
 
 
@@ -390,7 +456,7 @@ def main():
     parser = _get_arg_parser()
     args = parser.parse_args()
     if args.command:
-        client = pykafka.KafkaClient(hosts=args.host)
+        client = pykafka.KafkaClient(hosts=args.host, broker_version=args.broker_version)
         args.func(client, args)
     else:
         parser.print_help()
diff --git a/pykafka/protocol.py b/pykafka/protocol.py
index ceabac6..0bcb29f 100644
--- a/pykafka/protocol.py
+++ b/pykafka/protocol.py
@@ -2123,6 +2123,201 @@ class DescribeGroupsResponse(Response):
             self.groups[group.group_id] = group
 
 
+_CreateTopicRequest = namedtuple(
+    'CreateTopicRequest',
+    ['topic_name', 'num_partitions', 'replication_factor', 'replica_assignment',
+     'config_entries']
+)
+
+
+class CreateTopicRequest(_CreateTopicRequest):
+    def __new__(cls,
+                topic_name,
+                num_partitions,
+                replication_factor,
+                replica_assignment,
+                config_entries):
+        return super(CreateTopicRequest, cls).__new__(
+            cls, topic_name, num_partitions, replication_factor, replica_assignment,
+            config_entries)
+
+
+class CreateTopicsRequest(Request):
+    """A create topics request
+
+    Specification::
+
+    CreateTopics Request (Version: 0) => [create_topic_requests] timeout
+        create_topic_requests => topic num_partitions replication_factor [replica_assignment] [config_entries]
+            topic => STRING
+            num_partitions => INT32
+            replication_factor => INT16
+            replica_assignment => partition [replicas]
+                partition => INT32
+                replicas => INT32
+            config_entries => config_name config_value
+                config_name => STRING
+                config_value => NULLABLE_STRING
+        timeout => INT32
+    """
+    API_KEY = 19
+
+    def __init__(self, topic_requests, timeout=0):
+        self.topic_requests = topic_requests
+        self.timeout = timeout
+
+    def __len__(self):
+        """Length of the serialized message, in bytes"""
+        # header + len(topic_reqs)
+        size = self.HEADER_LEN + 4
+        for topic_req in self.topic_requests:
+            # len(topic_name) + topic_name
+            size += 2 + len(topic_req.topic_name)
+            # num_partitions + replication_factor + len(replica_assignment)
+            size += 4 + 2 + 4
+            for partition, replicas in topic_req.replica_assignment:
+                # partition + len(replicas) + replicas
+                size += 4 + 4 + 4 * len(replicas)
+            # len(config_entries)
+            size += 4
+            for config_name, config_value in topic_req.config_entries:
+                # len(config_name) + config_name + len(config_value) + config_value
+                size += 2 + len(config_name) + 2 + len(config_value)
+        # timeout
+        size += 4
+        return size
+
+    def get_bytes(self):
+        """Create a new create topics request"""
+        output = bytearray(len(self))
+        self._write_header(output)
+        offset = self.HEADER_LEN
+        fmt = '!i'
+        struct.pack_into(fmt, output, offset, len(self.topic_requests))
+        offset += struct.calcsize(fmt)
+        for topic_req in self.topic_requests:
+            fmt = '!h%dsihi' % len(topic_req.topic_name)
+            struct.pack_into(fmt, output, offset, len(topic_req.topic_name),
+                             topic_req.topic_name, topic_req.num_partitions,
+                             topic_req.replication_factor,
+                             len(topic_req.replica_assignment))
+            offset += struct.calcsize(fmt)
+            for partition, replicas in topic_req.replica_assignment:
+                fmt = '!ii'
+                struct.pack_into(fmt, output, offset, partition, len(replicas))
+                offset += struct.calcsize(fmt)
+                for replica in replicas:
+                    fmt = '!i'
+                    struct.pack_into(fmt, output, offset, replica)
+                    offset += struct.calcsize(fmt)
+            fmt = '!i'
+            struct.pack_into(fmt, output, offset, len(topic_req.config_entries))
+            offset += struct.calcsize(fmt)
+            for config_name, config_value in topic_req.config_entries:
+                fmt = '!h%dsh%ds' % (len(config_name), len(config_value))
+                struct.pack_into(fmt, output, offset, len(config_name), config_name,
+                                 len(config_value), config_value)
+                offset += struct.calcsize(fmt)
+        fmt = '!i'
+        struct.pack_into(fmt, output, offset, self.timeout)
+        offset += struct.calcsize(fmt)
+        return output
+
+
+class CreateTopicsResponse(Response):
+    """A create topics response
+
+    Specification::
+
+    CreateTopics Response (Version: 0) => [topic_errors]
+        topic_errors => topic error_code
+            topic => STRING
+            error_code => INT16
+    """
+    API_KEY = 19
+
+    def __init__(self, buff):
+        """Deserialize into a new Response
+
+        :param buff: Serialized message
+        :type buff: :class:`bytearray`
+        """
+        fmt = '[Sh]'
+        response = struct_helpers.unpack_from(fmt, buff, 0)
+        for _, error_code in response:
+            if error_code != 0:
+                self.raise_error(error_code, response)
+
+
+class DeleteTopicsRequest(Request):
+    """A delete topics request
+
+    Specification::
+
+    DeleteTopics Request (Version: 0) => [topics] timeout
+        topics => STRING
+        timeout => INT32
+    """
+    API_KEY = 20
+
+    def __init__(self, topics, timeout=0):
+        self.topics = topics
+        self.timeout = timeout
+
+    def __len__(self):
+        """Length of the serialized message, in bytes"""
+        # header + len(topics)
+        size = self.HEADER_LEN + 4
+        for topic in self.topics:
+            # len(topic) + group_id
+            size += 2 + len(topic)
+        # timeout
+        size += 4
+        return size
+
+    def get_bytes(self):
+        """Create a new delete topics request"""
+        output = bytearray(len(self))
+        self._write_header(output)
+        offset = self.HEADER_LEN
+        fmt = '!i'
+        struct.pack_into(fmt, output, offset, len(self.topics))
+        offset += struct.calcsize(fmt)
+        for topic in self.topics:
+            fmt = '!h%ds' % len(topic)
+            struct.pack_into(fmt, output, offset, len(topic), topic)
+            offset += struct.calcsize(fmt)
+        fmt = '!i'
+        struct.pack_into(fmt, output, offset, self.timeout)
+        offset += struct.calcsize(fmt)
+        return output
+
+
+class DeleteTopicsResponse(Response):
+    """A delete topics response
+
+    Specification::
+
+    DeleteTopics Response (Version: 0) => [topic_error_codes]
+        topic_error_codes => topic error_code
+            topic => STRING
+            error_code => INT16
+    """
+    API_KEY = 20
+
+    def __init__(self, buff):
+        """Deserialize into a new Response
+
+        :param buff: Serialized message
+        :type buff: :class:`bytearray`
+        """
+        fmt = '[Sh]'
+        response = struct_helpers.unpack_from(fmt, buff, 0)
+        for _, error_code in response:
+            if error_code != 0:
+                self.raise_error(error_code, response)
+
+
 class ApiVersionsRequest(Request):
     """An api versions request
 


Test patch:
diff --git a/tests/pykafka/test_protocol.py b/tests/pykafka/test_protocol.py
index dd1daa4..640c88a 100644
--- a/tests/pykafka/test_protocol.py
+++ b/tests/pykafka/test_protocol.py
@@ -1216,6 +1216,59 @@ class TestAdministrativeAPI(unittest2.TestCase):
         self.assertEqual(assignment.partition_assignment,
                          [(b'testtopic_replicated', [0, 1, 2, 8, 3, 9, 4, 5, 6, 7])])
 
+    def test_create_topics_request(self):
+        req = protocol.CreateTopicsRequest([protocol.CreateTopicRequest('mycooltopic', 4,
+                                                                        2, [], [])])
+        msg = req.get_bytes()
+        self.assertEqual(
+            msg,
+            bytearray(
+                b'\x00\x00\x004\x00\x13\x00\x00\x00\x00\x00\x00\x00\x07pykafka'  # header
+                b'\x00\x00\x00\x01'  # len(topic_reqs)
+                    b'\x00\x0b'  # len(topic)  # noqa
+                        b'mycooltopic'  # topic
+                    b'\x00\x00\x00\x04'  # num_partitions
+                    b'\x00\x02'  # replication_factor
+                    b'\x00\x00\x00\x00'  # len(replica_assignment)
+                    b'\x00\x00\x00\x00'  # len(config_entries)
+                b'\x00\x00\x00\x00'  # timeout
+            )
+        )
+
+    def test_create_topics_response(self):
+        response = protocol.CreateTopicsResponse(  # noqa
+            bytearray(
+                b'\x00\x00\x00\x01'  # len(topic_errors)
+                    b'\x00\t'  # len(topic) # noqa
+                        b'testtopic'  # topic_name
+                    b'\x00\x00'  # error_code
+            )
+        )
+
+    def test_delete_topics_request(self):
+        req = protocol.DeleteTopicsRequest(['mycooltopic'])
+        msg = req.get_bytes()
+        self.assertEqual(
+            msg,
+            bytearray(
+                b'\x00\x00\x00&\x00\x14\x00\x00\x00\x00\x00\x00\x00\x07pykafka'  # header
+                b'\x00\x00\x00\x01'  # len(topic_error_codes)  # noqa
+                    b'\x00\x0b'  # len(topic)
+                        b'mycooltopic'  # topic
+                b'\x00\x00\x00\x00'  # timeout
+            )
+        )
+
+    def test_delete_topics_response(self):
+        response = protocol.DeleteTopicsResponse(  # noqa
+            bytearray(
+                b'\x00\x00\x00\x01'  # len(topic_errors)
+                    b'\x00\t'  # len(topic) # noqa
+                        b'testtopic'  # topic_name
+                    b'\x00\x00'  # error_code
+            )
+        )
+
     def test_api_versions_request(self):
         req = protocol.ApiVersionsRequest()
         msg = req.get_bytes()