From c2bd7a4ad275a9a75264828084c1c10a714f85a0 Mon Sep 17 00:00:00 2001
From: Brad Charna <brad.charna@clutter.com>
Date: Fri, 11 Oct 2019 15:35:57 -0700
Subject: [PATCH] Add formatter for the logstash-logger gem

Allows for the added payload fields in lograge to surface in logs
outside of just controller logs.
---
 lib/lograge.rb                                | 20 +++++++++++++++-
 .../log_subscribers/action_controller.rb      |  4 ++++
 lib/lograge/log_subscribers/base.rb           |  4 ++++
 .../merged_with_lograge_payload.rb            | 23 +++++++++++++++++++
 spec/lograge_spec.rb                          |  4 ++++
 5 files changed, 54 insertions(+), 1 deletion(-)
 create mode 100644 lib/lograge/logstash_logger_formatters/merged_with_lograge_payload.rb

diff --git a/lib/lograge.rb b/lib/lograge.rb
index bde44f4f..e76b3b44 100644
--- a/lib/lograge.rb
+++ b/lib/lograge.rb
@@ -172,15 +172,33 @@ def setup_custom_payload
   end
 
   def extend_base_class(klass)
-    append_payload_method = klass.instance_method(:append_info_to_payload)
     custom_payload_method = lograge_config.custom_payload_method
+    define_append_info_to_payload_method(klass, custom_payload_method)
+    prepend_set_request_store_before_action(klass, custom_payload_method)
+  end
 
+  def define_append_info_to_payload_method(klass, custom_payload_method)
+    append_payload_method = klass.instance_method(:append_info_to_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
   end
 
+  def prepend_set_request_store_before_action(klass, custom_payload_method)
+    klass.class_eval do
+      prepend_before_action do |controller|
+        # TODO: it's not ideal that the user's custom_payload method would be called
+        # twice (below and in the append_info_to_payload).
+        # Would removing calling it in append_info_to_payload and accessing the
+        # payload via the request store in Lograge::LogSubscribers::Base#custom_options
+        # be okay? (i.e. modify the line
+        # `options.merge event.payload[:custom_payload] || {}` to read from request store)
+        RequestStore.store[:lograge_custom_payload] = custom_payload_method.call(controller)
+      end
+    end
+  end
+
   def set_lograge_log_options
     Lograge.logger = lograge_config.logger
     Lograge.custom_options = lograge_config.custom_options
diff --git a/lib/lograge/log_subscribers/action_controller.rb b/lib/lograge/log_subscribers/action_controller.rb
index 007239bc..032acefc 100644
--- a/lib/lograge/log_subscribers/action_controller.rb
+++ b/lib/lograge/log_subscribers/action_controller.rb
@@ -5,6 +5,10 @@ def process_action(event)
         process_main_event(event)
       end
 
+      def start_processing(event)
+        RequestStore.store[:lograge_event_payload] = event.payload.without(:headers, :params)
+      end
+
       def redirect_to(event)
         RequestStore.store[:lograge_location] = event.payload[:location]
       end
diff --git a/lib/lograge/log_subscribers/base.rb b/lib/lograge/log_subscribers/base.rb
index ae62623a..9a3b6ceb 100644
--- a/lib/lograge/log_subscribers/base.rb
+++ b/lib/lograge/log_subscribers/base.rb
@@ -58,6 +58,10 @@ def get_error_status_code(exception)
       end
 
       def custom_options(event)
+        # TODO: unsure if MergedWithLogragePayload should care about the user's
+        # custom_options for now (if it did, it seems that it would only make sense
+        # to call custom_options twice, one for the process_action event as is, plus one for
+        # the newly listened to start_processing.)
         options = Lograge.custom_options(event) || {}
         options.merge event.payload[:custom_payload] || {}
       end
diff --git a/lib/lograge/logstash_logger_formatters/merged_with_lograge_payload.rb b/lib/lograge/logstash_logger_formatters/merged_with_lograge_payload.rb
new file mode 100644
index 00000000..66ee7d72
--- /dev/null
+++ b/lib/lograge/logstash_logger_formatters/merged_with_lograge_payload.rb
@@ -0,0 +1,23 @@
+begin
+  require 'logstash-logger'
+  require 'logstash-event'
+rescue LoadError
+  puts 'You need to install the logstash-event and logstash-logger gems ' \
+    'to use Lograge::LogStashLoggerFormatters::MergedWithLogragePayload.'
+  raise
+end
+
+module Lograge
+  module LogStashLoggerFormatters
+    class MergedWithLogragePayload < LogStashLogger::Formatter::Base
+      private
+
+      def format_event(event)
+        lograge_event_payload = RequestStore.store[:lograge_event_payload] || {}
+        lograge_event_payload.merge!(RequestStore.store[:lograge_custom_payload] || {})
+        event.overwrite(event.fields.merge(lograge_event_payload))
+        "#{event.to_json}\n"
+      end
+    end
+  end
+end
diff --git a/spec/lograge_spec.rb b/spec/lograge_spec.rb
index ce3afeb5..451ade43 100644
--- a/spec/lograge_spec.rb
+++ b/spec/lograge_spec.rb
@@ -113,6 +113,10 @@ def append_info_to_payload(payload)
         def current_user_id
           '24601'
         end
+
+        def self.prepend_before_action(&block)
+          block.call(self.new)
+        end
       end
     end
     let(:payload) { { timestamp: Date.parse('5-11-1955') } }