From 7c3d7e2b70547252219c5fbde951f4f681f742e7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jan=20Strn=C3=A1dek?= <jan.strnadek@gmail.com>
Date: Sun, 10 Sep 2017 20:46:29 +0200
Subject: [PATCH] Fix #29 (#33)

* Fix #29

* Update request_logger.rb

If it is possible use status code from error, otherwise send 500.
---
 lib/grape_logging/loggers/response.rb         | 16 +++--
 .../middleware/request_logger.rb              | 70 +++++++++++++++----
 2 files changed, 68 insertions(+), 18 deletions(-)

diff --git a/lib/grape_logging/loggers/response.rb b/lib/grape_logging/loggers/response.rb
index 83a248a..d464b8d 100644
--- a/lib/grape_logging/loggers/response.rb
+++ b/lib/grape_logging/loggers/response.rb
@@ -6,14 +6,22 @@ module GrapeLogging
       end
 
       private
+
       # In some cases, response.body is not parseable by JSON.
       # For example, if you POST on a PUT endpoint, response.body is egal to """".
       # It's strange but it's the Grape behavior...
       def serialized_response_body(response)
-        begin
-          response.body.map{ |body| JSON.parse(body.to_s) }
-        rescue => e
-          response.body
+
+        if response.respond_to?(:body)
+          # Rack responses
+          begin
+            response.body.map{ |body| JSON.parse(body.to_s) }
+          rescue # No reason to have "=> e" here when we don't use it..
+            response.body
+          end
+        else
+          # Error & Exception responses
+          response
         end
       end
     end
diff --git a/lib/grape_logging/middleware/request_logger.rb b/lib/grape_logging/middleware/request_logger.rb
index 11ebdeb..a401ab1 100644
--- a/lib/grape_logging/middleware/request_logger.rb
+++ b/lib/grape_logging/middleware/request_logger.rb
@@ -9,6 +9,10 @@ module GrapeLogging
         GrapeLogging::Timings.append_db_runtime(event)
       end if defined?(ActiveRecord)
 
+      # Persist response status & response (body)
+      # to use int in parameters
+      attr_accessor :response_status, :response_body
+
       def initialize(app, options = {})
         super
 
@@ -23,34 +27,71 @@ module GrapeLogging
       def before
         reset_db_runtime
         start_time
-
         invoke_included_loggers(:before)
       end
 
-      def after
+      def after(status, response)
         stop_time
+
+        # Response status
+        @response_status = status
+        @response_body   = response
+
+        # Perform repotters
         @reporter.perform(collect_parameters)
+
+        # Invoke loggers
         invoke_included_loggers(:after)
         nil
       end
 
+      # Call stack and parse responses & status.
+      #
+      # @note Exceptions are logged as 500 status & re-raised.
       def call!(env)
-        super
+        @env = env
+
+        # Before hook
+        before
+
+        # Catch error
+        error = catch(:error) do
+          begin
+            @app_response = @app.call(@env)
+          rescue => e
+            # Log as 500 + message
+            after(e.respond_to?(:status) ? e.status : 500, e.message)
+
+            # Re-raise exception
+            raise e
+          end
+          nil
+        end
+
+        # Get status & response from app_response
+        # when no error occures.
+        if error
+          # Call with error & response
+          after(error[:status], error[:message])
+
+          # Throw again
+          throw(:error, error)
+        else
+          status, _, resp = *@app_response
+
+          # Call after hook properly
+          after(status, resp)
+        end
+
+        # Otherwise return original response
+        @app_response
       end
 
       protected
 
-      def response
-        begin
-          super
-        rescue
-          nil
-        end
-      end
-
       def parameters
         {
-          status: response.nil? ? 'fail' : response.status,
+          status: response_status,
           time: {
             total: total_runtime,
             db: db_runtime,
@@ -64,8 +105,9 @@ module GrapeLogging
       end
 
       private
+
       def request
-        @request ||= ::Rack::Request.new(env)
+        @request ||= ::Rack::Request.new(@env)
       end
 
       def total_runtime
@@ -95,7 +137,7 @@ module GrapeLogging
       def collect_parameters
         parameters.tap do |params|
           @included_loggers.each do |logger|
-            params.merge! logger.parameters(request, response) do |_, oldval, newval|
+            params.merge! logger.parameters(request, response_body) do |_, oldval, newval|
               oldval.respond_to?(:merge) ? oldval.merge(newval) : newval
             end
           end