deimos-temp-fork 0.0.1

data/lib/deimos/active_record_consume/batch_consumption.rb

@@ -0,0 +1,164 @@
+# 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|
+          attrs = if self.method(:record_attributes).parameters.size == 2
+                    record_attributes(m.payload, m.key)
+                  else
+                    record_attributes(m.payload)
+                  end
+
+          attrs&.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,79 @@
+# 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)
+        unless self.process_message?(payload)
+          Deimos.config.logger.debug(
+            message: 'Skipping processing of message',
+            payload: payload,
+            metadata: metadata
+          )
+          return
+        end
+
+        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
+
+        # for backwards compatibility
+        # TODO next major release we should deprecate this
+        attrs = if self.method(:record_attributes).parameters.size == 2
+                  record_attributes(payload.with_indifferent_access, key)
+                else
+                  record_attributes(payload.with_indifferent_access)
+                end
+        # don't use attributes= - bypass Rails < 5 attr_protected
+        attrs.each do |k, v|
+          record.send("#{k}=", v)
+        end
+        save_record(record)
+      end
+
+      # @param record [ActiveRecord::Base]
+      def save_record(record)
+        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

@@ -0,0 +1,67 @@
+# 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
+  # 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
+
+      # 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
+
+    # Setup
+    def initialize
+      @klass = self.class.config[:record_class]
+      @converter = ActiveRecordConsume::SchemaModelConverter.new(self.class.decoder, @klass)
+
+      if self.class.config[:key_schema]
+        @key_converter = ActiveRecordConsume::SchemaModelConverter.new(self.class.key_decoder, @klass)
+      end
+
+      @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]
+    # @param _key [String]
+    def record_attributes(payload, _key=nil)
+      @converter.convert(payload)
+    end
+
+    # Override this message to conditionally save records
+    # @param payload [Hash] The kafka message as a hash
+    # @return [Boolean] if true, record is created/update.
+    #   If false, record processing is skipped but message offset is still committed.
+    def process_message?(_payload)
+      true
+    end
+  end
+end

data/lib/deimos/active_record_producer.rb

@@ -0,0 +1,87 @@
+# frozen_string_literal: true
+
+require 'deimos/producer'
+
+module Deimos
+  # Class which automatically produces a record when given an ActiveRecord
+  # instance or a list of them. Just call `send_events` on a list of records
+  # and they will be auto-published. You can override `generate_payload`
+  # to make changes to the payload before it's published.
+  #
+  # You can also call this with a list of hashes representing attributes.
+  # This is common when using activerecord-import.
+  class ActiveRecordProducer < Producer
+    class << self
+      # Indicate the class this producer is working on.
+      # @param klass [Class]
+      # @param refetch [Boolean] if true, and we are given a hash instead of
+      # a record object, refetch the record to pass into the `generate_payload`
+      # method.
+      def record_class(klass, refetch: true)
+        config[:record_class] = klass
+        config[:refetch_record] = refetch
+      end
+
+      # @param record [ActiveRecord::Base]
+      # @param force_send [Boolean]
+      def send_event(record, force_send: false)
+        send_events([record], force_send: force_send)
+      end
+
+      # @param records [Array<ActiveRecord::Base>]
+      # @param force_send [Boolean]
+      def send_events(records, force_send: false)
+        primary_key = config[:record_class]&.primary_key
+        messages = records.map do |record|
+          if record.respond_to?(:attributes)
+            attrs = record.attributes.with_indifferent_access
+          else
+            attrs = record.with_indifferent_access
+            if config[:refetch_record] && attrs[primary_key]
+              record = config[:record_class].find(attrs[primary_key])
+            end
+          end
+          generate_payload(attrs, record).with_indifferent_access
+        end
+        self.publish_list(messages, force_send: force_send)
+      end
+
+      # Generate the payload, given a list of attributes or a record..
+      # Can be overridden or added to by subclasses.
+      # @param attributes [Hash]
+      # @param _record [ActiveRecord::Base] May be nil if refetch_record
+      # is not set.
+      # @return [Hash]
+      def generate_payload(attributes, _record)
+        fields = self.encoder.schema_fields
+        payload = attributes.stringify_keys
+        payload.delete_if do |k, _|
+          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