ruby-kafka-aws-iam 1.4.1

data/lib/kafka/protocol/leave_group_request.rb

@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+
+module Kafka
+  module Protocol
+    class LeaveGroupRequest
+      def initialize(group_id:, member_id:)
+        @group_id = group_id
+        @member_id = member_id
+      end
+
+      def api_key
+        LEAVE_GROUP_API
+      end
+
+      def response_class
+        LeaveGroupResponse
+      end
+
+      def encode(encoder)
+        encoder.write_string(@group_id)
+        encoder.write_string(@member_id)
+      end
+    end
+  end
+end

data/lib/kafka/protocol/leave_group_response.rb

@@ -0,0 +1,17 @@
+# frozen_string_literal: true
+
+module Kafka
+  module Protocol
+    class LeaveGroupResponse
+      attr_reader :error_code
+
+      def initialize(error_code:)
+        @error_code = error_code
+      end
+
+      def self.decode(decoder)
+        new(error_code: decoder.int16)
+      end
+    end
+  end
+end

data/lib/kafka/protocol/list_groups_request.rb

@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+module Kafka
+  module Protocol
+    class ListGroupsRequest
+      def api_key
+        LIST_GROUPS_API
+      end
+
+      def api_version
+        0
+      end
+
+      def response_class
+        Protocol::ListGroupsResponse
+      end
+
+      def encode(encoder)
+        # noop
+      end
+    end
+  end
+end

data/lib/kafka/protocol/list_groups_response.rb

@@ -0,0 +1,35 @@
+# frozen_string_literal: true
+
+module Kafka
+  module Protocol
+    class ListGroupsResponse
+      class GroupEntry
+        attr_reader :group_id, :protocol_type
+
+        def initialize(group_id:, protocol_type:)
+          @group_id = group_id
+          @protocol_type = protocol_type
+        end
+      end
+
+      attr_reader :error_code, :groups
+
+      def initialize(error_code:, groups:)
+        @error_code = error_code
+        @groups = groups
+      end
+
+      def self.decode(decoder)
+        error_code = decoder.int16
+        groups = decoder.array do
+          GroupEntry.new(
+            group_id: decoder.string,
+            protocol_type: decoder.string
+          )
+        end
+
+        new(error_code: error_code, groups: groups)
+      end
+    end
+  end
+end

data/lib/kafka/protocol/list_offset_request.rb

@@ -0,0 +1,53 @@
+# frozen_string_literal: true
+
+module Kafka
+  module Protocol
+    # A request to list the available offsets for a set of topics/partitions.
+    #
+    # ## API Specification
+    #
+    #     OffsetRequest => ReplicaId [TopicName [Partition Time MaxNumberOfOffsets]]
+    #       ReplicaId => int32
+    #       IsolationLevel => int8
+    #       TopicName => string
+    #       Partition => int32
+    #       Time => int64
+    #
+    class ListOffsetRequest
+      ISOLATION_READ_UNCOMMITTED = 0
+      ISOLATION_READ_COMMITTED = 1
+
+      # @param topics [Hash]
+      def initialize(topics:)
+        @replica_id = REPLICA_ID
+        @topics = topics
+      end
+
+      def api_version
+        2
+      end
+
+      def api_key
+        LIST_OFFSET_API
+      end
+
+      def response_class
+        Protocol::ListOffsetResponse
+      end
+
+      def encode(encoder)
+        encoder.write_int32(@replica_id)
+        encoder.write_int8(ISOLATION_READ_COMMITTED)
+
+        encoder.write_array(@topics) do |topic, partitions|
+          encoder.write_string(topic)
+
+          encoder.write_array(partitions) do |partition|
+            encoder.write_int32(partition.fetch(:partition))
+            encoder.write_int64(partition.fetch(:time))
+          end
+        end
+      end
+    end
+  end
+end

data/lib/kafka/protocol/list_offset_response.rb

@@ -0,0 +1,89 @@
+# frozen_string_literal: true
+
+module Kafka
+  module Protocol
+
+    # A response to a list offset request.
+    #
+    # ## API Specification
+    #
+    #     OffsetResponse => [TopicName [PartitionOffsets]]
+    #       ThrottleTimeMS => int32
+    #       PartitionOffsets => Partition ErrorCode Timestamp Offset
+    #       Partition => int32
+    #       ErrorCode => int16
+    #       Timestamp => int64
+    #       Offset => int64
+    #
+    class ListOffsetResponse
+      class TopicOffsetInfo
+        attr_reader :name, :partition_offsets
+
+        def initialize(name:, partition_offsets:)
+          @name = name
+          @partition_offsets = partition_offsets
+        end
+      end
+
+      class PartitionOffsetInfo
+        attr_reader :partition, :error_code, :timestamp, :offset
+
+        def initialize(partition:, error_code:, timestamp:, offset:)
+          @partition = partition
+          @error_code = error_code
+          @timestamp = timestamp
+          @offset = offset
+        end
+      end
+
+      attr_reader :topics
+
+      def initialize(topics:)
+        @topics = topics
+      end
+
+      def offset_for(topic, partition)
+        topic_info = @topics.find {|t| t.name == topic }
+
+        if topic_info.nil?
+          raise UnknownTopicOrPartition, "Unknown topic #{topic}"
+        end
+
+        partition_info = topic_info
+          .partition_offsets
+          .find {|p| p.partition == partition }
+
+        if partition_info.nil?
+          raise UnknownTopicOrPartition, "Unknown partition #{topic}/#{partition}"
+        end
+
+        Protocol.handle_error(partition_info.error_code)
+
+        partition_info.offset
+      end
+
+      def self.decode(decoder)
+        _throttle_time_ms = decoder.int32
+        topics = decoder.array do
+          name = decoder.string
+
+          partition_offsets = decoder.array do
+            PartitionOffsetInfo.new(
+              partition: decoder.int32,
+              error_code: decoder.int16,
+              timestamp: decoder.int64,
+              offset: decoder.int64
+            )
+          end
+
+          TopicOffsetInfo.new(
+            name: name,
+            partition_offsets: partition_offsets
+          )
+        end
+
+        new(topics: topics)
+      end
+    end
+  end
+end

data/lib/kafka/protocol/member_assignment.rb

@@ -0,0 +1,42 @@
+# frozen_string_literal: true
+
+module Kafka
+  module Protocol
+    class MemberAssignment
+      attr_reader :topics
+
+      def initialize(version: 0, topics: {}, user_data: nil)
+        @version = version
+        @topics = topics
+        @user_data = user_data
+      end
+
+      def assign(topic, partitions)
+        @topics[topic] ||= []
+        @topics[topic].concat(partitions)
+      end
+
+      def encode(encoder)
+        encoder.write_int16(@version)
+
+        encoder.write_array(@topics) do |topic, partitions|
+          encoder.write_string(topic)
+
+          encoder.write_array(partitions) do |partition|
+            encoder.write_int32(partition)
+          end
+        end
+
+        encoder.write_bytes(@user_data)
+      end
+
+      def self.decode(decoder)
+        new(
+          version: decoder.int16,
+          topics: Hash[decoder.array { [decoder.string, decoder.array { decoder.int32 }] }],
+          user_data: decoder.bytes,
+        )
+      end
+    end
+  end
+end

data/lib/kafka/protocol/message.rb

@@ -0,0 +1,172 @@
+# frozen_string_literal: true
+
+require "stringio"
+require "zlib"
+
+module Kafka
+  module Protocol
+
+    # ## API Specification
+    #
+    #     Message => Crc MagicByte Attributes Timestamp Key Value
+    #         Crc => int32
+    #         MagicByte => int8
+    #         Attributes => int8
+    #         Timestamp => int64, in ms
+    #         Key => bytes
+    #         Value => bytes
+    #
+    class Message
+      MAGIC_BYTE = 1
+
+      attr_reader :key, :value, :codec_id, :offset
+
+      attr_reader :bytesize, :create_time
+
+      def initialize(value:, key: nil, create_time: Time.now, codec_id: 0, offset: -1)
+        @key = key
+        @value = value
+        @codec_id = codec_id
+        @offset = offset
+        @create_time = create_time
+
+        @bytesize = @key.to_s.bytesize + @value.to_s.bytesize
+      end
+
+      def encode(encoder)
+        data = encode_with_crc
+
+        encoder.write_int64(offset)
+        encoder.write_bytes(data)
+      end
+
+      def ==(other)
+        @key == other.key &&
+          @value == other.value &&
+          @codec_id == other.codec_id &&
+          @offset == other.offset
+      end
+
+      def compressed?
+        @codec_id != 0
+      end
+
+      # @return [Array<Kafka::Protocol::Message>]
+      def decompress
+        codec = Compression.find_codec_by_id(@codec_id)
+
+        # For some weird reason we need to cut out the first 20 bytes.
+        data = codec.decompress(value)
+        message_set_decoder = Decoder.from_string(data)
+        message_set = MessageSet.decode(message_set_decoder)
+
+        correct_offsets(message_set.messages)
+      end
+
+      def self.decode(decoder)
+        offset = decoder.int64
+        message_decoder = Decoder.from_string(decoder.bytes)
+
+        _crc = message_decoder.int32
+        magic_byte = message_decoder.int8
+        attributes = message_decoder.int8
+
+        # The magic byte indicates the message format version. There are situations
+        # where an old message format can be returned from a newer version of Kafka,
+        # because old messages are not necessarily rewritten on upgrades.
+        case magic_byte
+        when 0
+          # No timestamp in the pre-0.10 message format.
+          timestamp = nil
+        when 1
+          timestamp = message_decoder.int64
+
+          # If the timestamp is set to zero, it's because the message has been upgraded
+          # from the Kafka 0.9 disk format to the Kafka 0.10 format. The former didn't
+          # have a timestamp attribute, so we'll just set the timestamp to nil.
+          timestamp = nil if timestamp.zero?
+        else
+          raise Kafka::Error, "Invalid magic byte: #{magic_byte}"
+        end
+
+        key = message_decoder.bytes
+        value = message_decoder.bytes
+
+        # The codec id is encoded in the three least significant bits of the
+        # attributes.
+        codec_id = attributes & 0b111
+
+        # The timestamp will be nil if the message was written in the Kafka 0.9 log format.
+        create_time = timestamp && Time.at(timestamp / 1000.0)
+
+        new(key: key, value: value, codec_id: codec_id, offset: offset, create_time: create_time)
+      end
+
+      # Ensure the backward compatibility of Message format from Kafka 0.11.x
+      def is_control_record
+        false
+      end
+
+      def headers
+        {}
+      end
+
+      private
+
+      # Offsets may be relative with regards to wrapped message offset, but there are special cases.
+      #
+      # Cases when client will receive corrected offsets:
+      #   - When fetch request is version 0, kafka will correct relative offset on broker side before replying fetch response
+      #   - When messages is stored in 0.9 format on disk (broker configured to do so).
+      #
+      # All other cases, compressed inner messages should have relative offset, with below attributes:
+      #   - The container message should have the 'real' offset
+      #   - The container message's offset should be the 'real' offset of the last message in the compressed batch
+      def correct_offsets(messages)
+        max_relative_offset = messages.last.offset
+
+        # The offsets are already correct, do nothing.
+        return messages if max_relative_offset == offset
+
+        # The contained messages have relative offsets, and needs to be corrected.
+        base_offset = offset - max_relative_offset
+
+        messages.map do |message|
+          Message.new(
+            offset: message.offset + base_offset,
+            value: message.value,
+            key: message.key,
+            create_time: message.create_time,
+            codec_id: message.codec_id
+          )
+        end
+      end
+
+      def encode_with_crc
+        buffer = StringIO.new
+        encoder = Encoder.new(buffer)
+
+        data = encode_without_crc
+        crc = Zlib.crc32(data)
+
+        encoder.write_int32(crc)
+        encoder.write(data)
+
+        buffer.string
+      end
+
+      def encode_without_crc
+        buffer = StringIO.new
+        encoder = Encoder.new(buffer)
+
+        encoder.write_int8(MAGIC_BYTE)
+        encoder.write_int8(@codec_id)
+        encoder.write_int64((@create_time.to_f * 1000).to_i)
+        encoder.write_bytes(@key)
+        encoder.write_bytes(@value)
+
+        buffer.string
+      end
+    end
+  end
+end

data/lib/kafka/protocol/message_set.rb

@@ -0,0 +1,55 @@
+# frozen_string_literal: true
+
+module Kafka
+  module Protocol
+    class MessageSet
+      attr_reader :messages
+
+      def initialize(messages: [])
+        @messages = messages
+      end
+
+      def size
+        @messages.size
+      end
+
+      def ==(other)
+        messages == other.messages
+      end
+
+      def encode(encoder)
+        # Messages in a message set are *not* encoded as an array. Rather,
+        # they are written in sequence.
+        @messages.each do |message|
+          message.encode(encoder)
+        end
+      end
+
+      def self.decode(decoder)
+        fetched_messages = []
+
+        until decoder.eof?
+          begin
+            message = Message.decode(decoder)
+
+            if message.compressed?
+              fetched_messages.concat(message.decompress)
+            else
+              fetched_messages << message
+            end
+          rescue EOFError
+            if fetched_messages.empty?
+              # If the first message in the set is truncated, it's likely because the
+              # message is larger than the maximum size that we have asked for.
+              raise MessageTooLargeToRead
+            else
+              # We tried to decode a partial message at the end of the set; just skip it.
+            end
+          end
+        end
+
+        new(messages: fetched_messages)
+      end
+    end
+  end
+end

data/lib/kafka/protocol/metadata_request.rb

@@ -0,0 +1,31 @@
+# frozen_string_literal: true
+
+module Kafka
+  module Protocol
+    class MetadataRequest
+
+      # A request for cluster metadata.
+      #
+      # @param topics [Array<String>]
+      def initialize(topics: [])
+        @topics = topics
+      end
+
+      def api_key
+        TOPIC_METADATA_API
+      end
+
+      def api_version
+        1
+      end
+
+      def response_class
+        Protocol::MetadataResponse
+      end
+
+      def encode(encoder)
+        encoder.write_array(@topics) {|topic| encoder.write_string(topic) }
+      end
+    end
+  end
+end