<s> from __future__ import absolute_import <EOL> import logging <EOL> import struct <EOL> import six <EOL> from six . moves import xrange <EOL> import kafka . protocol . commit <EOL> import kafka . protocol . fetch <EOL> import kafka . protocol . message <EOL> import kafka . protocol . metadata <EOL> import kafka . protocol . offset <EOL> import kafka . protocol . produce <EOL> import kafka . structs <EOL> from kafka . codec import ( <EOL> gzip_encode , gzip_decode , snappy_encode , snappy_decode ) <EOL> from kafka . errors import ProtocolError , ChecksumError , UnsupportedCodecError <EOL> from kafka . structs import ConsumerMetadataResponse <EOL> from kafka . util import ( <EOL> crc32 , read_short_string , read_int_string , relative_unpack , <EOL> write_short_string , write_int_string , group_by_topic_and_partition ) <EOL> log = logging . getLogger ( __name__ ) <EOL> ATTRIBUTE_CODEC_MASK = <NUM_LIT> <EOL> CODEC_NONE = <NUM_LIT> <EOL> CODEC_GZIP = <NUM_LIT> <EOL> CODEC_SNAPPY = <NUM_LIT> <EOL> ALL_CODECS = ( CODEC_NONE , CODEC_GZIP , CODEC_SNAPPY ) <EOL> class KafkaProtocol ( object ) : <EOL> """<STR_LIT>""" <EOL> PRODUCE_KEY = <NUM_LIT:0> <EOL> FETCH_KEY = <NUM_LIT:1> <EOL> OFFSET_KEY = <NUM_LIT:2> <EOL> METADATA_KEY = <NUM_LIT:3> <EOL> OFFSET_COMMIT_KEY = <NUM_LIT:8> <EOL> OFFSET_FETCH_KEY = <NUM_LIT:9> <EOL> CONSUMER_METADATA_KEY = <NUM_LIT:10> <EOL> @ classmethod <EOL> def _encode_message_header ( cls , client_id , correlation_id , request_key , <EOL> version = <NUM_LIT:0> ) : <EOL> """<STR_LIT>""" <EOL> return struct . pack ( '<STR_LIT>' % len ( client_id ) , <EOL> request_key , <EOL> version , <EOL> correlation_id , <EOL> len ( client_id ) , <EOL> client_id ) <EOL> @ classmethod <EOL> def _encode_message_set ( cls , messages ) : <EOL> """<STR_LIT>""" <EOL> message_set = [ ] <EOL> for message in messages : <EOL> encoded_message = KafkaProtocol . _encode_message ( message ) <EOL> message_set . append ( struct . pack ( '<STR_LIT>' % len ( encoded_message ) , <NUM_LIT:0> , <EOL> len ( encoded_message ) , <EOL> encoded_message ) ) <EOL> return b'<STR_LIT>' . join ( message_set ) <EOL> @ classmethod <EOL> def _encode_message ( cls , message ) : <EOL> """<STR_LIT>""" <EOL> if message . magic == <NUM_LIT:0> : <EOL> msg = b'<STR_LIT>' . join ( [ <EOL> struct . pack ( '<STR_LIT>' , message . magic , message . attributes ) , <EOL> write_int_string ( message . key ) , <EOL> write_int_string ( message . value ) <EOL> ] ) <EOL> crc = crc32 ( msg ) <EOL> msg = struct . pack ( '<STR_LIT>' % len ( msg ) , crc , msg ) <EOL> else : <EOL> raise ProtocolError ( "<STR_LIT>" % message . magic ) <EOL> return msg <EOL> @ classmethod <EOL> def encode_produce_request ( cls , payloads = ( ) , acks = <NUM_LIT:1> , timeout = <NUM_LIT:1000> ) : <EOL> """<STR_LIT>""" <EOL> if acks not in ( <NUM_LIT:1> , <NUM_LIT:0> , - <NUM_LIT:1> ) : <EOL> raise ValueError ( '<STR_LIT>' % acks ) <EOL> return kafka . protocol . produce . ProduceRequest [ <NUM_LIT:0> ] ( <EOL> required_acks = acks , <EOL> timeout = timeout , <EOL> topics = [ ( <EOL> topic , <EOL> [ ( <EOL> partition , <EOL> [ ( <NUM_LIT:0> , <NUM_LIT:0> , kafka . protocol . message . Message ( msg . value , key = msg . key , <EOL> magic = msg . magic , <EOL> attributes = msg . attributes ) ) <EOL> for msg in payload . messages ] ) <EOL> for partition , payload in topic_payloads . items ( ) ] ) <EOL> for topic , topic_payloads in group_by_topic_and_partition ( payloads ) . items ( ) ] ) <EOL> @ classmethod <EOL> def decode_produce_response ( cls , response ) : <EOL> """<STR_LIT>""" <EOL> return [ <EOL> kafka . structs . ProduceResponsePayload ( topic , partition , error , offset ) <EOL> for topic , partitions in response . topics <EOL> for partition , error , offset in partitions <EOL> ] <EOL> @ classmethod <EOL> def encode_fetch_request ( cls , payloads = ( ) , max_wait_time = <NUM_LIT:100> , min_bytes = <NUM_LIT> ) : <EOL> """<STR_LIT>""" <EOL> return kafka . protocol . fetch . FetchRequest [ <NUM_LIT:0> ] ( <EOL> replica_id = - <NUM_LIT:1> , <EOL> max_wait_time = max_wait_time , <EOL> min_bytes = min_bytes , <EOL> topics = [ ( <EOL> topic , <EOL> [ ( <EOL> partition , <EOL> payload . offset , <EOL> payload . max_bytes ) <EOL> for partition , payload in topic_payloads . items ( ) ] ) <EOL> for topic , topic_payloads in group_by_topic_and_partition ( payloads ) . items ( ) ] ) <EOL> @ classmethod <EOL> def decode_fetch_response ( cls , response ) : <EOL> """<STR_LIT>""" <EOL> return [ <EOL> kafka . structs . FetchResponsePayload ( <EOL> topic , partition , error , highwater_offset , [ <EOL> kafka . structs . OffsetAndMessage ( offset , message ) <EOL> for offset , _ , message in messages ] ) <EOL> for topic , partitions in response . topics <EOL> for partition , error , highwater_offset , messages in partitions <EOL> ] <EOL> @ classmethod <EOL> def encode_offset_request ( cls , payloads = ( ) ) : <EOL> return kafka . protocol . offset . OffsetRequest [ <NUM_LIT:0> ] ( <EOL> replica_id = - <NUM_LIT:1> , <EOL> topics = [ ( <EOL> topic , <EOL> [ ( <EOL> partition , <EOL> payload . time , <EOL> payload . max_offsets ) <EOL> for partition , payload in six . iteritems ( topic_payloads ) ] ) <EOL> for topic , topic_payloads in six . iteritems ( group_by_topic_and_partition ( payloads ) ) ] ) <EOL> @ classmethod <EOL> def decode_offset_response ( cls , response ) : <EOL> """<STR_LIT>""" <EOL> return [ <EOL> kafka . structs . OffsetResponsePayload ( topic , partition , error , tuple ( offsets ) ) <EOL> for topic , partitions in response . topics <EOL> for partition , error , offsets in partitions <EOL> ] <EOL> @ classmethod <EOL> def encode_metadata_request ( cls , topics = ( ) , payloads = None ) : <EOL> """<STR_LIT>""" <EOL> if payloads is not None : <EOL> topics = payloads <EOL> return kafka . protocol . metadata . MetadataRequest [ <NUM_LIT:0> ] ( topics ) <EOL> @ classmethod <EOL> def decode_metadata_response ( cls , response ) : <EOL> return response <EOL> @ classmethod <EOL> def encode_consumer_metadata_request ( cls , client_id , correlation_id , payloads ) : <EOL> """<STR_LIT>""" <EOL> message = [ ] <EOL> message . append ( cls . _encode_message_header ( client_id , correlation_id , <EOL> KafkaProtocol . CONSUMER_METADATA_KEY ) ) <EOL> message . append ( struct . pack ( '<STR_LIT>' % len ( payloads ) , len ( payloads ) , payloads ) ) <EOL> msg = b'<STR_LIT>' . join ( message ) <EOL> return write_int_string ( msg ) <EOL> @ classmethod <EOL> def decode_consumer_metadata_response ( cls , data ) : <EOL> """<STR_LIT>""" <EOL> ( ( correlation_id , error , nodeId ) , cur ) = relative_unpack ( '<STR_LIT>' , data , <NUM_LIT:0> ) <EOL> ( host , cur ) = read_short_string ( data , cur ) <EOL> ( ( port , ) , cur ) = relative_unpack ( '<STR_LIT>' , data , cur ) <EOL> return ConsumerMetadataResponse ( error , nodeId , host , port ) <EOL> @ classmethod <EOL> def encode_offset_commit_request ( cls , group , payloads ) : <EOL> """<STR_LIT>""" <EOL> return kafka . protocol . commit . OffsetCommitRequest [ <NUM_LIT:0> ] ( <EOL> consumer_group = group , <EOL> topics = [ ( <EOL> topic , <EOL> [ ( <EOL> partition , <EOL> payload . offset , <EOL> payload . metadata ) <EOL> for partition , payload in six . iteritems ( topic_payloads ) ] ) <EOL> for topic , topic_payloads in six . iteritems ( group_by_topic_and_partition ( payloads ) ) ] ) <EOL> @ classmethod <EOL> def decode_offset_commit_response ( cls , response ) : <EOL> """<STR_LIT>""" <EOL> return [ <EOL> kafka . structs . OffsetCommitResponsePayload ( topic , partition , error ) <EOL> for topic , partitions in response . topics <EOL> for partition , error in partitions <EOL> ] <EOL> @ classmethod <EOL> def encode_offset_fetch_request ( cls , group , payloads , from_kafka = False ) : <EOL> """<STR_LIT>""" <EOL> version = <NUM_LIT:1> if from_kafka else <NUM_LIT:0> <EOL> return kafka . protocol . commit . OffsetFetchRequest [ version ] ( <EOL> consumer_group = group , <EOL> topics = [ ( <EOL> topic , <EOL> list ( topic_payloads . keys ( ) ) ) <EOL> for topic , topic_payloads in six . iteritems ( group_by_topic_and_partition ( payloads ) ) ] ) <EOL> @ classmethod <EOL> def decode_offset_fetch_response ( cls , response ) : <EOL> """<STR_LIT>""" <EOL> return [ <EOL> kafka . structs . OffsetFetchResponsePayload ( <EOL> topic , partition , offset , metadata , error <EOL> ) <EOL> for topic , partitions in response . topics <EOL> for partition , offset , metadata , error in partitions <EOL> ] <EOL> def create_message ( payload , key = None ) : <EOL> """<STR_LIT>""" <EOL> return kafka . structs . Message ( <NUM_LIT:0> , <NUM_LIT:0> , key , payload ) <EOL> def create_gzip_message ( payloads , key = None , compresslevel = None ) : <EOL> """<STR_LIT>""" <EOL> message_set = KafkaProtocol . _encode_message_set ( <EOL> [ create_message ( payload , pl_key ) for payload , pl_key in payloads ] ) <EOL> gzipped = gzip_encode ( message_set , compresslevel = compresslevel ) <EOL> codec = ATTRIBUTE_CODEC_MASK & CODEC_GZIP <EOL> return kafka . structs . Message ( <NUM_LIT:0> , <NUM_LIT> | codec , key , gzipped ) <EOL> def create_snappy_message ( payloads , key = None ) : <EOL> """<STR_LIT>""" <EOL> message_set = KafkaProtocol . _encode_message_set ( <EOL> [ create_message ( payload , pl_key ) for payload , pl_key in payloads ] ) <EOL> snapped = snappy_encode ( message_set ) <EOL> codec = ATTRIBUTE_CODEC_MASK & CODEC_SNAPPY <EOL> return kafka . structs . Message ( <NUM_LIT:0> , <NUM_LIT> | codec , key , snapped ) <EOL> def create_message_set ( messages , codec = CODEC_NONE , key = None , compresslevel = None ) : <EOL> """<STR_LIT>""" <EOL> if codec == CODEC_NONE : <EOL> return [ create_message ( m , k ) for m , k in messages ] <EOL> elif codec == CODEC_GZIP : <EOL> return [ create_gzip_message ( messages , key , compresslevel ) ] <EOL> elif codec == CODEC_SNAPPY : <EOL> return [ create_snappy_message ( messages , key ) ] <EOL> else : <EOL> raise UnsupportedCodecError ( "<STR_LIT>" % codec ) </s>
