1 files changed, 7 insertions, 14 deletions
diff --git a/lib/bulk_imports/pipeline.rb b/lib/bulk_imports/pipeline.rb
index 06b81b5da14..1d55ad95887 100644
--- a/lib/bulk_imports/pipeline.rb
+++ b/lib/bulk_imports/pipeline.rb
@@ -4,12 +4,17 @@ module BulkImports
   module Pipeline
     extend ActiveSupport::Concern
     include Gitlab::ClassAttributes
+    include Runner
 
-    included do
-      include Runner
+    def initialize(context)
+      @context = context
+    end
 
+    included do
       private
 
+      attr_reader :context
+
       def extractor
         @extractor ||= instantiate(self.class.get_extractor)
       end
@@ -22,10 +27,6 @@ module BulkImports
         @loaders ||= instantiate(self.class.get_loader)
       end
 
-      def after_run
-        @after_run ||= self.class.after_run_callback
-      end
-
       def pipeline
         @pipeline ||= self.class.name
       end
@@ -52,10 +53,6 @@ module BulkImports
         class_attributes[:loader] = { klass: klass, options: options }
       end
 
-      def after_run(&block)
-        class_attributes[:after_run] = block
-      end
-
       def get_extractor
         class_attributes[:extractor]
       end
@@ -68,10 +65,6 @@ module BulkImports
         class_attributes[:loader]
       end
 
-      def after_run_callback
-        class_attributes[:after_run]
-      end
-
       def abort_on_failure!
         class_attributes[:abort_on_failure] = true
       end