From d236967f3a1d39d080e0844027c0d9628db68f28 Mon Sep 17 00:00:00 2001
From: Thorsten Ball <mrnugget@gmail.com>
Date: Fri, 13 Oct 2017 15:09:52 +0200
Subject: [PATCH] Allow multiple base controller classes to be set

This changes the `base_controller_class` configuration option to also
accept an array of class names, so that the custom payload is appended
in multiple controller base classes.

Why? Because some applications have controllers that inherit from
multiple base classes. Example: `ActionController::API` and
`ActionController::Base`. If you have an an application with an API and
also an admin panel, that would certainly be the case.

Until now, using a custom payload was only possible in the controllers
that inherit from the configured `base_controller_class`. With this
change, the custom payload feature can be used in all controllers -- if
their base classes are set.
---
 README.md            | 13 ++++++++++++-
 lib/lograge.rb       | 18 +++++++++++++++---
 spec/lograge_spec.rb |  9 ++++++++-
 3 files changed, 35 insertions(+), 5 deletions(-)

diff --git a/README.md b/README.md
index e1845ca0..98fd2956 100644
--- a/README.md
+++ b/README.md
@@ -65,7 +65,9 @@ Rails.application.configure do
 end
 ```
 
-If you're using Rails 5's API-only mode and inherit from `ActionController::API`:
+If you're using Rails 5's API-only mode and inherit from
+`ActionController::API`, you must define it as the controller base class which
+lograge will patch:
 
 ```ruby
 # config/initializers/lograge.rb
@@ -74,6 +76,15 @@ Rails.application.configure do
 end
 ```
 
+If you use multiple base controller classes in your application, specify an array:
+
+```ruby
+# config/initializers/lograge.rb
+Rails.application.configure do
+  config.lograge.base_controller_class = ['ActionController::API', 'ActionController::Base']
+end
+```
+
 You can also add a hook for own custom data
 
 ```ruby
diff --git a/lib/lograge.rb b/lib/lograge.rb
index 914ca7eb..b7376f5d 100644
--- a/lib/lograge.rb
+++ b/lib/lograge.rb
@@ -144,11 +144,23 @@ def attach_to_action_controller
 
   def setup_custom_payload
     return unless lograge_config.custom_payload_method.respond_to?(:call)
-    base_controller_class = lograge_config.base_controller_class.try(:constantize) || ActionController::Base
-    append_payload_method = base_controller_class.instance_method(:append_info_to_payload)
+
+    base_controller_classes = Array(lograge_config.base_controller_class)
+    base_controller_classes.map! { |klass| klass.try(:constantize) }
+    if base_controller_classes.empty?
+      base_controller_classes << ActionController::Base
+    end
+
+    base_controller_classes.each do |base_controller_class|
+      extend_base_controller_class(base_controller_class)
+    end
+  end
+
+  def extend_base_controller_class(klass)
+    append_payload_method = klass.instance_method(:append_info_to_payload)
     custom_payload_method = lograge_config.custom_payload_method
 
-    base_controller_class.send(:define_method, :append_info_to_payload) do |payload|
+    klass.send(:define_method, :append_info_to_payload) do |payload|
       append_payload_method.bind(self).call(payload)
       payload[:custom_payload] = custom_payload_method.call(self)
     end
diff --git a/spec/lograge_spec.rb b/spec/lograge_spec.rb
index 9f0fd553..ce3afeb5 100644
--- a/spec/lograge_spec.rb
+++ b/spec/lograge_spec.rb
@@ -129,11 +129,12 @@ def current_user_id
 
     context 'when base_controller_class option is set' do
       let(:controller_class) { 'ActionController::API' }
+      let(:base_controller_class) { controller_class }
       let(:app_config) do
         config_obj = ActiveSupport::OrderedOptions.new.tap do |config|
           config.action_dispatch = double(rack_cache: false)
           config.lograge = Lograge::OrderedOptions.new
-          config.lograge.base_controller_class = controller_class
+          config.lograge.base_controller_class = base_controller_class
           config.lograge.custom_payload do |c|
             { user_id: c.current_user_id }
           end
@@ -142,6 +143,12 @@ def current_user_id
       end
 
       it { should eq(payload.merge(appended: true, custom_payload: { user_id: '24601' })) }
+
+      context 'when base_controller_class is an array' do
+        let(:base_controller_class) { [controller_class] }
+
+        it { should eq(payload.merge(appended: true, custom_payload: { user_id: '24601' })) }
+      end
     end
   end