kafka-python 3.0.0__py3-none-any.whl

kafka/benchmarks/varint_speed.py

@@ -0,0 +1,426 @@
+#!/usr/bin/env python
+import pyperf
+
+
+test_data = [
+    (b"\x00", 0),
+    (b"\x01", -1),
+    (b"\x02", 1),
+    (b"\x7E", 63),
+    (b"\x7F", -64),
+    (b"\x80\x01", 64),
+    (b"\x81\x01", -65),
+    (b"\xFE\x7F", 8191),
+    (b"\xFF\x7F", -8192),
+    (b"\x80\x80\x01", 8192),
+    (b"\x81\x80\x01", -8193),
+    (b"\xFE\xFF\x7F", 1048575),
+    (b"\xFF\xFF\x7F", -1048576),
+    (b"\x80\x80\x80\x01", 1048576),
+    (b"\x81\x80\x80\x01", -1048577),
+    (b"\xFE\xFF\xFF\x7F", 134217727),
+    (b"\xFF\xFF\xFF\x7F", -134217728),
+    (b"\x80\x80\x80\x80\x01", 134217728),
+    (b"\x81\x80\x80\x80\x01", -134217729),
+    (b"\xFE\xFF\xFF\xFF\x7F", 17179869183),
+    (b"\xFF\xFF\xFF\xFF\x7F", -17179869184),
+    (b"\x80\x80\x80\x80\x80\x01", 17179869184),
+    (b"\x81\x80\x80\x80\x80\x01", -17179869185),
+    (b"\xFE\xFF\xFF\xFF\xFF\x7F", 2199023255551),
+    (b"\xFF\xFF\xFF\xFF\xFF\x7F", -2199023255552),
+    (b"\x80\x80\x80\x80\x80\x80\x01", 2199023255552),
+    (b"\x81\x80\x80\x80\x80\x80\x01", -2199023255553),
+    (b"\xFE\xFF\xFF\xFF\xFF\xFF\x7F", 281474976710655),
+    (b"\xFF\xFF\xFF\xFF\xFF\xFF\x7F", -281474976710656),
+    (b"\x80\x80\x80\x80\x80\x80\x80\x01", 281474976710656),
+    (b"\x81\x80\x80\x80\x80\x80\x80\x01", -281474976710657),
+    (b"\xFE\xFF\xFF\xFF\xFF\xFF\xFF\x7F", 36028797018963967),
+    (b"\xFF\xFF\xFF\xFF\xFF\xFF\xFF\x7F", -36028797018963968),
+    (b"\x80\x80\x80\x80\x80\x80\x80\x80\x01", 36028797018963968),
+    (b"\x81\x80\x80\x80\x80\x80\x80\x80\x01", -36028797018963969),
+    (b"\xFE\xFF\xFF\xFF\xFF\xFF\xFF\xFF\x7F", 4611686018427387903),
+    (b"\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\x7F", -4611686018427387904),
+    (b"\x80\x80\x80\x80\x80\x80\x80\x80\x80\x01", 4611686018427387904),
+    (b"\x81\x80\x80\x80\x80\x80\x80\x80\x80\x01", -4611686018427387905),
+]
+
+
+BENCH_VALUES_ENC = [
+    60,  # 1 byte
+    -8192,  # 2 bytes
+    1048575,  # 3 bytes
+    134217727,  # 4 bytes
+    -17179869184,  # 5 bytes
+    2199023255551,  # 6 bytes
+]
+
+BENCH_VALUES_DEC = [
+    b"\x7E",  # 1 byte
+    b"\xFF\x7F",  # 2 bytes
+    b"\xFE\xFF\x7F",  # 3 bytes
+    b"\xFF\xFF\xFF\x7F",  # 4 bytes
+    b"\x80\x80\x80\x80\x01",  # 5 bytes
+    b"\xFE\xFF\xFF\xFF\xFF\x7F",  # 6 bytes
+]
+BENCH_VALUES_DEC = list(map(bytearray, BENCH_VALUES_DEC))
+
+
+def _assert_valid_enc(enc_func):
+    for encoded, decoded in test_data:
+        assert enc_func(decoded) == encoded, decoded
+
+
+def _assert_valid_dec(dec_func):
+    for encoded, decoded in test_data:
+        res, pos = dec_func(bytearray(encoded))
+        assert res == decoded, (decoded, res)
+        assert pos == len(encoded), (decoded, pos)
+
+
+def _assert_valid_size(size_func):
+    for encoded, decoded in test_data:
+        assert size_func(decoded) == len(encoded), decoded
+
+
+def encode_varint_1(num):
+    """ Encode an integer to a varint presentation. See
+    https://developers.google.com/protocol-buffers/docs/encoding?csw=1#varints
+    on how those can be produced.
+
+        Arguments:
+            num (int): Value to encode
+
+        Returns:
+            bytearray: Encoded presentation of integer with length from 1 to 10
+                 bytes
+    """
+    # Shift sign to the end of number
+    num = (num << 1) ^ (num >> 63)
+    # Max 10 bytes. We assert those are allocated
+    buf = bytearray(10)
+
+    for i in range(10):
+        # 7 lowest bits from the number and set 8th if we still have pending
+        # bits left to encode
+        buf[i] = num & 0x7f | (0x80 if num > 0x7f else 0)
+        num = num >> 7
+        if num == 0:
+            break
+    else:
+        # Max size of endcoded double is 10 bytes for unsigned values
+        raise ValueError("Out of double range")
+    return buf[:i + 1]
+
+
+def int2byte(i):
+    return bytes((i,))
+
+def encode_varint_2(value):
+    value = (value << 1) ^ (value >> 63)
+
+    bits = value & 0x7f
+    value >>= 7
+    res = b""
+    while value:
+        res += int2byte(0x80 | bits)
+        bits = value & 0x7f
+        value >>= 7
+    return res + int2byte(bits)
+
+
+def encode_varint_3(value, buf):
+    append = buf.append
+    value = (value << 1) ^ (value >> 63)
+
+    bits = value & 0x7f
+    value >>= 7
+    while value:
+        append(0x80 | bits)
+        bits = value & 0x7f
+        value >>= 7
+    append(bits)
+    return value
+
+
+def encode_varint_4(value):
+    value = (value << 1) ^ (value >> 63)
+
+    if value <= 0x7f:  # 1 byte
+        return int2byte(value)
+    if value <= 0x3fff:  # 2 bytes
+        return int2byte(0x80 | (value & 0x7f)) + int2byte(value >> 7)
+    if value <= 0x1fffff:  # 3 bytes
+        return int2byte(0x80 | (value & 0x7f)) + \
+            int2byte(0x80 | ((value >> 7) & 0x7f)) + \
+            int2byte(value >> 14)
+    if value <= 0xfffffff:  # 4 bytes
+        return int2byte(0x80 | (value & 0x7f)) + \
+            int2byte(0x80 | ((value >> 7) & 0x7f)) + \
+            int2byte(0x80 | ((value >> 14) & 0x7f)) + \
+            int2byte(value >> 21)
+    if value <= 0x7ffffffff:  # 5 bytes
+        return int2byte(0x80 | (value & 0x7f)) + \
+            int2byte(0x80 | ((value >> 7) & 0x7f)) + \
+            int2byte(0x80 | ((value >> 14) & 0x7f)) + \
+            int2byte(0x80 | ((value >> 21) & 0x7f)) + \
+            int2byte(value >> 28)
+    else:
+        # Return to general algorithm
+        bits = value & 0x7f
+        value >>= 7
+        res = b""
+        while value:
+            res += int2byte(0x80 | bits)
+            bits = value & 0x7f
+            value >>= 7
+        return res + int2byte(bits)
+
+
+def encode_varint_5(value, buf, pos=0):
+    value = (value << 1) ^ (value >> 63)
+
+    bits = value & 0x7f
+    value >>= 7
+    while value:
+        buf[pos] = 0x80 | bits
+        bits = value & 0x7f
+        value >>= 7
+        pos += 1
+    buf[pos] = bits
+    return pos + 1
+
+def encode_varint_6(value, buf):
+    append = buf.append
+    value = (value << 1) ^ (value >> 63)
+
+    if value <= 0x7f:  # 1 byte
+        append(value)
+        return 1
+    if value <= 0x3fff:  # 2 bytes
+        append(0x80 | (value & 0x7f))
+        append(value >> 7)
+        return 2
+    if value <= 0x1fffff:  # 3 bytes
+        append(0x80 | (value & 0x7f))
+        append(0x80 | ((value >> 7) & 0x7f))
+        append(value >> 14)
+        return 3
+    if value <= 0xfffffff:  # 4 bytes
+        append(0x80 | (value & 0x7f))
+        append(0x80 | ((value >> 7) & 0x7f))
+        append(0x80 | ((value >> 14) & 0x7f))
+        append(value >> 21)
+        return 4
+    if value <= 0x7ffffffff:  # 5 bytes
+        append(0x80 | (value & 0x7f))
+        append(0x80 | ((value >> 7) & 0x7f))
+        append(0x80 | ((value >> 14) & 0x7f))
+        append(0x80 | ((value >> 21) & 0x7f))
+        append(value >> 28)
+        return 5
+    else:
+        # Return to general algorithm
+        bits = value & 0x7f
+        value >>= 7
+        i = 0
+        while value:
+            append(0x80 | bits)
+            bits = value & 0x7f
+            value >>= 7
+            i += 1
+    append(bits)
+    return i
+
+
+def size_of_varint_1(value):
+    """ Number of bytes needed to encode an integer in variable-length format.
+    """
+    value = (value << 1) ^ (value >> 63)
+    res = 0
+    while True:
+        res += 1
+        value = value >> 7
+        if value == 0:
+            break
+    return res
+
+
+def size_of_varint_2(value):
+    """ Number of bytes needed to encode an integer in variable-length format.
+    """
+    value = (value << 1) ^ (value >> 63)
+    if value <= 0x7f:
+        return 1
+    if value <= 0x3fff:
+        return 2
+    if value <= 0x1fffff:
+        return 3
+    if value <= 0xfffffff:
+        return 4
+    if value <= 0x7ffffffff:
+        return 5
+    if value <= 0x3ffffffffff:
+        return 6
+    if value <= 0x1ffffffffffff:
+        return 7
+    if value <= 0xffffffffffffff:
+        return 8
+    if value <= 0x7fffffffffffffff:
+        return 9
+    return 10
+
+
+def _read_byte(memview, pos):
+    """ Read a byte from memoryview as an integer
+
+        Raises:
+            IndexError: if position is out of bounds
+    """
+    return memview[pos]
+
+
+def decode_varint_1(buffer, pos=0):
+    """ Decode an integer from a varint presentation. See
+    https://developers.google.com/protocol-buffers/docs/encoding?csw=1#varints
+    on how those can be produced.
+
+        Arguments:
+            buffer (bytes-like): any object acceptable by ``memoryview``
+            pos (int): optional position to read from
+
+        Returns:
+            (int, int): Decoded int value and next read position
+    """
+    value = 0
+    shift = 0
+    memview = memoryview(buffer)
+    for i in range(pos, pos + 10):
+        try:
+            byte = _read_byte(memview, i)
+        except IndexError:
+            raise ValueError("End of byte stream")
+        if byte & 0x80 != 0:
+            value |= (byte & 0x7f) << shift
+            shift += 7
+        else:
+            value |= byte << shift
+            break
+    else:
+        # Max size of endcoded double is 10 bytes for unsigned values
+        raise ValueError("Out of double range")
+    # Normalize sign
+    return (value >> 1) ^ -(value & 1), i + 1
+
+
+def decode_varint_2(buffer, pos=0):
+    result = 0
+    shift = 0
+    while 1:
+        b = buffer[pos]
+        result |= ((b & 0x7f) << shift)
+        pos += 1
+        if not (b & 0x80):
+            # result = result_type(() & mask)
+            return ((result >> 1) ^ -(result & 1), pos)
+        shift += 7
+        if shift >= 64:
+            raise ValueError("Out of int64 range")
+
+
+def decode_varint_3(buffer, pos=0):
+    result = buffer[pos]
+    if not (result & 0x81):
+        return (result >> 1), pos + 1
+    if not (result & 0x80):
+        return (result >> 1) ^ (~0), pos + 1
+
+    result &= 0x7f
+    pos += 1
+    shift = 7
+    while 1:
+        b = buffer[pos]
+        result |= ((b & 0x7f) << shift)
+        pos += 1
+        if not (b & 0x80):
+            return ((result >> 1) ^ -(result & 1), pos)
+        shift += 7
+        if shift >= 64:
+            raise ValueError("Out of int64 range")
+
+
+if __name__ == '__main__':
+    _assert_valid_enc(encode_varint_1)
+    _assert_valid_enc(encode_varint_2)
+
+    for encoded, decoded in test_data:
+        res = bytearray()
+        encode_varint_3(decoded, res)
+        assert res == encoded
+
+    _assert_valid_enc(encode_varint_4)
+
+    # import dis
+    # dis.dis(encode_varint_4)
+
+    for encoded, decoded in test_data:
+        res = bytearray(10)
+        written = encode_varint_5(decoded, res)
+        assert res[:written] == encoded
+
+    for encoded, decoded in test_data:
+        res = bytearray()
+        encode_varint_6(decoded, res)
+        assert res == encoded
+
+    _assert_valid_size(size_of_varint_1)
+    _assert_valid_size(size_of_varint_2)
+    _assert_valid_dec(decode_varint_1)
+    _assert_valid_dec(decode_varint_2)
+    _assert_valid_dec(decode_varint_3)
+
+    # import dis
+    # dis.dis(decode_varint_3)
+
+    runner = pyperf.Runner()
+    # Encode algorithms returning a bytes result
+    for bench_func in [
+            encode_varint_1,
+            encode_varint_2,
+            encode_varint_4]:
+        for i, value in enumerate(BENCH_VALUES_ENC):
+            runner.bench_func(
+                '{}_{}byte'.format(bench_func.__name__, i + 1),
+                bench_func, value)
+
+    # Encode algorithms writing to the buffer
+    for bench_func in [
+            encode_varint_3,
+            encode_varint_5,
+            encode_varint_6]:
+        for i, value in enumerate(BENCH_VALUES_ENC):
+            fname = bench_func.__name__
+            runner.timeit(
+                '{}_{}byte'.format(fname, i + 1),
+                stmt="{}({}, buffer)".format(fname, value),
+                setup="from __main__ import {}; buffer = bytearray(10)".format(
+                    fname)
+            )
+
+    # Size algorithms
+    for bench_func in [
+            size_of_varint_1,
+            size_of_varint_2]:
+        for i, value in enumerate(BENCH_VALUES_ENC):
+            runner.bench_func(
+                '{}_{}byte'.format(bench_func.__name__, i + 1),
+                bench_func, value)
+
+    # Decode algorithms
+    for bench_func in [
+            decode_varint_1,
+            decode_varint_2,
+            decode_varint_3]:
+        for i, value in enumerate(BENCH_VALUES_DEC):
+            runner.bench_func(
+                '{}_{}byte'.format(bench_func.__name__, i + 1),
+                bench_func, value)

kafka/cli/__init__.py

@@ -0,0 +1,36 @@
+import sys
+
+USAGE = """\
+usage: kafka-python {admin,consumer,producer} [options]
+
+Top-level entry point for the kafka-python command-line tools.
+
+  kafka-python admin     -- administrative client (topics, configs, groups, ...)
+  kafka-python consumer  -- console consumer
+  kafka-python producer  -- console producer
+
+Run `kafka-python <subcommand> --help` for subcommand-specific options.
+"""
+
+
+def run_cli(args=None):
+    if args is None:
+        args = sys.argv[1:]
+
+    if not args or args[0] in ('-h', '--help'):
+        sys.stdout.write(USAGE)
+        return 0 if args else 1
+
+    sub, rest = args[0], args[1:]
+    if sub == 'admin':
+        from kafka.cli.admin import run_cli as sub_run
+    elif sub == 'consumer':
+        from kafka.cli.consumer import run_cli as sub_run
+    elif sub == 'producer':
+        from kafka.cli.producer import run_cli as sub_run
+    else:
+        sys.stderr.write('kafka-python: unknown subcommand: %r\n' % sub)
+        sys.stderr.write(USAGE)
+        return 2
+
+    return sub_run(rest, prog='kafka-python %s' % sub)

kafka/cli/admin/__init__.py

@@ -0,0 +1,117 @@
+import argparse
+import json
+import logging
+from pprint import pprint
+
+from kafka.admin.client import KafkaAdminClient
+from .acls import ACLsCommandGroup
+from .cluster import ClusterCommandGroup
+from .configs import ConfigsCommandGroup
+from .groups import GroupsCommandGroup
+from .partitions import PartitionsCommandGroup
+from .topics import TopicsCommandGroup
+from .transactions import TransactionsCommandGroup
+from .users import UsersCommandGroup
+from ..common import add_common_cli_args, configure_logging, build_connect_kwargs
+from kafka.errors import BrokerResponseError
+
+
+DEFAULT_COMMAND_GROUPS = (
+    ACLsCommandGroup, ClusterCommandGroup, ConfigsCommandGroup,
+    TopicsCommandGroup, PartitionsCommandGroup, GroupsCommandGroup,
+    TransactionsCommandGroup, UsersCommandGroup,
+)
+
+
+def main_parser(prog=None):
+    return build_parser(DEFAULT_COMMAND_GROUPS, prog=prog)
+
+
+def build_parser(groups=(), prog=None):
+    parser = argparse.ArgumentParser(
+        prog=prog or 'python -m kafka.admin',
+        description='Kafka Admin Client',
+    )
+    add_common_cli_args(parser, bootstrap_required=False)
+    parser.add_argument_group('output').add_argument(
+        '--format', type=str, default='raw',
+        help='output format: raw|json')
+    groups_sub = parser.add_subparsers(dest='group', metavar='GROUP', title='Available command groups')
+    for group in groups:
+        group_parser = groups_sub.add_parser(group.GROUP, help=group.HELP)
+        group_parser.set_defaults(group=group_parser) # refcycle..
+        commands_sub = group_parser.add_subparsers(dest='command', metavar='COMMAND', title='Available commands')
+        for cmd in group.COMMANDS:
+            command_parser = commands_sub.add_parser(cmd.COMMAND, help=cmd.HELP)
+            options = command_parser.add_argument_group('command options')
+            cmd.add_arguments(options)
+            command_parser.set_defaults(command=cmd.command)
+    return parser
+
+
+def run_cli(args=None, prog=None):
+    parser = main_parser(prog=prog)
+    config = parser.parse_args(args)
+    if not config.group:
+        parser.print_help()
+        return 1
+    elif not config.command:
+        config.group.print_help()
+        return 1
+    if config.format not in ('raw', 'json'):
+        print(f'Unrecognized format: {config.format}')
+        return 1
+
+    configure_logging(config)
+    logger = logging.getLogger(__name__)
+
+    try:
+        kwargs = build_connect_kwargs(config)
+    except ValueError as exc:
+        parser.print_usage()
+        print(f'{parser.prog}: {exc}')
+        return 1
+    client = KafkaAdminClient(**kwargs)
+
+    try:
+        result = config.command(client, config)
+        if config.format == 'raw':
+            pprint(result)
+        elif config.format == 'json':
+            if hasattr(result, 'to_dict'):
+                result = result.to_dict()
+            print(json.dumps(result))
+        return 0
+    except BrokerResponseError as exc:
+        print(exc)
+        return 1
+    except ValueError as exc:
+        print(exc.args[0])
+        return 1
+    except AttributeError as exc:
+        logger.exception(exc)
+        parser.print_help()
+        return 1
+    except Exception:
+        logger.critical('Error!', exc_info=True)
+        return 1
+    finally:
+        client.close()
+
+
+# Commands TODO:
+    # --dry-run support
+    # --trace ?
+
+    # [cluster]
+    # unregister-broker (UnregisterBroker)
+    # add-raft-voter (AddRaftVoter)
+    # remove-raft-voter (RemoveRaftVoter)
+    # describe-quotas (DescribeClientQuotas)
+    # alter-quotas (AlterClientQuotas)
+
+    # [tokens] *DelegationTokenRequest
+    # create
+    # describe
+    # renew
+    # expire

kafka/cli/admin/acls/__init__.py

@@ -0,0 +1,9 @@
+from .create import CreateACLs
+from .delete import DeleteACLs
+from .describe import DescribeACLs
+
+
+class ACLsCommandGroup:
+    GROUP = 'acls'
+    HELP = 'Manage Kafka ACLs'
+    COMMANDS = [DescribeACLs, CreateACLs, DeleteACLs]

kafka/cli/admin/acls/common.py

@@ -0,0 +1,76 @@
+from kafka.admin._acls import (
+    ACL, ACLFilter, ACLOperation, ACLPermissionType, ACLResourcePatternType,
+    ResourceType, ResourcePattern, ResourcePatternFilter
+)
+
+
+def _enum_names(enum_cls):
+    return [e.name.lower() for e in enum_cls]
+
+
+def add_acl_filter_args(parser):
+    """Add arguments for building an ACLFilter (used by describe and delete)."""
+    parser.add_argument('--principal', type=str, default=None, help='ACL principal (e.g. User:alice)')
+    parser.add_argument('--host', type=str, default=None, help='ACL host (default: match any)')
+    parser.add_argument('--operation', type=str, default='any',
+                        choices=_enum_names(ACLOperation),
+                        help='ACL operation (default: any)')
+    parser.add_argument('--permission-type', type=str, default='any',
+                        choices=_enum_names(ACLPermissionType),
+                        help='ACL permission type (default: any)')
+    parser.add_argument('--resource-type', type=str, default='any',
+                        choices=_enum_names(ResourceType),
+                        help='Resource type (default: any)')
+    parser.add_argument('--resource-name', type=str, default=None, help='Resource name (* wildcards accepted)')
+    parser.add_argument('--pattern-type', type=str, default='any',
+                        choices=_enum_names(ACLResourcePatternType),
+                        help='Resource pattern type (default: any)')
+
+
+def add_acl_args(parser, required=False):
+    """Add arguments for building a concrete ACL (used by create)."""
+    parser.add_argument('--principal', type=str, required=required, help='ACL principal (e.g. User:alice)')
+    parser.add_argument('--host', type=str, default='*', help='ACL host (default: *)')
+    parser.add_argument('--operation', type=str, required=required,
+                        choices=_enum_names(ACLOperation),
+                        help='ACL operation')
+    parser.add_argument('--permission-type', type=str, default='allow',
+                        choices=_enum_names(ACLPermissionType),
+                        help='ACL permission type (default: allow)')
+    parser.add_argument('--resource-type', type=str, required=required,
+                        choices=_enum_names(ResourceType),
+                        help='Resource type')
+    parser.add_argument('--resource-name', type=str, required=required, help='Resource name (* wildcards accepted)')
+    parser.add_argument('--pattern-type', type=str, default='literal',
+                        choices=_enum_names(ACLResourcePatternType),
+                        help='Resource pattern type (default: literal)')
+
+
+def acl_filter_from_args(args):
+    """Build an ACLFilter from parsed CLI arguments."""
+    return ACLFilter(
+        principal=args.principal,
+        host=args.host,
+        operation=ACLOperation[args.operation.upper()],
+        permission_type=ACLPermissionType[args.permission_type.upper()],
+        resource_pattern=ResourcePatternFilter(
+            resource_type=ResourceType[args.resource_type.upper()],
+            resource_name=args.resource_name,
+            pattern_type=ACLResourcePatternType[args.pattern_type.upper()],
+        ),
+    )
+
+
+def acl_from_args(args):
+    """Build an ACL from parsed CLI arguments."""
+    return ACL(
+        principal=args.principal,
+        host=args.host or '*',
+        operation=ACLOperation[args.operation.upper()],
+        permission_type=ACLPermissionType[args.permission_type.upper()],
+        resource_pattern=ResourcePattern(
+            resource_type=ResourceType[args.resource_type.upper()],
+            resource_name=args.resource_name,
+            pattern_type=ACLResourcePatternType[args.pattern_type.upper()],
+        ),
+    )

kafka/cli/admin/acls/create.py

@@ -0,0 +1,19 @@
+from .common import add_acl_args, acl_from_args
+
+
+class CreateACLs:
+    COMMAND = 'create'
+    HELP = 'Create Kafka ACLs'
+
+    @classmethod
+    def add_arguments(cls, parser):
+        add_acl_args(parser, required=True)
+
+    @classmethod
+    def command(cls, client, args):
+        acl = acl_from_args(args)
+        result = client.create_acls([acl])
+        return {
+            'succeeded': [repr(a) for a in result['succeeded']],
+            'failed': [str(e) for e in result['failed']],
+        }