1 files changed, 9 insertions, 5 deletions
diff --git a/app/models/data_list.rb b/app/models/data_list.rb
index adad8e3013e..e99364b2709 100644
--- a/app/models/data_list.rb
+++ b/app/models/data_list.rb
@@ -1,22 +1,26 @@
 # frozen_string_literal: true
 
 class DataList
-  def initialize(batch, data_fields_hash, klass)
+  def initialize(batch, data_fields_hash, data_fields_klass)
     @batch = batch
     @data_fields_hash = data_fields_hash
-    @klass = klass
+    @data_fields_klass = data_fields_klass
   end
 
   def to_array
-    [klass, columns, values]
+    [data_fields_klass, columns, values]
   end
 
   private
 
-  attr_reader :batch, :data_fields_hash, :klass
+  attr_reader :batch, :data_fields_hash, :data_fields_klass
 
   def columns
-    data_fields_hash.keys << 'service_id'
+    data_fields_hash.keys << data_fields_foreign_key
+  end
+
+  def data_fields_foreign_key
+    data_fields_klass.reflections['integration'].foreign_key
   end
 
   def values