Allow handlers to use rescue_from in a way familiar to ActionControll…

…er users

ActiveSupport::Rescuable does close to what we need, but rather than calling render or redirect, handlers return an object (which may be a Twirp::Error). We refine ActiveSupport::Rescuable to return the result of the rescue_from block/method.

lib/twirp/rails.rb

@@ -15,6 +15,7 @@ class Error < StandardError; end
 require_relative "rails/configuration"
 require_relative "rails/dispatcher"
 require_relative "rails/engine"
+require_relative "rails/rescuable"
 require_relative "rails/handler"
 
 module Twirp

lib/twirp/rails/handler.rb

@@ -4,6 +4,8 @@ module Twirp
   module Rails
     class Handler
       include Twirp::Rails::Callbacks
+      include ActiveSupport::Rescuable
+      using Twirp::Rails::Rescuable
 
       attr_reader :request, :env
       attr_reader :action_name
@@ -34,7 +36,11 @@ def process_action(name)
         ActiveSupport::Notifications.instrument("handler_run_callbacks.twirp_rails", handler: self.class.name, action: action_name, env: @env, request: @request) do
           run_callbacks(:process_action) do
             ActiveSupport::Notifications.instrument("handler_run.twirp_rails", handler: self.class.name, action: action_name, env: @env, request: @request) do |payload|
-              payload[:response] = send_action(name)
+              payload[:response] = begin
+                send_action(name)
+              rescue => exception
+                rescue_with_handler_and_return(exception) || raise
+              end
             end
           end
         end

lib/twirp/rails/rescuable.rb

@@ -0,0 +1,31 @@
+# frozen_string_literal: true
+
+module Twirp
+  module Rails
+    module Rescuable
+      refine ::ActiveSupport::Rescuable::ClassMethods do
+        # A slightly altered version of ActiveSupport::Rescuable#rescue_with_handler
+        # that returns the result rather than the handled exception
+        def rescue_with_handler_and_return(exception, object: self, visited_exceptions: [])
+          visited_exceptions << exception
+
+          if (handler = handler_for_rescue(exception, object: object))
+            handler.call exception
+          elsif exception
+            if visited_exceptions.include?(exception.cause)
+              nil
+            else
+              rescue_with_handler(exception.cause, object: object, visited_exceptions: visited_exceptions)
+            end
+          end
+        end
+      end
+
+      refine ::ActiveSupport::Rescuable do
+        def rescue_with_handler_and_return(exception)
+          self.class.rescue_with_handler_and_return exception, object: self
+        end
+      end
+    end
+  end
+end

spec/rails_app/app/handlers/haberdasher_handler.rb

@@ -3,6 +3,10 @@ class HaberdasherHandler < Twirp::Rails::Handler
   before_action :track_request_ip
   before_action :reject_giant_hats
 
+  rescue_from "ArgumentError" do |error|
+    Twirp::Error.invalid_argument(error.message)
+  end
+
   def make_hat
     # We can return a Twirp::Error when appropriate
     if request.inches < 12

spec/requests/haberdasher_spec.rb

@@ -145,4 +145,23 @@ def make_hat_success_request
       expect(response.status).to eq(304)
     end
   end
+
+  describe "Rescuable" do
+    it "rescues from ArgumentError" do
+      size = Twirp::Example::Haberdasher::Size.new(inches: 100)
+
+      # Fake an exception
+      expect(Twirp::Example::Haberdasher::Hat).to receive(:new).and_raise(ArgumentError.new("is way too large"))
+
+      post "/twirp/twirp.example.haberdasher.Haberdasher/MakeHat",
+        params: size.to_proto, headers: {
+          :accept => "application/protobuf",
+          "Content-Type" => "application/protobuf"
+        }
+
+      expect(response.status).to eq(400)
+      expect(response.content_type).to eq("application/json")
+      expect(response.body).to eq('{"code":"invalid_argument","msg":"is way too large"}')
+    end
+  end
 end