deimos-ruby 1.6.1 → 1.8.0.pre.beta1

data/lib/deimos/active_record_consume/batch_consumption.rb

@@ -0,0 +1,159 @@
+# frozen_string_literal: true
+
+require 'deimos/active_record_consume/batch_slicer'
+require 'deimos/utils/deadlock_retry'
+require 'deimos/message'
+
+module Deimos
+  module ActiveRecordConsume
+    # Methods for consuming batches of messages and saving them to the database
+    # in bulk ActiveRecord operations.
+    module BatchConsumption
+      # Handle a batch of Kafka messages. Batches are split into "slices",
+      # which are groups of independent messages that can be processed together
+      # in a single database operation.
+      # If two messages in a batch have the same key, we cannot process them
+      # in the same operation as they would interfere with each other. Thus
+      # they are split
+      # @param payloads [Array<Hash>] Decoded payloads.
+      # @param metadata [Hash] Information about batch, including keys.
+      def consume_batch(payloads, metadata)
+        messages = payloads.
+          zip(metadata[:keys]).
+          map { |p, k| Deimos::Message.new(p, nil, key: k) }
+
+        tags = %W(topic:#{metadata[:topic]})
+
+        Deimos.instrument('ar_consumer.consume_batch', tags) do
+          # The entire batch should be treated as one transaction so that if
+          # any message fails, the whole thing is rolled back or retried
+          # if there is deadlock
+          Deimos::Utils::DeadlockRetry.wrap(tags) do
+            if @compacted || self.class.config[:no_keys]
+              update_database(compact_messages(messages))
+            else
+              uncompacted_update(messages)
+            end
+          end
+        end
+      end
+
+      # Get unique key for the ActiveRecord instance from the incoming key.
+      # Override this method (with super) to customize the set of attributes that
+      # uniquely identifies each record in the database.
+      # @param key [String] The encoded key.
+      # @return [Hash] The key attributes.
+      def record_key(key)
+        decoded_key = decode_key(key)
+
+        if decoded_key.nil?
+          {}
+        elsif decoded_key.is_a?(Hash)
+          @key_converter.convert(decoded_key)
+        else
+          { @klass.primary_key => decoded_key }
+        end
+      end
+
+    protected
+
+      # Perform database operations for a batch of messages without compaction.
+      # All messages are split into slices containing only unique keys, and
+      # each slice is handles as its own batch.
+      # @param messages [Array<Message>] List of messages.
+      def uncompacted_update(messages)
+        BatchSlicer.
+          slice(messages).
+          each(&method(:update_database))
+      end
+
+      # Perform database operations for a group of messages.
+      # All messages with payloads are passed to upsert_records.
+      # All tombstones messages are passed to remove_records.
+      # @param messages [Array<Message>] List of messages.
+      def update_database(messages)
+        # Find all upserted records (i.e. that have a payload) and all
+        # deleted record (no payload)
+        removed, upserted = messages.partition(&:tombstone?)
+
+        upsert_records(upserted) if upserted.any?
+        remove_records(removed) if removed.any?
+      end
+
+      # Upsert any non-deleted records
+      # @param messages [Array<Message>] List of messages for a group of
+      # records to either be updated or inserted.
+      def upsert_records(messages)
+        key_cols = key_columns(messages)
+
+        # Create payloads with payload + key attributes
+        upserts = messages.map do |m|
+          record_attributes(m.payload, m.key)&.
+            merge(record_key(m.key))
+        end
+
+        # If overridden record_attributes indicated no record, skip
+        upserts.compact!
+
+        options = if key_cols.empty?
+                    {} # Can't upsert with no key, just do regular insert
+                  else
+                    {
+                      on_duplicate_key_update: {
+                        # conflict_target must explicitly list the columns for
+                        # Postgres and SQLite. Not required for MySQL, but this
+                        # ensures consistent behaviour.
+                        conflict_target: key_cols,
+                        columns: :all
+                      }
+                    }
+                  end
+
+        @klass.import!(upserts, options)
+      end
+
+      # Delete any records with a tombstone.
+      # @param messages [Array<Message>] List of messages for a group of
+      # deleted records.
+      def remove_records(messages)
+        clause = deleted_query(messages)
+
+        clause.delete_all
+      end
+
+      # Create an ActiveRecord relation that matches all of the passed
+      # records. Used for bulk deletion.
+      # @param records [Array<Message>] List of messages.
+      # @return ActiveRecord::Relation Matching relation.
+      def deleted_query(records)
+        keys = records.
+          map { |m| record_key(m.key)[@klass.primary_key] }.
+          reject(&:nil?)
+
+        @klass.unscoped.where(@klass.primary_key => keys)
+      end
+
+      # Get the set of attribute names that uniquely identify messages in the
+      # batch. Requires at least one record.
+      # @param records [Array<Message>] Non-empty list of messages.
+      # @return [Array<String>] List of attribute names.
+      # @raise If records is empty.
+      def key_columns(records)
+        raise 'Cannot determine key from empty batch' if records.empty?
+
+        first_key = records.first.key
+        record_key(first_key).keys
+      end
+
+      # Compact a batch of messages, taking only the last message for each
+      # unique key.
+      # @param batch [Array<Message>] Batch of messages.
+      # @return [Array<Message>] Compacted batch.
+      def compact_messages(batch)
+        return batch unless batch.first&.key.present?
+
+        batch.reverse.uniq(&:key).reverse!
+      end
+    end
+  end
+end

data/lib/deimos/active_record_consume/batch_slicer.rb

@@ -0,0 +1,27 @@
+# frozen_string_literal: true
+
+module Deimos
+  module ActiveRecordConsume
+    # Helper class for breaking down batches into independent groups for
+    # processing
+    class BatchSlicer
+      # Split the batch into a series of independent slices. Each slice contains
+      # messages that can be processed in any order (i.e. they have distinct
+      # keys). Messages with the same key will be separated into different
+      # slices that maintain the correct order.
+      # E.g. Given messages A1, A2, B1, C1, C2, C3, they will be sliced as:
+      # [[A1, B1, C1], [A2, C2], [C3]]
+      def self.slice(messages)
+        ops = messages.group_by(&:key)
+
+        # Find maximum depth
+        depth = ops.values.map(&:length).max || 0
+
+        # Generate slices for each depth
+        depth.times.map do |i|
+          ops.values.map { |arr| arr.dig(i) }.compact
+        end
+      end
+    end
+  end
+end

data/lib/deimos/active_record_consume/message_consumption.rb

@@ -0,0 +1,58 @@
+# frozen_string_literal: true
+
+module Deimos
+  module ActiveRecordConsume
+    # Methods for consuming individual messages and saving them to the database
+    # as ActiveRecord instances.
+    module MessageConsumption
+      # Find the record specified by the given payload and key.
+      # Default is to use the primary key column and the value of the first
+      # field in the key.
+      # @param klass [Class < ActiveRecord::Base]
+      # @param _payload [Hash]
+      # @param key [Object]
+      # @return [ActiveRecord::Base]
+      def fetch_record(klass, _payload, key)
+        klass.unscoped.where(klass.primary_key => key).first
+      end
+
+      # Assign a key to a new record.
+      # @param record [ActiveRecord::Base]
+      # @param _payload [Hash]
+      # @param key [Object]
+      def assign_key(record, _payload, key)
+        record[record.class.primary_key] = key
+      end
+
+      # :nodoc:
+      def consume(payload, metadata)
+        key = metadata.with_indifferent_access[:key]
+        klass = self.class.config[:record_class]
+        record = fetch_record(klass, (payload || {}).with_indifferent_access, key)
+        if payload.nil?
+          destroy_record(record)
+          return
+        end
+        if record.blank?
+          record = klass.new
+          assign_key(record, payload, key)
+        end
+        attrs = record_attributes(payload.with_indifferent_access, key)
+        # don't use attributes= - bypass Rails < 5 attr_protected
+        attrs.each do |k, v|
+          record.send("#{k}=", v)
+        end
+        record.created_at ||= Time.zone.now if record.respond_to?(:created_at)
+        record.updated_at = Time.zone.now if record.respond_to?(:updated_at)
+        record.save!
+      end
+
+      # Destroy a record that received a null payload. Override if you need
+      # to do something other than a straight destroy (e.g. mark as archived).
+      # @param record [ActiveRecord::Base]
+      def destroy_record(record)
+        record&.destroy
+      end
+    end
+  end
+end

data/lib/deimos/active_record_consume/schema_model_converter.rb

@@ -0,0 +1,52 @@
+# frozen_string_literal: true
+
+module Deimos
+  module ActiveRecordConsume
+    # Convert a message with a schema to an ActiveRecord model
+    class SchemaModelConverter
+      # Create new converter
+      # @param decoder [SchemaBackends::Base] Incoming message schema.
+      # @param klass [ActiveRecord::Base] Model to map to.
+      def initialize(decoder, klass)
+        @decoder = decoder
+        @klass = klass
+      end
+
+      # Convert a message from a decoded hash to a set of ActiveRecord
+      # attributes. Attributes that don't exist in the model will be ignored.
+      # @param payload [Hash] Decoded message payload.
+      # @return [Hash] Model attributes.
+      def convert(payload)
+        attributes = {}
+        @decoder.schema_fields.each do |field|
+          column = @klass.columns.find { |c| c.name == field.name }
+          next if column.nil?
+          next if %w(updated_at created_at).include?(field.name)
+
+          attributes[field.name] = _coerce_field(column, payload[field.name])
+        end
+        attributes
+      end
+
+    private
+
+      # @param column [ActiveRecord::ConnectionAdapters::Column]
+      # @param val [Object]
+      def _coerce_field(column, val)
+        return nil if val.nil?
+
+        if column.type == :datetime
+          int_val = begin
+                      val.is_a?(Integer) ? val : (val.is_a?(String) && Integer(val))
+                    rescue StandardError
+                      nil
+                    end
+
+          return Time.zone.at(int_val) if int_val
+        end
+
+        val
+      end
+    end
+  end
+end

data/lib/deimos/active_record_consumer.rb

@@ -1,101 +1,59 @@
 # frozen_string_literal: true
 
+require 'deimos/active_record_consume/batch_consumption'
+require 'deimos/active_record_consume/message_consumption'
+require 'deimos/active_record_consume/schema_model_converter'
 require 'deimos/consumer'
 
 module Deimos
-  # Consumer that automatically saves the payload into the database.
+  # Basic ActiveRecord consumer class. Consumes messages and upserts them to
+  # the database. For tombstones (null payloads), deletes corresponding
+  # records from the database. Can operate in either message-by-message mode
+  # or in batch mode.
+  #
+  # In batch mode, ActiveRecord callbacks will be skipped and messages will
+  # be batched to minimize database calls.
+
+  # To configure batch vs. message mode, change the delivery mode of your
+  # Phobos listener.
+  # Message-by-message -> use `delivery: message` or `delivery: batch`
+  # Batch -> use `delivery: inline_batch`
   class ActiveRecordConsumer < Consumer
+    include ActiveRecordConsume::MessageConsumption
+    include ActiveRecordConsume::BatchConsumption
+
     class << self
       # param klass [Class < ActiveRecord::Base] the class used to save to the
       # database.
       def record_class(klass)
         config[:record_class] = klass
       end
-    end
 
-    # Find the record specified by the given payload and key.
-    # Default is to use the primary key column and the value of the first
-    # field in the key.
-    # @param klass [Class < ActiveRecord::Base]
-    # @param _payload [Hash]
-    # @param key [Object]
-    # @return [ActiveRecord::Base]
-    def fetch_record(klass, _payload, key)
-      klass.unscoped.where(klass.primary_key => key).first
+      # param val [Boolean] Turn pre-compaction of the batch on or off. If true,
+      # only the last message for each unique key in a batch is processed.
+      def compacted(val)
+        config[:compacted] = val
+      end
     end
 
-    # Assign a key to a new record.
-    # @param record [ActiveRecord::Base]
-    # @param _payload [Hash]
-    # @param key [Object]
-    def assign_key(record, _payload, key)
-      record[record.class.primary_key] = key
-    end
+    # Setup
+    def initialize
+      @klass = self.class.config[:record_class]
+      @converter = ActiveRecordConsume::SchemaModelConverter.new(self.class.decoder, @klass)
 
-    # :nodoc:
-    def consume(payload, metadata)
-      key = metadata.with_indifferent_access[:key]
-      klass = self.class.config[:record_class]
-      record = fetch_record(klass, (payload || {}).with_indifferent_access, key)
-      if payload.nil?
-        destroy_record(record)
-        return
+      if self.class.config[:key_schema]
+        @key_converter = ActiveRecordConsume::SchemaModelConverter.new(self.class.key_decoder, @klass)
       end
-      if record.blank?
-        record = klass.new
-        assign_key(record, payload, key)
-      end
-      attrs = record_attributes(payload.with_indifferent_access)
-      # don't use attributes= - bypass Rails < 5 attr_protected
-      attrs.each do |k, v|
-        record.send("#{k}=", v)
-      end
-      record.created_at ||= Time.zone.now if record.respond_to?(:created_at)
-      record.updated_at ||= Time.zone.now if record.respond_to?(:updated_at)
-      record.save!
-    end
 
-    # Destroy a record that received a null payload. Override if you need
-    # to do something other than a straight destroy (e.g. mark as archived).
-    # @param record [ActiveRecord::Base]
-    def destroy_record(record)
-      record&.destroy
+      @compacted = self.class.config[:compacted] != false
     end
 
     # Override this method (with `super`) if you want to add/change the default
     # attributes set to the new/existing record.
     # @param payload [Hash]
-    def record_attributes(payload)
-      klass = self.class.config[:record_class]
-      attributes = {}
-      self.class.decoder.schema_fields.each do |field|
-        column = klass.columns.find { |c| c.name == field.name }
-        next if column.nil?
-        next if %w(updated_at created_at).include?(field.name)
-
-        attributes[field.name] = _coerce_field(column, payload[field.name])
-      end
-      attributes
-    end
-
-  private
-
-    # @param column [ActiveRecord::ConnectionAdapters::Column]
-    # @param val [Object]
-    def _coerce_field(column, val)
-      return nil if val.nil?
-
-      if column.type == :datetime
-        int_val = begin
-                    val.is_a?(Integer) ? val : (val.is_a?(String) && Integer(val))
-                  rescue StandardError
-                    nil
-                  end
-
-        return Time.zone.at(int_val) if int_val
-      end
-
-      val
+    # @param _key [String]
+    def record_attributes(payload, _key=nil)
+      @converter.convert(payload)
     end
   end
 end

data/lib/deimos/active_record_producer.rb

@@ -59,6 +59,29 @@ module Deimos
           k.to_sym != :payload_key && !fields.map(&:name).include?(k)
         end
       end
+
+      # Query to use when polling the database with the DbPoller. Add
+      # includes, joins, or wheres as necessary, or replace entirely.
+      # @param time_from [Time] the time to start the query from.
+      # @param time_to [Time] the time to end the query.
+      # @param column_name [Symbol] the column name to look for.
+      # @param min_id [Numeric] the minimum ID (i.e. all IDs must be greater
+      # than this value).
+      # @return [ActiveRecord::Relation]
+      def poll_query(time_from:, time_to:, column_name: :updated_at, min_id:)
+        klass = config[:record_class]
+        table = ActiveRecord::Base.connection.quote_table_name(klass.table_name)
+        column = ActiveRecord::Base.connection.quote_column_name(column_name)
+        primary = ActiveRecord::Base.connection.quote_column_name(klass.primary_key)
+        klass.where(
+          "((#{table}.#{column} = ? AND #{table}.#{primary} > ?) \
+           OR #{table}.#{column} > ?) AND #{table}.#{column} <= ?",
+          time_from,
+          min_id,
+          time_from,
+          time_to
+        )
+      end
     end
   end
 end