Extended rescue_from to take a block.

Author: dblock

README.markdown

@@ -93,6 +93,22 @@ The error format can be specified using `error_format`. Available formats are `:
       error_format :json
     end
 
+You can rescue all exceptions with a code block. The `rack_response` wrapper automatically sets the default error code and content-type.
+
+    class Twitter::API < Grape::API
+      rescue_from :all do |e|
+        rack_response({ :message => "rescued from #{e.class.name}" })
+      end
+    end
+
+You can also rescue specific exceptions with a code block and handle the Rack response at the lowest level.
+
+    class Twitter::API < Grape::API
+      rescue_from :all do |e|
+        Rack::Response.new([ e.message ], 500, { "Content-type" => "text/error" ).finish
+      end
+    end
+
 ## Note on Patches/Pull Requests
 
 * Fork the project.

lib/grape/api.rb

@@ -105,9 +105,15 @@ def default_error_status(new_status = nil)
       # @overload rescue_from(*exception_classes, options = {})
       #   @param [Array] exception_classes A list of classes that you want to rescue, or
       #     the symbol :all to rescue from all exceptions.
+      #   @param [Block] block Execution block to handle the given exception.
       #   @param [Hash] options Options for the rescue usage.
       #   @option options [Boolean] :backtrace Include a backtrace in the rescue response.
-      def rescue_from(*args)
+      def rescue_from(*args, &block)
+        if block_given?
+          args.each do |arg|
+            set(:rescue_handlers, { arg => block })
+          end
+        end
         set(:rescue_options, args.pop) if args.last.is_a?(Hash)
         set(:rescue_all, true) and return if args.include?(:all)
         set(:rescued_errors, args)
@@ -159,7 +165,7 @@ def http_basic(options = {}, &block)
 
       def http_digest(options = {}, &block)
         options[:realm] ||= "API Authorization"
-	options[:opaque] ||= "secret"
+        options[:opaque] ||= "secret"
         auth :http_digest, options, &block
       end
 
@@ -262,9 +268,10 @@ def build_endpoint(&block)
           :rescue_all => settings[:rescue_all], 
           :rescued_errors => settings[:rescued_errors], 
           :format => settings[:error_format] || :txt, 
-          :rescue_options => settings[:rescue_options]
-        b.use Rack::Auth::Basic, settings[:auth][:realm], &settings[:auth][:proc] if settings[:auth] && settings[:auth][:type] == :http_basic
-	b.use Rack::Auth::Digest::MD5, settings[:auth][:realm], settings[:auth][:opaque], &settings[:auth][:proc] if settings[:auth] && settings[:auth][:type] == :http_digest
+          :rescue_options => settings[:rescue_options],
+          :rescue_handlers => settings[:rescue_handlers] || {}
+        b.use Rack::Auth::Basic, settings[:auth][:realm], &settings[:auth][:proc] if settings[:auth] && settings[:auth][:type] == :http_basic	
+        b.use Rack::Auth::Digest::MD5, settings[:auth][:realm], settings[:auth][:opaque], &settings[:auth][:proc] if settings[:auth] && settings[:auth][:type] == :http_digest
         b.use Grape::Middleware::Prefixer, :prefix => prefix if prefix        
         b.use Grape::Middleware::Versioner, :versions => (version if version.is_a?(Array)) if version
         b.use Grape::Middleware::Formatter, :default_format => default_format || :json

lib/grape/middleware/error.rb

@@ -12,9 +12,9 @@ def default_options
         :default_message => "",
         :format => :txt,
         :formatters => {},
-
         :rescue_all => false, # true to rescue all exceptions        
-        :rescue_options => {:backtrace => false}, # true to display backtrace
+        :rescue_options => { :backtrace => false }, # true to display backtrace
+        :rescue_handlers => {}, # rescue handler blocks
         :rescued_errors => []
       }
       end
@@ -45,20 +45,29 @@ def call!(env)
           })
         rescue Exception => e
           raise unless options[:rescue_all] || (options[:rescued_errors] || []).include?(e.class)
-          error_response({ :message => e.message, :backtrace => e.backtrace })
+          handler = options[:rescue_handlers][e.class] || options[:rescue_handlers][:all]
+          handler.nil? ? handle_error(e) : self.instance_exec(e, &handler)
         end
 
       end
 
+      def handle_error(e)
+        error_response({ :message => e.message, :backtrace => e.backtrace })
+      end
+      
       def error_response(error = {})
         status = error[:status] || options[:default_status]
         message = error[:message] || options[:default_message]
         headers = {'Content-Type' => content_type}
         headers.merge!(error[:headers]) if error[:headers].is_a?(Hash)
         backtrace = error[:backtrace] || []
-        Rack::Response.new([format_message(message, backtrace, status)], status, headers).finish
+        rack_response(format_message(message, backtrace, status), status, headers)
       end
-      
+
+      def rack_response(message, status = options[:default_status], headers = { 'Content-Type' => content_type })
+        Rack::Response.new([ message ], status, headers).finish
+      end
+            
       def format_message(message, backtrace, status)
         formatter = formatter_for(options[:format])
         throw :error, :status => 406, :message => "The requested format #{options[:format]} is not supported." unless formatter        

spec/grape/api_spec.rb

@@ -617,4 +617,54 @@ def two
       last_response.status.should eql 403
     end
   end
+  
+  describe ".rescue_from klass, block" do
+    it 'should rescue Exception' do
+      subject.rescue_from RuntimeError do |e|
+        rack_response({ :message => "rescued from #{e.message}" }, 202)
+      end
+      subject.get '/exception' do
+        raise "rain!"
+      end
+      get '/exception'
+      last_response.status.should eql 202
+      last_response.body.should == '{:message=>"rescued from rain!"}'
+    end
+    it 'should rescue an error via rescue_from :all' do
+      class ConnectionError < RuntimeError; end
+      subject.rescue_from :all do |e|
+        rack_response({ :message => "rescued from #{e.class.name}" }, 500)
+      end
+      subject.get '/exception' do
+        raise ConnectionError
+      end
+      get '/exception'
+      last_response.status.should eql 500
+      last_response.body.should == '{:message=>"rescued from ConnectionError"}'
+    end
+    it 'should rescue a specific error' do
+      class ConnectionError < RuntimeError; end
+      subject.rescue_from ConnectionError do |e|
+        rack_response({ :message => "rescued from #{e.class.name}" }, 500)
+      end
+      subject.get '/exception' do
+        raise ConnectionError
+      end
+      get '/exception'
+      last_response.status.should eql 500
+      last_response.body.should == '{:message=>"rescued from ConnectionError"}'
+    end
+    it 'should not rescue a different error' do
+      class CommunicationError < RuntimeError; end
+      subject.rescue_from RuntimeError do |e|
+        rack_response({ :message => "rescued from #{e.class.name}" }, 500)
+      end
+      subject.get '/uncaught' do
+        raise CommunicationError
+      end
+      lambda { get '/uncaught' }.should raise_error(CommunicationError)
+    end
+  end
+  
+  
 end

spec/grape/middleware/exception_spec.rb

@@ -116,5 +116,6 @@ def app
       get '/'
       last_response.status.should == 401    
     end 
+    
   end 
 end