From 8bdfb4ce5573d78f33a251a6df61fe58b6c6b6b6 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 cfa1a5bc..f9fda0fd 100644
--- a/lib/lograge.rb
+++ b/lib/lograge.rb
@@ -174,15 +174,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 5afde029..ff3a018b 100644
--- a/lib/lograge/log_subscribers/action_controller.rb
+++ b/lib/lograge/log_subscribers/action_controller.rb
@@ -7,6 +7,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 46cc7a30..fb6efb77 100644
--- a/lib/lograge/log_subscribers/base.rb
+++ b/lib/lograge/log_subscribers/base.rb
@@ -69,6 +69,10 @@ def extract_allocations(event)
       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 f97f4a72..7fba7862 100644
--- a/spec/lograge_spec.rb
+++ b/spec/lograge_spec.rb
@@ -118,6 +118,10 @@ def current_user_id
           '24601'
         end
 
+        def self.prepend_before_action(&block)
+          block.call(new)
+        end
+
         class << self
           def logger; end
         end