Merge pull request #1079 from zbelzer/support_streaming

Add `stream` method to InsideRoute to leverage Rack::Chunked

Author: dblock

.rubocop_todo.yml

@@ -35,7 +35,7 @@ Metrics/MethodLength:
 # Offense count: 8
 # Configuration parameters: CountComments.
 Metrics/ModuleLength:
-  Max: 243
+  Max: 271
 
 # Offense count: 17
 Metrics/PerceivedComplexity:

CHANGELOG.md

@@ -7,6 +7,7 @@ Next Release
 * [#1039](https://github.com/intridea/grape/pull/1039): Added support for custom parameter types - [@rnubel](https://github.com/rnubel).
 * [#1047](https://github.com/intridea/grape/pull/1047): Adds `given` to DSL::Parameters, allowing for dependent params - [@rnubel](https://github.com/rnubel).
 * [#1064](https://github.com/intridea/grape/pull/1064): Add public `Grape::Exception::ValidationErrors#full_messages` - [@romanlehnert](https://github.com/romanlehnert).
+* [#1079](https://github.com/intridea/grape/pull/1079): Added `stream` method to take advantage of `Rack::Chunked` [@zbelzer](https://github.com/zbelzer).
 * Your contribution here!
 
 #### Fixes

README.md

@@ -2105,6 +2105,8 @@ end
 Use `body false` to return `204 No Content` without any data or content-type.
 
 You can also set the response to a file-like object with `file`.
+Note: Rack will read your entire Enumerable before returning a response. If
+you would like to stream the response, see `stream`.
 
 ```ruby
 class FileStreamer
@@ -2126,6 +2128,16 @@ class API < Grape::API
 end
 ```
 
+If you want a file-like object to be streamed using Rack::Chunked, use `stream`.
+
+```ruby
+class API < Grape::API
+  get '/' do
+    stream FileStreamer.new('file.bin')
+  end
+end
+```
+
 ## Authentication
 
 ### Basic and Digest Auth

lib/grape.rb

@@ -122,6 +122,7 @@ module Util
     autoload :StackableValues
     autoload :InheritableSetting
     autoload :StrictHashConfiguration
+    autoload :FileResponse
   end
 
   module DSL

lib/grape/dsl/inside_route.rb

@@ -177,12 +177,34 @@ def body(value = nil)
       #   GET /file # => "contents of file"
       def file(value = nil)
         if value
-          @file = value
+          @file = Grape::Util::FileResponse.new(value)
         else
           @file
         end
       end
 
+      # Allows you to define the response as a streamable object.
+      #
+      # If Content-Length and Transfer-Encoding are blank (among other conditions),
+      # Rack assumes this response can be streamed in chunks.
+      #
+      # @example
+      #   get '/stream' do
+      #     stream FileStreamer.new(...)
+      #   end
+      #
+      #   GET /stream # => "chunked contents of file"
+      #
+      # See:
+      # * https://github.com/rack/rack/blob/99293fa13d86cd48021630fcc4bd5acc9de5bdc3/lib/rack/chunked.rb
+      # * https://github.com/rack/rack/blob/99293fa13d86cd48021630fcc4bd5acc9de5bdc3/lib/rack/etag.rb
+      def stream(value = nil)
+        header 'Content-Length', nil
+        header 'Transfer-Encoding', nil
+        header 'Cache-Control', 'no-cache' # Skips ETag generation (reading the response up front)
+        file(value)
+      end
+
       # Allows you to make use of Grape Entities by setting
       # the response body to the serializable hash of the
       # entity provided in the `:with` option. This has the

lib/grape/middleware/formatter.rb

@@ -18,26 +18,49 @@ def before
 
       def after
         status, headers, bodies = *@app_response
-        # allow content-type to be explicitly overwritten
-        api_format = mime_types[headers[Grape::Http::Headers::CONTENT_TYPE]] || env['api.format']
-        formatter = Grape::Formatter::Base.formatter_for api_format, options
-        begin
-          bodymap = if bodies.respond_to?(:collect)
-                      bodies.collect do |body|
-                        formatter.call body, env
-                      end
-                    else
-                      bodies
-                    end
-        rescue Grape::Exceptions::InvalidFormatter => e
-          throw :error, status: 500, message: e.message
+
+        if bodies.is_a?(Grape::Util::FileResponse)
+          headers = ensure_content_type(headers)
+
+          response =
+            Rack::Response.new([], status, headers) do |resp|
+              resp.body = bodies.file
+            end
+        else
+          # Allow content-type to be explicitly overwritten
+          api_format = mime_types[headers[Grape::Http::Headers::CONTENT_TYPE]] || env['api.format']
+          formatter = Grape::Formatter::Base.formatter_for(api_format, options)
+
+          begin
+            bodymap = bodies.collect do |body|
+              formatter.call(body, env)
+            end
+
+            headers = ensure_content_type(headers)
+
+            response = Rack::Response.new(bodymap, status, headers)
+          rescue Grape::Exceptions::InvalidFormatter => e
+            throw :error, status: 500, message: e.message
+          end
         end
-        headers[Grape::Http::Headers::CONTENT_TYPE] = content_type_for(env['api.format']) unless headers[Grape::Http::Headers::CONTENT_TYPE]
-        Rack::Response.new(bodymap, status, headers)
+
+        response
       end
 
       private
 
+      # Set the content type header for the API format if it is not already present.
+      #
+      # @param headers [Hash]
+      # @return [Hash]
+      def ensure_content_type(headers)
+        if headers[Grape::Http::Headers::CONTENT_TYPE]
+          headers
+        else
+          headers.merge(Grape::Http::Headers::CONTENT_TYPE => content_type_for(env['api.format']))
+        end
+      end
+
       def request
         @request ||= Rack::Request.new(env)
       end

lib/grape/util/file_response.rb

@@ -0,0 +1,21 @@
+module Grape
+  module Util
+    # A simple class used to identify responses which represent files and do not
+    # need to be formatted or pre-read by Rack::Response
+    class FileResponse
+      attr_reader :file
+
+      # @param file [Object]
+      def initialize(file)
+        @file = file
+      end
+
+      # Equality provided mostly for tests.
+      #
+      # @return [Boolean]
+      def ==(other)
+        file == other.file
+      end
+    end
+  end
+end

spec/grape/api_spec.rb

@@ -795,6 +795,37 @@ def subject.enable_root_route!
       expect(last_response.body).to eq(file)
     end
 
+    it 'returns the content of the file with file' do
+      file_content = 'This is some file content'
+      test_file = Tempfile.new('test')
+      test_file.write file_content
+      test_file.rewind
+
+      subject.get('/file') { file test_file }
+      get '/file'
+      expect(last_response.headers['Content-Length']).to eq('25')
+      expect(last_response.headers['Content-Type']).to eq('text/plain')
+      expect(last_response.body).to eq(file_content)
+    end
+
+    it 'streams the content of the file with stream' do
+      test_stream = Enumerator.new do |blk|
+        blk.yield 'This is some'
+        blk.yield ' file content'
+      end
+
+      subject.use Rack::Chunked
+      subject.get('/stream') { stream test_stream }
+      get '/stream', {}, 'HTTP_VERSION' => 'HTTP/1.1'
+
+      expect(last_response.headers['Content-Type']).to eq('text/plain')
+      expect(last_response.headers['Content-Length']).to eq(nil)
+      expect(last_response.headers['Cache-Control']).to eq('no-cache')
+      expect(last_response.headers['Transfer-Encoding']).to eq('chunked')
+
+      expect(last_response.body).to eq("c\r\nThis is some\r\nd\r\n file content\r\n0\r\n\r\n")
+    end
+
     it 'sets content type for error' do
       subject.get('/error') { error!('error in plain text', 500) }
       get '/error'

spec/grape/dsl/inside_route_spec.rb

@@ -201,8 +201,43 @@ def initialize
             subject.file 'file'
           end
 
-          it 'returns value' do
-            expect(subject.file).to eq 'file'
+          it 'returns value wrapped in FileResponse' do
+            expect(subject.file).to eq Grape::Util::FileResponse.new('file')
+          end
+        end
+
+        it 'returns default' do
+          expect(subject.file).to be nil
+        end
+      end
+
+      describe '#stream' do
+        describe 'set' do
+          before do
+            subject.header 'Cache-Control', 'cache'
+            subject.header 'Content-Length', 123
+            subject.header 'Transfer-Encoding', 'base64'
+            subject.stream 'file'
+          end
+
+          it 'returns value wrapped in FileResponse' do
+            expect(subject.stream).to eq Grape::Util::FileResponse.new('file')
+          end
+
+          it 'also sets result of file to value wrapped in FileResponse' do
+            expect(subject.file).to eq Grape::Util::FileResponse.new('file')
+          end
+
+          it 'sets Cache-Control header to no-cache' do
+            expect(subject.header['Cache-Control']).to eq 'no-cache'
+          end
+
+          it 'sets Content-Length header to nil' do
+            expect(subject.header['Content-Length']).to eq nil
+          end
+
+          it 'sets Transfer-Encoding header to nil' do
+            expect(subject.header['Transfer-Encoding']).to eq nil
           end
         end