CCOL-2039: Post process valid and invalid records from batch consumption

.github/workflows/ci.yml

@@ -17,7 +17,7 @@ jobs:
       BUNDLE_WITHOUT: development:test
 
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v3
 
       - name: Set up Ruby 2.7
         uses: ruby/setup-ruby@v1
@@ -39,7 +39,7 @@ jobs:
         ruby: [ '2.6', '2.7', '3.0', '3.1' ]
 
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v3
       - uses: ruby/setup-ruby@v1
         with:
           ruby-version: ${{ matrix.ruby }}

CHANGELOG.md

@@ -12,6 +12,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Feature: Add individual configuration for custom `bulk_import_id_generator` proc per consumer
 - Feature: Add global `replace_assocations` value for  for all consumers
 - Feature: Add individual `replace_assocations` value for for individual consumers
+- Feature: `should_consume?` method accepts BatchRecord associations
+- Feature: Reintroduce `filter_records` for bulk filtering of records prior to insertion
+- Feature: Return valid and invalid records saved during consumption for further processing in `batch_consumption.valid_records` and `batch_consumption.invalid_records` ActiveSupport Notifications
 
 # 1.22.5 - 2023-07-18
 - Fix: Fixed buffer overflow crash with DB producer.

lib/deimos/active_record_consume/batch_consumption.rb

@@ -28,18 +28,13 @@ def consume_batch(payloads, metadata)
           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
+        @tags = %W(topic:#{metadata[:topic]})
+
+        Deimos.instrument('ar_consumer.consume_batch', @tags) do
+          if @compacted || self.class.config[:no_keys]
+            update_database(compact_messages(messages))
+          else
+            uncompacted_update(messages)
           end
         end
       end
@@ -93,8 +88,9 @@ def deleted_query(records)
       end
 
       # @param _record [ActiveRecord::Base]
+      # @param _associations [Hash]
       # @return [Boolean]
-      def should_consume?(_record)
+      def should_consume?(_record, _associations=nil)
         true
       end
 
@@ -155,19 +151,35 @@ def update_database(messages)
       # @return [void]
       def upsert_records(messages)
         record_list = build_records(messages)
-        record_list.filter!(self.method(:should_consume?).to_proc)
+        invalid = filter_records(record_list)
+        unless invalid.blank?
+          ActiveSupport::Notifications.instrument('batch_consumption.invalid_records', {
+                                                    records: invalid,
+                                                    consumer: self.class
+                                                  })
 
+        end
         return if record_list.empty?
 
         key_col_proc = self.method(:key_columns).to_proc
         col_proc = self.method(:columns).to_proc
 
         updater = MassUpdater.new(@klass,
+                                  @tags,
                                   key_col_proc: key_col_proc,
                                   col_proc: col_proc,
                                   replace_associations: self.class.replace_associations,
                                   bulk_import_id_generator: self.class.bulk_import_id_generator)
-        updater.mass_update(record_list)
+        ActiveSupport::Notifications.instrument('batch_consumption.valid_records', {
+                                                  records: updater.mass_update(record_list),
+                                                  consumer: self.class
+                                                })
+      end
+
+      # @param record_list [BatchRecordList]
+      # @return [Array<BatchRecord>]
+      def filter_records(record_list)
+        record_list.filter!(self.method(:should_consume?).to_proc)
       end
 
       # Process messages prior to saving to database

lib/deimos/active_record_consume/batch_record_list.rb

@@ -17,10 +17,19 @@ def initialize(records)
         self.bulk_import_column = records.first&.bulk_import_column&.to_sym
       end
 
-      # Filter out any invalid records.
+      # Filter and return removed invalid batch records by the specified method
       # @param method [Proc]
+      # @return [Array<BatchRecord>]
       def filter!(method)
-        self.batch_records.delete_if { |record| !method.call(record.record) }
+        self.batch_records, invalid = self.batch_records.partition do |batch_record|
+          case method.parameters.size
+          when 2
+            method.call(batch_record.record, batch_record.associations)
+          else
+            method.call(batch_record.record)
+          end
+        end
+        invalid
       end
 
       # Get the original ActiveRecord objects.

lib/deimos/active_record_consume/mass_updater.rb

@@ -19,8 +19,10 @@ def default_cols(klass)
       # @param key_col_proc [Proc<Class < ActiveRecord::Base>]
       # @param col_proc [Proc<Class < ActiveRecord::Base>]
       # @param replace_associations [Boolean]
-      def initialize(klass, key_col_proc: nil, col_proc: nil, replace_associations: true, bulk_import_id_generator: nil)
+      def initialize(klass, tags, key_col_proc: nil, col_proc: nil,
+                     replace_associations: true, bulk_import_id_generator: nil)
         @klass = klass
+        @tags = tags
         @replace_associations = replace_associations
         @bulk_import_id_generator = bulk_import_id_generator
 
@@ -83,9 +85,16 @@ def import_associations(record_list)
       end
 
       # @param record_list [BatchRecordList]
+      # @return [Array<ActiveRecord::Base>]
       def mass_update(record_list)
-        save_records_to_database(record_list)
-        import_associations(record_list) if record_list.associations.any?
+        # 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
+          save_records_to_database(record_list)
+          import_associations(record_list) if record_list.associations.any?
+        end
+        record_list.records
       end
 
     end

spec/active_record_batch_consumer_association_spec.rb

@@ -96,12 +96,17 @@ def publish_batch(messages)
         key_config plain: true
         record_class Widget
 
-        def should_consume?(record)
+        def should_consume?(record, associations)
           if self.should_consume_proc
-            return self.should_consume_proc.call(record)
+            case self.should_consume_proc.parameters.size
+            when 2
+              self.should_consume_proc.call(record, associations)
+            else
+              self.should_consume_proc.call(record)
+            end
+          else
+            true
           end
-
-          true
         end
 
         def record_attributes(payload, _key)
@@ -269,7 +274,7 @@ def columns(record_class)
 
     context 'with invalid models' do
       before(:each) do
-        consumer_class.should_consume_proc = proc { |val| val.some_int <= 10 }
+        consumer_class.should_consume_proc = proc { |record| record.some_int <= 10 }
       end
 
       it 'should only save valid models' do
@@ -280,5 +285,27 @@ def columns(record_class)
         expect(Widget.count).to eq(2)
       end
     end
+
+    context 'with invalid associations' do
+
+      before(:each) do
+        consumer_class.should_consume_proc = proc { |record, associations|
+          record.some_int <= 10 && associations['detail']['title'] != 'invalid'
+        }
+      end
+
+      it 'should only save valid associations' do
+        publish_batch([
+                        { key: 2,
+                          payload: { test_id: 'xyz', some_int: 5, title: 'valid' } },
+                        { key: 3,
+                          payload: { test_id: 'abc', some_int: 15, title: 'valid' } },
+                        { key: 4,
+                          payload: { test_id: 'abc', some_int: 9, title: 'invalid' } }
+                      ])
+        expect(Widget.count).to eq(2)
+        expect(Widget.second.some_int).to eq(5)
+      end
+    end
            end
 end