From 90049a41075de0864f20431b1065fb0421a4865c Mon Sep 17 00:00:00 2001
From: David Heinemeier Hansson <david@loudthinking.com>
Date: Thu, 18 Feb 2021 22:35:36 +0100
Subject: [PATCH] Add send_stream to do for dynamic streams what send_data does
 for static files (#41488)

---
 actionpack/CHANGELOG.md                       | 14 ++++++++
 .../lib/action_controller/metal/live.rb       | 35 +++++++++++++++++++
 .../test/controller/live_stream_test.rb       | 28 +++++++++++++++
 3 files changed, 77 insertions(+)

diff --git a/actionpack/CHANGELOG.md b/actionpack/CHANGELOG.md
index c950d21f98..961b8c9cde 100644
--- a/actionpack/CHANGELOG.md
+++ b/actionpack/CHANGELOG.md
@@ -1,3 +1,17 @@
+*   Add `ActionController::Live#send_stream` that makes it more convenient to send generated streams:
+
+    ```ruby
+    send_stream(filename: "subscribers.csv") do |stream|
+      stream.write "email_address,updated_at\n"
+    
+      @subscribers.find_each do |subscriber|
+        stream.write "#{subscriber.email_address},#{subscriber.updated_at}\n"
+      end
+    end
+    ```
+    
+    *DHH*
+
 *   `ActionDispatch::Request#content_type` now returned Content-Type header as it is.
 
     Previously, `ActionDispatch::Request#content_type` returned value does NOT contain charset part.
diff --git a/actionpack/lib/action_controller/metal/live.rb b/actionpack/lib/action_controller/metal/live.rb
index de3c944bb5..062cc2ba3f 100644
--- a/actionpack/lib/action_controller/metal/live.rb
+++ b/actionpack/lib/action_controller/metal/live.rb
@@ -282,6 +282,41 @@ module ActionController
       response.close if response
     end
 
+    # Sends a stream to the browser, which is helpful when you're generating exports or other running data where you
+    # don't want the entire file buffered in memory first. Similar to send_data, but where the data is generated live.
+    #
+    # Options:
+    # * <tt>:filename</tt> - suggests a filename for the browser to use.
+    # * <tt>:type</tt> - specifies an HTTP content type.
+    #   You can specify either a string or a symbol for a registered type with <tt>Mime::Type.register</tt>, for example :json.
+    #   If omitted, type will be inferred from the file extension specified in <tt>:filename</tt>.
+    #   If no content type is registered for the extension, the default type 'application/octet-stream' will be used.
+    # * <tt>:disposition</tt> - specifies whether the file will be shown inline or downloaded.
+    #   Valid values are 'inline' and 'attachment' (default).
+    #
+    # Example of generating a csv export:
+    #
+    #    send_stream(filename: "subscribers.csv") do |stream|
+    #      stream.write "email_address,updated_at\n"
+    #
+    #      @subscribers.find_each do |subscriber|
+    #        stream.write "#{subscriber.email_address},#{subscriber.updated_at}\n"
+    #      end
+    #    end
+    def send_stream(filename:, disposition: "attachment", type: nil)
+      response.headers["Content-Type"] =
+        (type.is_a?(Symbol) ? Mime[type].to_s : type) ||
+        Mime::Type.lookup_by_extension(File.extname(filename).downcase.delete(".")) ||
+        "application/octet-stream"
+
+      response.headers["Content-Disposition"] =
+        ActionDispatch::Http::ContentDisposition.format(disposition: disposition, filename: filename)
+
+      yield response.stream
+    ensure
+      response.stream.close
+    end
+
     private
       # Spawn a new thread to serve up the controller in. This is to get
       # around the fact that Rack isn't based around IOs and we need to use
diff --git a/actionpack/test/controller/live_stream_test.rb b/actionpack/test/controller/live_stream_test.rb
index 6e77c64518..48542a0ef8 100644
--- a/actionpack/test/controller/live_stream_test.rb
+++ b/actionpack/test/controller/live_stream_test.rb
@@ -144,6 +144,18 @@ module ActionController
         response.stream.close
       end
 
+      def basic_send_stream
+        send_stream(filename: "my.csv") do |stream|
+          stream.write "name,age\ndavid,41"
+        end
+      end
+
+      def send_stream_with_options
+        send_stream(filename: "export", disposition: "inline", type: :json) do |stream|
+          stream.write %[{ name: "David", age: 41 }]
+        end
+      end
+
       def blocking_stream
         response.headers["Content-Type"] = "text/event-stream"
         %w{ hello world }.each do |word|
@@ -300,6 +312,22 @@ module ActionController
       assert_equal "text/event-stream", @response.headers["Content-Type"]
     end
 
+    def test_send_stream
+      get :basic_send_stream
+      assert_equal "name,age\ndavid,41", @response.body
+      assert_equal "text/csv", @response.headers["Content-Type"]
+      assert_match "attachment", @response.headers["Content-Disposition"]
+      assert_match "my.csv", @response.headers["Content-Disposition"]
+    end
+
+    def test_send_stream_with_optons
+      get :send_stream_with_options
+      assert_equal %[{ name: "David", age: 41 }], @response.body
+      assert_equal "application/json", @response.headers["Content-Type"]
+      assert_match "inline", @response.headers["Content-Disposition"]
+      assert_match "export", @response.headers["Content-Disposition"]
+    end
+
     def test_delayed_autoload_after_write_within_interlock_hook
       # Simulate InterlockHook
       ActiveSupport::Dependencies.interlock.start_running