ruby-kafka-aws-iam 1.4.1

data/lib/kafka/protocol/record_batch.rb

@@ -0,0 +1,223 @@
+require 'bigdecimal'
+require 'digest/crc32'
+require 'kafka/protocol/record'
+
+module Kafka
+  module Protocol
+    class RecordBatch
+      MAGIC_BYTE = 2
+      # The size of metadata before the real record data
+      RECORD_BATCH_OVERHEAD = 49
+      # Masks to extract information from attributes
+      CODEC_ID_MASK = 0b00000111
+      IN_TRANSACTION_MASK = 0b00010000
+      IS_CONTROL_BATCH_MASK = 0b00100000
+      TIMESTAMP_TYPE_MASK = 0b001000
+
+      attr_reader :records, :first_offset, :first_timestamp, :partition_leader_epoch, :in_transaction, :is_control_batch, :last_offset_delta, :max_timestamp, :producer_id, :producer_epoch, :first_sequence
+
+      attr_accessor :codec_id
+
+      def initialize(
+          records: [],
+          first_offset: 0,
+          first_timestamp: Time.now,
+          partition_leader_epoch: 0,
+          codec_id: 0,
+          in_transaction: false,
+          is_control_batch: false,
+          last_offset_delta: 0,
+          producer_id: -1,
+          producer_epoch: 0,
+          first_sequence: 0,
+          max_timestamp: Time.now
+      )
+        @records = Array(records)
+        @first_offset = first_offset
+        @first_timestamp = first_timestamp
+        @codec_id = codec_id
+
+        # Records verification
+        @last_offset_delta = last_offset_delta
+        @max_timestamp = max_timestamp
+
+        # Transaction information
+        @producer_id = producer_id
+        @producer_epoch = producer_epoch
+
+        @first_sequence = first_sequence
+        @partition_leader_epoch = partition_leader_epoch
+        @in_transaction = in_transaction
+        @is_control_batch = is_control_batch
+
+        mark_control_record
+      end
+
+      def size
+        @records.size
+      end
+
+      def last_offset
+        @first_offset + @last_offset_delta
+      end
+
+      def attributes
+        0x0000 | @codec_id |
+          (@in_transaction ? IN_TRANSACTION_MASK : 0x0) |
+          (@is_control_batch ? IS_CONTROL_BATCH_MASK : 0x0)
+      end
+
+      def encode(encoder)
+        encoder.write_int64(@first_offset)
+
+        record_batch_buffer = StringIO.new
+        record_batch_encoder = Encoder.new(record_batch_buffer)
+
+        record_batch_encoder.write_int32(@partition_leader_epoch)
+        record_batch_encoder.write_int8(MAGIC_BYTE)
+
+        body = encode_record_batch_body
+        crc = ::Digest::CRC32c.checksum(body)
+
+        record_batch_encoder.write_int32(crc)
+        record_batch_encoder.write(body)
+
+        encoder.write_bytes(record_batch_buffer.string)
+      end
+
+      def encode_record_batch_body
+        buffer = StringIO.new
+        encoder = Encoder.new(buffer)
+
+        encoder.write_int16(attributes)
+        encoder.write_int32(@last_offset_delta)
+        encoder.write_int64((@first_timestamp.to_f * 1000).to_i)
+        encoder.write_int64((@max_timestamp.to_f * 1000).to_i)
+
+        encoder.write_int64(@producer_id)
+        encoder.write_int16(@producer_epoch)
+        encoder.write_int32(@first_sequence)
+
+        encoder.write_int32(@records.length)
+
+        records_array = encode_record_array
+        if compressed?
+          codec = Compression.find_codec_by_id(@codec_id)
+          records_array = codec.compress(records_array)
+        end
+        encoder.write(records_array)
+
+        buffer.string
+      end
+
+      def encode_record_array
+        buffer = StringIO.new
+        encoder = Encoder.new(buffer)
+        @records.each do |record|
+          record.encode(encoder)
+        end
+        buffer.string
+      end
+
+      def compressed?
+        @codec_id != 0
+      end
+
+      def fulfill_relative_data
+        first_record = records.min_by { |record| record.create_time }
+        @first_timestamp = first_record.nil? ? Time.now : first_record.create_time
+
+        last_record = records.max_by { |record| record.create_time }
+        @max_timestamp = last_record.nil? ? Time.now : last_record.create_time
+
+        records.each_with_index do |record, index|
+          record.offset_delta = index
+          record.timestamp_delta = ((record.create_time - first_timestamp) * 1000).to_i
+        end
+        @last_offset_delta = records.length - 1
+      end
+
+      def ==(other)
+        records == other.records &&
+          first_offset == other.first_offset &&
+          partition_leader_epoch == other.partition_leader_epoch &&
+          in_transaction == other.in_transaction &&
+          is_control_batch == other.is_control_batch &&
+          last_offset_delta == other.last_offset_delta &&
+          producer_id == other.producer_id &&
+          producer_epoch == other.producer_epoch &&
+          first_sequence == other.first_sequence
+      end
+
+      def self.decode(decoder)
+        first_offset = decoder.int64
+
+        record_batch_raw = decoder.bytes
+        record_batch_decoder = Decoder.from_string(record_batch_raw)
+
+        partition_leader_epoch = record_batch_decoder.int32
+        # Currently, the magic byte is used to distingush legacy MessageSet and
+        # RecordBatch. Therefore, we don't care about magic byte here yet.
+        _magic_byte = record_batch_decoder.int8
+        _crc = record_batch_decoder.int32
+
+        attributes = record_batch_decoder.int16
+        codec_id = attributes & CODEC_ID_MASK
+        in_transaction = (attributes & IN_TRANSACTION_MASK) > 0
+        is_control_batch = (attributes & IS_CONTROL_BATCH_MASK) > 0
+        log_append_time = (attributes & TIMESTAMP_TYPE_MASK) != 0
+
+        last_offset_delta = record_batch_decoder.int32
+        first_timestamp = Time.at(record_batch_decoder.int64 / BigDecimal(1000))
+        max_timestamp = Time.at(record_batch_decoder.int64 / BigDecimal(1000))
+
+        producer_id = record_batch_decoder.int64
+        producer_epoch = record_batch_decoder.int16
+        first_sequence = record_batch_decoder.int32
+
+        records_array_length = record_batch_decoder.int32
+        records_array_raw = record_batch_decoder.read(
+          record_batch_raw.size - RECORD_BATCH_OVERHEAD
+        )
+        if codec_id != 0
+          codec = Compression.find_codec_by_id(codec_id)
+          records_array_raw = codec.decompress(records_array_raw)
+        end
+
+        records_array_decoder = Decoder.from_string(records_array_raw)
+        records_array = []
+        until records_array_decoder.eof?
+          record = Record.decode(records_array_decoder)
+          record.offset = first_offset + record.offset_delta
+          record.create_time = log_append_time && max_timestamp ? max_timestamp : first_timestamp + record.timestamp_delta / BigDecimal(1000)
+          records_array << record
+        end
+
+        raise InsufficientDataMessage if records_array.length != records_array_length
+
+        new(
+          records: records_array,
+          first_offset: first_offset,
+          first_timestamp: first_timestamp,
+          partition_leader_epoch: partition_leader_epoch,
+          in_transaction: in_transaction,
+          is_control_batch: is_control_batch,
+          last_offset_delta: last_offset_delta,
+          producer_id: producer_id,
+          producer_epoch: producer_epoch,
+          first_sequence: first_sequence,
+          max_timestamp: max_timestamp
+        )
+      rescue EOFError
+        raise InsufficientDataMessage, 'Partial trailing record detected!'
+      end
+
+      def mark_control_record
+        if is_control_batch
+          record = @records.first
+          record.is_control_record = true unless record.nil?
+        end
+      end
+    end
+  end
+end

data/lib/kafka/protocol/request_message.rb

@@ -0,0 +1,26 @@
+# frozen_string_literal: true
+
+module Kafka
+  module Protocol
+    class RequestMessage
+      API_VERSION = 0
+
+      def initialize(api_key:, api_version: API_VERSION, correlation_id:, client_id:, request:)
+        @api_key = api_key
+        @api_version = api_version
+        @correlation_id = correlation_id
+        @client_id = client_id
+        @request = request
+      end
+
+      def encode(encoder)
+        encoder.write_int16(@api_key)
+        encoder.write_int16(@api_version)
+        encoder.write_int32(@correlation_id)
+        encoder.write_string(@client_id)
+
+        @request.encode(encoder)
+      end
+    end
+  end
+end

data/lib/kafka/protocol/sasl_handshake_request.rb

@@ -0,0 +1,33 @@
+# frozen_string_literal: true
+
+module Kafka
+  module Protocol
+
+    # SaslHandshake Request (Version: 0) => mechanism
+    #  mechanism => string
+
+    class SaslHandshakeRequest
+
+      SUPPORTED_MECHANISMS = %w(AWS_MSK_IAM GSSAPI PLAIN SCRAM-SHA-256 SCRAM-SHA-512 OAUTHBEARER)
+
+      def initialize(mechanism)
+        unless SUPPORTED_MECHANISMS.include?(mechanism)
+          raise Kafka::Error, "Unsupported SASL mechanism #{mechanism}. Supported are #{SUPPORTED_MECHANISMS.join(', ')}"
+        end
+        @mechanism = mechanism
+      end
+
+      def api_key
+        SASL_HANDSHAKE_API
+      end
+
+      def response_class
+        SaslHandshakeResponse
+      end
+
+      def encode(encoder)
+        encoder.write_string(@mechanism)
+      end
+    end
+  end
+end

data/lib/kafka/protocol/sasl_handshake_response.rb

@@ -0,0 +1,28 @@
+# frozen_string_literal: true
+
+module Kafka
+  module Protocol
+
+    # SaslHandshake Response (Version: 0) => error_code [enabled_mechanisms]
+    #  error_code => int16
+    #  enabled_mechanisms => array of strings
+
+    class SaslHandshakeResponse
+      attr_reader :error_code
+
+      attr_reader :enabled_mechanisms
+
+      def initialize(error_code:, enabled_mechanisms:)
+        @error_code = error_code
+        @enabled_mechanisms = enabled_mechanisms
+      end
+
+      def self.decode(decoder)
+        new(
+          error_code: decoder.int16,
+          enabled_mechanisms: decoder.array { decoder.string }
+        )
+      end
+    end
+  end
+end

data/lib/kafka/protocol/sync_group_request.rb

@@ -0,0 +1,33 @@
+# frozen_string_literal: true
+
+module Kafka
+  module Protocol
+    class SyncGroupRequest
+      def initialize(group_id:, generation_id:, member_id:, group_assignment: {})
+        @group_id = group_id
+        @generation_id = generation_id
+        @member_id = member_id
+        @group_assignment = group_assignment
+      end
+
+      def api_key
+        SYNC_GROUP_API
+      end
+
+      def response_class
+        SyncGroupResponse
+      end
+
+      def encode(encoder)
+        encoder.write_string(@group_id)
+        encoder.write_int32(@generation_id)
+        encoder.write_string(@member_id)
+
+        encoder.write_array(@group_assignment) do |member_id, member_assignment|
+          encoder.write_string(member_id)
+          encoder.write_bytes(Encoder.encode_with(member_assignment))
+        end
+      end
+    end
+  end
+end

data/lib/kafka/protocol/sync_group_response.rb

@@ -0,0 +1,26 @@
+# frozen_string_literal: true
+
+require "kafka/protocol/member_assignment"
+
+module Kafka
+  module Protocol
+    class SyncGroupResponse
+      attr_reader :error_code, :member_assignment
+
+      def initialize(error_code:, member_assignment:)
+        @error_code = error_code
+        @member_assignment = member_assignment
+      end
+
+      def self.decode(decoder)
+        error_code = decoder.int16
+        member_assignment_bytes = decoder.bytes
+
+        new(
+          error_code: error_code,
+          member_assignment: member_assignment_bytes ? MemberAssignment.decode(Decoder.from_string(member_assignment_bytes)) : nil
+        )
+      end
+    end
+  end
+end

data/lib/kafka/protocol/txn_offset_commit_request.rb

@@ -0,0 +1,46 @@
+# frozen_string_literal: true
+
+module Kafka
+  module Protocol
+    class TxnOffsetCommitRequest
+
+      def api_key
+        TXN_OFFSET_COMMIT_API
+      end
+
+      def api_version
+        2
+      end
+
+      def response_class
+        TxnOffsetCommitResponse
+      end
+
+      def initialize(transactional_id:, group_id:, producer_id:, producer_epoch:, offsets:)
+        @transactional_id = transactional_id
+        @producer_id = producer_id
+        @producer_epoch = producer_epoch
+        @group_id = group_id
+        @offsets = offsets
+      end
+
+      def encode(encoder)
+        encoder.write_string(@transactional_id.to_s)
+        encoder.write_string(@group_id)
+        encoder.write_int64(@producer_id)
+        encoder.write_int16(@producer_epoch)
+
+        encoder.write_array(@offsets) do |topic, partitions|
+          encoder.write_string(topic)
+          encoder.write_array(partitions) do |partition, offset|
+            encoder.write_int32(partition)
+            encoder.write_int64(offset[:offset])
+            encoder.write_string(nil) # metadata
+            encoder.write_int32(offset[:leader_epoch])
+          end
+        end
+      end
+
+    end
+  end
+end

data/lib/kafka/protocol/txn_offset_commit_response.rb

@@ -0,0 +1,47 @@
+# frozen_string_literal: true
+
+module Kafka
+  module Protocol
+    class TxnOffsetCommitResponse
+      class PartitionError
+        attr_reader :partition, :error_code
+
+        def initialize(partition:, error_code:)
+          @partition = partition
+          @error_code = error_code
+        end
+      end
+
+      class TopicPartitionsError
+        attr_reader :topic, :partitions
+
+        def initialize(topic:, partitions:)
+          @topic = topic
+          @partitions = partitions
+        end
+      end
+
+      attr_reader :errors
+
+      def initialize(errors:)
+        @errors = errors
+      end
+
+      def self.decode(decoder)
+        _throttle_time_ms = decoder.int32
+        errors = decoder.array do
+          TopicPartitionsError.new(
+            topic: decoder.string,
+            partitions: decoder.array do
+              PartitionError.new(
+                partition: decoder.int32,
+                error_code: decoder.int16
+              )
+            end
+          )
+        end
+        new(errors: errors)
+      end
+    end
+  end
+end

data/lib/kafka/protocol.rb

@@ -0,0 +1,225 @@
+# frozen_string_literal: true
+
+module Kafka
+
+  # The protocol layer of the library.
+  #
+  # The Kafka protocol (https://kafka.apache.org/protocol) defines a set of API
+  # requests, each with a well-known numeric API key, as well as a set of error
+  # codes with specific meanings.
+  #
+  # This module, and the classes contained in it, implement the client side of
+  # the protocol.
+  module Protocol
+    # The replica id of non-brokers is always -1.
+    REPLICA_ID = -1
+
+    PRODUCE_API               = 0
+    FETCH_API                 = 1
+    LIST_OFFSET_API           = 2
+    TOPIC_METADATA_API        = 3
+    OFFSET_COMMIT_API         = 8
+    OFFSET_FETCH_API          = 9
+    FIND_COORDINATOR_API      = 10
+    JOIN_GROUP_API            = 11
+    HEARTBEAT_API             = 12
+    LEAVE_GROUP_API           = 13
+    SYNC_GROUP_API            = 14
+    DESCRIBE_GROUPS_API       = 15
+    LIST_GROUPS_API           = 16
+    SASL_HANDSHAKE_API        = 17
+    API_VERSIONS_API          = 18
+    CREATE_TOPICS_API         = 19
+    DELETE_TOPICS_API         = 20
+    INIT_PRODUCER_ID_API      = 22
+    ADD_PARTITIONS_TO_TXN_API = 24
+    ADD_OFFSETS_TO_TXN_API    = 25
+    END_TXN_API               = 26
+    TXN_OFFSET_COMMIT_API     = 28
+    DESCRIBE_CONFIGS_API      = 32
+    ALTER_CONFIGS_API         = 33
+    CREATE_PARTITIONS_API     = 37
+
+    # A mapping from numeric API keys to symbolic API names.
+    APIS = {
+      PRODUCE_API               => :produce,
+      FETCH_API                 => :fetch,
+      LIST_OFFSET_API           => :list_offset,
+      TOPIC_METADATA_API        => :topic_metadata,
+      OFFSET_COMMIT_API         => :offset_commit,
+      OFFSET_FETCH_API          => :offset_fetch,
+      FIND_COORDINATOR_API      => :find_coordinator,
+      JOIN_GROUP_API            => :join_group,
+      HEARTBEAT_API             => :heartbeat,
+      LEAVE_GROUP_API           => :leave_group,
+      SYNC_GROUP_API            => :sync_group,
+      SASL_HANDSHAKE_API        => :sasl_handshake,
+      API_VERSIONS_API          => :api_versions,
+      CREATE_TOPICS_API         => :create_topics,
+      DELETE_TOPICS_API         => :delete_topics,
+      INIT_PRODUCER_ID_API      => :init_producer_id_api,
+      ADD_PARTITIONS_TO_TXN_API => :add_partitions_to_txn_api,
+      ADD_OFFSETS_TO_TXN_API    => :add_offsets_to_txn_api,
+      END_TXN_API               => :end_txn_api,
+      TXN_OFFSET_COMMIT_API     => :txn_offset_commit_api,
+      DESCRIBE_CONFIGS_API      => :describe_configs_api,
+      CREATE_PARTITIONS_API     => :create_partitions
+    }
+
+    # A mapping from numeric error codes to exception classes.
+    ERRORS = {
+      -1 => UnknownError,
+       1 => OffsetOutOfRange,
+       2 => CorruptMessage,
+       3 => UnknownTopicOrPartition,
+       4 => InvalidMessageSize,
+       5 => LeaderNotAvailable,
+       6 => NotLeaderForPartition,
+       7 => RequestTimedOut,
+       8 => BrokerNotAvailable,
+       9 => ReplicaNotAvailable,
+      10 => MessageSizeTooLarge,
+      11 => StaleControllerEpoch,
+      12 => OffsetMetadataTooLarge,
+      13 => NetworkException,
+      14 => CoordinatorLoadInProgress,
+      15 => CoordinatorNotAvailable,
+      16 => NotCoordinatorForGroup,
+      17 => InvalidTopic,
+      18 => RecordListTooLarge,
+      19 => NotEnoughReplicas,
+      20 => NotEnoughReplicasAfterAppend,
+      21 => InvalidRequiredAcks,
+      22 => IllegalGeneration,
+      23 => InconsistentGroupProtocol,
+      24 => InvalidGroupId,
+      25 => UnknownMemberId,
+      26 => InvalidSessionTimeout,
+      27 => RebalanceInProgress,
+      28 => InvalidCommitOffsetSize,
+      29 => TopicAuthorizationFailed,
+      30 => GroupAuthorizationFailed,
+      31 => ClusterAuthorizationFailed,
+      32 => InvalidTimestamp,
+      33 => UnsupportedSaslMechanism,
+      34 => InvalidSaslState,
+      35 => UnsupportedVersion,
+      36 => TopicAlreadyExists,
+      37 => InvalidPartitions,
+      38 => InvalidReplicationFactor,
+      39 => InvalidReplicaAssignment,
+      40 => InvalidConfig,
+      41 => NotController,
+      42 => InvalidRequest,
+      43 => UnsupportedForMessageFormat,
+      44 => PolicyViolation,
+      45 => OutOfOrderSequenceNumberError,
+      46 => DuplicateSequenceNumberError,
+      47 => InvalidProducerEpochError,
+      48 => InvalidTxnStateError,
+      49 => InvalidProducerIDMappingError,
+      50 => InvalidTransactionTimeoutError,
+      51 => ConcurrentTransactionError,
+      52 => TransactionCoordinatorFencedError
+    }
+
+    # A mapping from int to corresponding resource type in symbol.
+    # https://github.com/apache/kafka/blob/trunk/clients/src/main/java/org/apache/kafka/common/resource/ResourceType.java
+    RESOURCE_TYPE_UNKNOWN          = 0
+    RESOURCE_TYPE_ANY              = 1
+    RESOURCE_TYPE_TOPIC            = 2
+    RESOURCE_TYPE_GROUP            = 3
+    RESOURCE_TYPE_CLUSTER          = 4
+    RESOURCE_TYPE_TRANSACTIONAL_ID = 5
+    RESOURCE_TYPE_DELEGATION_TOKEN = 6
+    RESOURCE_TYPES = {
+      RESOURCE_TYPE_UNKNOWN          => :unknown,
+      RESOURCE_TYPE_ANY              => :any,
+      RESOURCE_TYPE_TOPIC            => :topic,
+      RESOURCE_TYPE_GROUP            => :group,
+      RESOURCE_TYPE_CLUSTER          => :cluster,
+      RESOURCE_TYPE_TRANSACTIONAL_ID => :transactional_id,
+      RESOURCE_TYPE_DELEGATION_TOKEN => :delegation_token,
+    }
+
+    # Coordinator types. Since Kafka 0.11.0, there are types of coordinators:
+    # Group and Transaction
+    COORDINATOR_TYPE_GROUP = 0
+    COORDINATOR_TYPE_TRANSACTION = 1
+
+    # Handles an error code by either doing nothing (if there was no error) or
+    # by raising an appropriate exception.
+    #
+    # @param error_code Integer
+    # @raise [ProtocolError]
+    # @return [nil]
+    def self.handle_error(error_code, error_message = nil)
+      if error_code == 0
+        # No errors, yay!
+      elsif error = ERRORS[error_code]
+        raise error, error_message
+      else
+        raise UnknownError, "Unknown error with code #{error_code} #{error_message}"
+      end
+    end
+
+    # Returns the symbolic name for an API key.
+    #
+    # @param api_key Integer
+    # @return [Symbol]
+    def self.api_name(api_key)
+      APIS.fetch(api_key, :unknown)
+    end
+  end
+end
+
+require "kafka/protocol/metadata_request"
+require "kafka/protocol/metadata_response"
+require "kafka/protocol/produce_request"
+require "kafka/protocol/produce_response"
+require "kafka/protocol/fetch_request"
+require "kafka/protocol/fetch_response"
+require "kafka/protocol/list_offset_request"
+require "kafka/protocol/list_offset_response"
+require "kafka/protocol/add_offsets_to_txn_request"
+require "kafka/protocol/add_offsets_to_txn_response"
+require "kafka/protocol/txn_offset_commit_request"
+require "kafka/protocol/txn_offset_commit_response"
+require "kafka/protocol/find_coordinator_request"
+require "kafka/protocol/find_coordinator_response"
+require "kafka/protocol/join_group_request"
+require "kafka/protocol/join_group_response"
+require "kafka/protocol/sync_group_request"
+require "kafka/protocol/sync_group_response"
+require "kafka/protocol/leave_group_request"
+require "kafka/protocol/leave_group_response"
+require "kafka/protocol/heartbeat_request"
+require "kafka/protocol/heartbeat_response"
+require "kafka/protocol/offset_fetch_request"
+require "kafka/protocol/offset_fetch_response"
+require "kafka/protocol/offset_commit_request"
+require "kafka/protocol/offset_commit_response"
+require "kafka/protocol/api_versions_request"
+require "kafka/protocol/api_versions_response"
+require "kafka/protocol/sasl_handshake_request"
+require "kafka/protocol/sasl_handshake_response"
+require "kafka/protocol/create_topics_request"
+require "kafka/protocol/create_topics_response"
+require "kafka/protocol/delete_topics_request"
+require "kafka/protocol/delete_topics_response"
+require "kafka/protocol/describe_configs_request"
+require "kafka/protocol/describe_configs_response"
+require "kafka/protocol/alter_configs_request"
+require "kafka/protocol/alter_configs_response"
+require "kafka/protocol/create_partitions_request"
+require "kafka/protocol/create_partitions_response"
+require "kafka/protocol/list_groups_request"
+require "kafka/protocol/list_groups_response"
+require "kafka/protocol/describe_groups_request"
+require "kafka/protocol/describe_groups_response"
+require "kafka/protocol/init_producer_id_request"
+require "kafka/protocol/init_producer_id_response"
+require "kafka/protocol/add_partitions_to_txn_request"
+require "kafka/protocol/add_partitions_to_txn_response"
+require "kafka/protocol/end_txn_request"
+require "kafka/protocol/end_txn_response"