deimos-ruby 1.6.3 → 1.8.1.pre.beta1

data/lib/deimos/consumer.rb

@@ -1,95 +1,104 @@
 # frozen_string_literal: true
 
-require 'deimos/base_consumer'
-require 'deimos/shared_config'
-require 'phobos/handler'
-require 'active_support/all'
+require 'deimos/consume/batch_consumption'
+require 'deimos/consume/message_consumption'
 
-# Class to consume messages coming from the pipeline topic
+# Class to consume messages coming from a Kafka topic
 # Note: According to the docs, instances of your handler will be created
-# for every incoming message. This class should be lightweight.
+# for every incoming message/batch. This class should be lightweight.
 module Deimos
-  # Parent consumer class.
-  class Consumer < BaseConsumer
-    include Phobos::Handler
+  # Basic consumer class. Inherit from this class and override either consume
+  # or consume_batch, depending on the delivery mode of your listener.
+  # `consume` -> use `delivery :message` or `delivery :batch`
+  # `consume_batch` -> use `delivery :inline_batch`
+  class Consumer
+    include Consume::MessageConsumption
+    include Consume::BatchConsumption
+    include SharedConfig
 
-    # :nodoc:
-    def around_consume(payload, metadata)
-      decoded_payload = payload.dup
-      new_metadata = metadata.dup
-      benchmark = Benchmark.measure do
-        _with_error_span(payload, metadata) do
-          new_metadata[:key] = decode_key(metadata[:key]) if self.class.config[:key_configured]
-          decoded_payload = payload ? self.class.decoder.decode(payload) : nil
-          _received_message(decoded_payload, new_metadata)
-          yield decoded_payload, new_metadata
-        end
+    class << self
+      # @return [Deimos::SchemaBackends::Base]
+      def decoder
+        @decoder ||= Deimos.schema_backend(schema: config[:schema],
+                                           namespace: config[:namespace])
+      end
+
+      # @return [Deimos::SchemaBackends::Base]
+      def key_decoder
+        @key_decoder ||= Deimos.schema_backend(schema: config[:key_schema],
+                                               namespace: config[:namespace])
       end
-      _handle_success(benchmark.real, decoded_payload, new_metadata)
     end
 
-    # Consume incoming messages.
-    # @param _payload [String]
-    # @param _metadata [Hash]
-    def consume(_payload, _metadata)
-      raise NotImplementedError
+    # Helper method to decode an encoded key.
+    # @param key [String]
+    # @return [Object] the decoded key.
+    def decode_key(key)
+      return nil if key.nil?
+
+      config = self.class.config
+      unless config[:key_configured]
+        raise 'No key config given - if you are not decoding keys, please use '\
+          '`key_config plain: true`'
+      end
+
+      if config[:key_field]
+        self.class.decoder.decode_key(key, config[:key_field])
+      elsif config[:key_schema]
+        self.class.key_decoder.decode(key, schema: config[:key_schema])
+      else # no encoding
+        key
+      end
     end
 
   private
 
-    def _received_message(payload, metadata)
-      Deimos.config.logger.info(
-        message: 'Got Kafka event',
-        payload: payload,
-        metadata: metadata
+    def _with_span
+      @span = Deimos.config.tracer&.start(
+        'deimos-consumer',
+        resource: self.class.name.gsub('::', '-')
       )
-      Deimos.config.metrics&.increment('handler', tags: %W(
-                                         status:received
+      yield
+    ensure
+      Deimos.config.tracer&.finish(@span)
+    end
+
+    def _report_time_delayed(payload, metadata)
+      return if payload.nil? || payload['timestamp'].blank?
+
+      begin
+        time_delayed = Time.now.in_time_zone - payload['timestamp'].to_datetime
+      rescue ArgumentError
+        Deimos.config.logger.info(
+          message: "Error parsing timestamp! #{payload['timestamp']}"
+        )
+        return
+      end
+      Deimos.config.metrics&.histogram('handler', time_delayed, tags: %W(
+                                         time:time_delayed
                                          topic:#{metadata[:topic]}
                                        ))
-      _report_time_delayed(payload, metadata)
     end
 
-    # @param exception [Throwable]
-    # @param payload [Hash]
-    # @param metadata [Hash]
-    def _handle_error(exception, payload, metadata)
-      Deimos.config.metrics&.increment(
-        'handler',
-        tags: %W(
-          status:error
-          topic:#{metadata[:topic]}
-        )
-      )
-      Deimos.config.logger.warn(
-        message: 'Error consuming message',
-        handler: self.class.name,
-        metadata: metadata,
-        data: payload,
-        error_message: exception.message,
-        error: exception.backtrace
-      )
-      super
+    # Overrideable method to determine if a given error should be considered
+    # "fatal" and always be reraised.
+    # @param _error [Exception]
+    # @param _payload [Hash]
+    # @param _metadata [Hash]
+    # @return [Boolean]
+    def fatal_error?(_error, _payload, _metadata)
+      false
     end
 
-    # @param time_taken [Float]
+    # @param exception [Exception]
     # @param payload [Hash]
     # @param metadata [Hash]
-    def _handle_success(time_taken, payload, metadata)
-      Deimos.config.metrics&.histogram('handler', time_taken, tags: %W(
-                                         time:consume
-                                         topic:#{metadata[:topic]}
-                                       ))
-      Deimos.config.metrics&.increment('handler', tags: %W(
-                                         status:success
-                                         topic:#{metadata[:topic]}
-                                       ))
-      Deimos.config.logger.info(
-        message: 'Finished processing Kafka event',
-        payload: payload,
-        time_elapsed: time_taken,
-        metadata: metadata
-      )
+    def _error(exception, payload, metadata)
+      Deimos.config.tracer&.set_error(@span, exception)
+
+      raise if Deimos.config.consumers.reraise_errors ||
+               Deimos.config.consumers.fatal_error&.call(exception, payload, metadata) ||
+               fatal_error?(exception, payload, metadata)
     end
   end
 end

data/lib/deimos/kafka_message.rb

@@ -42,7 +42,7 @@ module Deimos
       messages.map do |m|
         {
           key: m.key.present? ? decoder&.decode_key(m.key) || m.key : nil,
-          payload: decoder&.decoder&.decode(self.message) || self.message
+          payload: decoder&.decoder&.decode(m.message) || m.message
         }
       end
     end

data/lib/deimos/kafka_topic_info.rb

@@ -13,8 +13,8 @@ module Deimos
       def lock(topic, lock_id)
         # Try to create it - it's fine if it already exists
         begin
-          self.create(topic: topic)
-        rescue ActiveRecord::RecordNotUnique # rubocop:disable Lint/SuppressedException
+          self.create(topic: topic, last_processed_at: Time.zone.now)
+        rescue ActiveRecord::RecordNotUnique
           # continue on
         end
 
@@ -52,7 +52,26 @@ module Deimos
       # @param lock_id [String]
       def clear_lock(topic, lock_id)
         self.where(topic: topic, locked_by: lock_id).
-          update_all(locked_by: nil, locked_at: nil, error: false, retries: 0)
+          update_all(locked_by: nil,
+                     locked_at: nil,
+                     error: false,
+                     retries: 0,
+                     last_processed_at: Time.zone.now)
+      end
+
+      # Update all topics that aren't currently locked and have no messages
+      # waiting. It's OK if some messages get inserted in the middle of this
+      # because the point is that at least within a few milliseconds of each
+      # other, it wasn't locked and had no messages, meaning the topic
+      # was in a good state.
+      # @param except_topics [Array<String>] the list of topics we've just
+      # realized had messages in them, meaning all other topics were empty.
+      def ping_empty_topics(except_topics)
+        records = KafkaTopicInfo.where(locked_by: nil).
+          where('topic not in(?)', except_topics)
+        records.each do |info|
+          info.update_attribute(:last_processed_at, Time.zone.now)
+        end
       end
 
       # The producer calls this if it gets an error sending messages. This

data/lib/deimos/message.rb

@@ -10,7 +10,7 @@ module Deimos
     # @param producer [Class]
     def initialize(payload, producer, topic: nil, key: nil, partition_key: nil)
       @payload = payload&.with_indifferent_access
-      @producer_name = producer.name
+      @producer_name = producer&.name
       @topic = topic
       @key = key
       @partition_key = partition_key
@@ -70,5 +70,10 @@ module Deimos
     def ==(other)
       self.to_h == other.to_h
     end
+
+    # @return [Boolean] True if this message is a tombstone
+    def tombstone?
+      payload.nil?
+    end
   end
 end

data/lib/deimos/metrics/provider.rb

@@ -1,6 +1,5 @@
 # frozen_string_literal: true
 
-# rubocop:disable Lint/UnusedMethodArgument
 module Deimos
   module Metrics
     # Base class for all metrics providers.
@@ -35,4 +34,3 @@ module Deimos
     end
   end
 end
-# rubocop:enable Lint/UnusedMethodArgument

data/lib/deimos/poll_info.rb

@@ -0,0 +1,9 @@
+# frozen_string_literal: true
+
+module Deimos
+  # ActiveRecord class to record the last time we polled the database.
+  # For use with DbPoller.
+  class PollInfo < ActiveRecord::Base
+    self.table_name = 'deimos_poll_info'
+  end
+end

data/lib/deimos/schema_backends/avro_base.rb

@@ -33,6 +33,30 @@ module Deimos
         decode(payload, schema: @key_schema['name'])[field_name]
       end
 
+      # :nodoc:
+      def sql_type(field)
+        type = field.type.type
+        return type if %w(array map record).include?(type)
+
+        if type == :union
+          non_null = field.type.schemas.reject { |f| f.type == :null }
+          if non_null.size > 1
+            warn("WARNING: #{field.name} has more than one non-null type. Picking the first for the SQL type.")
+          end
+          return non_null.first.type
+        end
+        return type.to_sym if %w(float boolean).include?(type)
+        return :integer if type == 'int'
+        return :bigint if type == 'long'
+
+        if type == 'double'
+          warn('Avro `double` type turns into SQL `float` type. Please ensure you have the correct `limit` set.')
+          return :float
+        end
+
+        :string
+      end
+
       # @override
       def coerce_field(field, value)
         AvroSchemaCoercer.new(avro_schema).coerce_type(field.type, value)
@@ -40,7 +64,10 @@ module Deimos
 
       # @override
       def schema_fields
-        avro_schema.fields.map { |field| SchemaField.new(field.name, field.type) }
+        avro_schema.fields.map do |field|
+          enum_values = field.type.type == 'enum' ? field.type.symbols : []
+          SchemaField.new(field.name, field.type, enum_values)
+        end
       end
 
       # @override

data/lib/deimos/schema_backends/base.rb

@@ -3,13 +3,15 @@
 module Deimos
   # Represents a field in the schema.
   class SchemaField
-    attr_accessor :name, :type
+    attr_accessor :name, :type, :enum_values
 
     # @param name [String]
     # @param type [Object]
-    def initialize(name, type)
+    # @param enum_values [Array<String>]
+    def initialize(name, type, enum_values=[])
       @name = name
       @type = type
+      @enum_values = enum_values
     end
   end
 
@@ -109,6 +111,17 @@ module Deimos
         raise NotImplementedError
       end
 
+      # Given a field definition, return the SQL type that might be used in
+      # ActiveRecord table creation - e.g. for Avro, a `long` type would
+      # return `:bigint`. There are also special values that need to be returned:
+      # `:array`, `:map` and `:record`, for types representing those structures.
+      # `:enum` is also recognized.
+      # @param field [SchemaField]
+      # @return [Symbol]
+      def sql_type(field)
+        raise NotImplementedError
+      end
+
       # Encode a message key. To be defined by subclass.
       # @param key [String|Hash] the value to use as the key.
       # @param key_id [Symbol|String] the field name of the key.

data/lib/deimos/tracing/provider.rb

@@ -1,6 +1,5 @@
 # frozen_string_literal: true
 
-# rubocop:disable Lint/UnusedMethodArgument
 module Deimos
   module Tracing
     # Base class for all tracing providers.
@@ -28,4 +27,3 @@ module Deimos
     end
   end
 end
-# rubocop:enable Lint/UnusedMethodArgument

data/lib/deimos/utils/db_poller.rb

@@ -0,0 +1,149 @@
+# frozen_string_literal: true
+
+require 'deimos/poll_info'
+require 'sigurd/executor'
+require 'sigurd/signal_handler'
+
+module Deimos
+  module Utils
+    # Class which continually polls the database and sends Kafka messages.
+    class DbPoller
+      BATCH_SIZE = 1000
+
+      # Needed for Executor so it can identify the worker
+      attr_reader :id
+
+      # Begin the DB Poller process.
+      def self.start!
+        if Deimos.config.db_poller_objects.empty?
+          raise('No pollers configured!')
+        end
+
+        pollers = Deimos.config.db_poller_objects.map do |poller_config|
+          self.new(poller_config)
+        end
+        executor = Sigurd::Executor.new(pollers,
+                                        sleep_seconds: 5,
+                                        logger: Deimos.config.logger)
+        signal_handler = Sigurd::SignalHandler.new(executor)
+        signal_handler.run!
+      end
+
+      # @param config [Deimos::Configuration::ConfigStruct]
+      def initialize(config)
+        @config = config
+        @id = SecureRandom.hex
+        begin
+          @producer = @config.producer_class.constantize
+        rescue NameError
+          raise "Class #{@config.producer_class} not found!"
+        end
+        unless @producer < Deimos::ActiveRecordProducer
+          raise "Class #{@producer.class.name} is not an ActiveRecordProducer!"
+        end
+      end
+
+      # Start the poll:
+      # 1) Grab the current PollInfo from the database indicating the last
+      # time we ran
+      # 2) On a loop, process all the recent updates between the last time
+      # we ran and now.
+      def start
+        # Don't send asynchronously
+        if Deimos.config.producers.backend == :kafka_async
+          Deimos.config.producers.backend = :kafka
+        end
+        Deimos.config.logger.info('Starting...')
+        @signal_to_stop = false
+        retrieve_poll_info
+        loop do
+          if @signal_to_stop
+            Deimos.config.logger.info('Shutting down')
+            break
+          end
+          process_updates
+          sleep 0.1
+        end
+      end
+
+      # Grab the PollInfo or create if it doesn't exist.
+      def retrieve_poll_info
+        ActiveRecord::Base.connection.reconnect!
+        new_time = @config.start_from_beginning ? Time.new(0) : Time.zone.now
+        @info = Deimos::PollInfo.find_by_producer(@config.producer_class) ||
+                Deimos::PollInfo.create!(producer: @config.producer_class,
+                                         last_sent: new_time,
+                                         last_sent_id: 0)
+      end
+
+      # Stop the poll.
+      def stop
+        Deimos.config.logger.info('Received signal to stop')
+        @signal_to_stop = true
+      end
+
+      # Indicate whether this current loop should process updates. Most loops
+      # will busy-wait (sleeping 0.1 seconds) until it's ready.
+      # @return [Boolean]
+      def should_run?
+        Time.zone.now - @info.last_sent - @config.delay_time >= @config.run_every
+      end
+
+      # @param record [ActiveRecord::Base]
+      # @return [ActiveSupport::TimeWithZone]
+      def last_updated(record)
+        record.public_send(@config.timestamp_column)
+      end
+
+      # Send messages for updated data.
+      def process_updates
+        return unless should_run?
+
+        time_from = @config.full_table ? Time.new(0) : @info.last_sent.in_time_zone
+        time_to = Time.zone.now - @config.delay_time
+        Deimos.config.logger.info("Polling #{@producer.topic} from #{time_from} to #{time_to}")
+        message_count = 0
+        batch_count = 0
+
+        # poll_query gets all the relevant data from the database, as defined
+        # by the producer itself.
+        loop do
+          Deimos.config.logger.debug("Polling #{@producer.topic}, batch #{batch_count + 1}")
+          batch = fetch_results(time_from, time_to).to_a
+          break if batch.empty?
+
+          batch_count += 1
+          process_batch(batch)
+          message_count += batch.size
+          time_from = last_updated(batch.last)
+        end
+        Deimos.config.logger.info("Poll #{@producer.topic} complete at #{time_to} (#{message_count} messages, #{batch_count} batches}")
+      end
+
+      # @param time_from [ActiveSupport::TimeWithZone]
+      # @param time_to [ActiveSupport::TimeWithZone]
+      # @return [ActiveRecord::Relation]
+      def fetch_results(time_from, time_to)
+        id = @producer.config[:record_class].primary_key
+        quoted_timestamp = ActiveRecord::Base.connection.quote_column_name(@config.timestamp_column)
+        quoted_id = ActiveRecord::Base.connection.quote_column_name(id)
+        @producer.poll_query(time_from: time_from,
+                             time_to: time_to,
+                             column_name: @config.timestamp_column,
+                             min_id: @info.last_sent_id).
+          limit(BATCH_SIZE).
+          order("#{quoted_timestamp}, #{quoted_id}")
+      end
+
+      # @param batch [Array<ActiveRecord::Base>]
+      def process_batch(batch)
+        record = batch.last
+        id_method = record.class.primary_key
+        last_id = record.public_send(id_method)
+        last_updated_at = last_updated(record)
+        @producer.send_events(batch)
+        @info.update_attributes!(last_sent: last_updated_at, last_sent_id: last_id)
+      end
+    end
+  end
+end