lib/marlowe/rails.rb in marlowe-1.0.3 vs lib/marlowe/rails.rb in marlowe-2.0

@@ -1,12 +1,29 @@
+# frozen_string_literal: true
+
 module Marlowe
-  class Railtie < Rails::Railtie
+  class Railtie < Rails::Railtie # :nodoc:
     initializer 'marlowe.configure_rails_initialization' do
-      app.middleware.insert_before Rails::Rack::Logger, Marlowe::Middleware, correlation_header: Rails.application.config.try(:marlowe_correlation_header)
+      config = app.config
+
+      opts = {
+        header: config.try(:marlowe_header) || config.try(:marlowe_correlation_header),
+        handler: config.try(:marlowe_request_id_handler),
+        return: config.try(:marlowe_return_request_id),
+        action_dispatch: config.try(:marlowe_replace_action_dispatch_request_id)
+      }.compact
+
+      if opts[:action_dispatch]
+        app.middleware.insert_before ActionDispatch::RequestId,
+          Marlowe::Middleware, opts
+        app.middleware.delete ActionDispatch::RequestId
+      else
+        app.middleware.insert_before Rails::Rack::Logger, Marlowe::Middleware,
+          opts
+      end
     end
 
-    #:nodoc:
-    def app
+    def app #:nodoc:
       Rails.application
     end
   end
 end