lib/rocket_job/batch/io.rb in rocketjob-4.1.0 vs lib/rocket_job/batch/io.rb in rocketjob-4.1.1

@@ -63,9 +63,21 @@
         count             = input(category).upload(file_name_or_io, file_name: file_name, **args, &block)
         self.record_count = (record_count || 0) + count
         count
       end
 
+      def upload_arel(arel, *column_names, category: :main, &block)
+        count             = input(category).upload_arel(arel, *column_names, &block)
+        self.record_count = (record_count || 0) + count
+        count
+      end
+
+      def upload_mongo_query(criteria, *column_names, category: :main, &block)
+        count             = input(category).upload_mongo_query(criteria, *column_names, &block)
+        self.record_count = (record_count || 0) + count
+        count
+      end
+
       # Upload the supplied slices for processing by workers
       #
       # Updates the record_count after adding the records
       #
       # Returns [Integer] the number of records uploaded