#434

backend/app/controllers/counts_exports_controller.rb

@@ -12,6 +12,15 @@ class CountsExportsController < ApplicationController
     render json: { error: 'export_worker_unavailable' }, status: :service_unavailable
   rescue CountsExportDownload::Busy
     head :too_many_requests
+  rescue CountsAxes::TooManyAxes
+    render json: { error: 'too_many_axes', max_axes: CountsAxes::MAX_AXES }, status: :bad_request
+  rescue CountsAxes::PrivateCondition
+    render json: { error: 'private_condition' }, status: :forbidden
+  rescue CountsTableQuery::TooManyColumns
+    render json: { error: 'too_many_columns', max_columns: CountsTableQuery::MAX_COLUMNS },
+        status: :unprocessable_entity
+  rescue ActiveRecord::StatementTimeout
+    render json: { error: 'counts_query_timeout' }, status: :service_unavailable
   rescue ArgumentError, TypeError, KeyError, JSON::ParserError
     head :bad_request
   end
@@ -48,6 +57,15 @@ class CountsExportsController < ApplicationController
     head :gateway_timeout
   rescue CountsExportDownload::StorageFull, Errno::ENOSPC
     head :insufficient_storage
+  rescue CountsAxes::TooManyAxes
+    render json: { error: 'too_many_axes', max_axes: CountsAxes::MAX_AXES }, status: :bad_request
+  rescue CountsAxes::PrivateCondition
+    render json: { error: 'private_condition' }, status: :forbidden
+  rescue CountsTableQuery::TooManyColumns
+    render json: { error: 'too_many_columns', max_columns: CountsTableQuery::MAX_COLUMNS },
+        status: :unprocessable_entity
+  rescue ActiveRecord::StatementTimeout
+    render json: { error: 'counts_query_timeout' }, status: :service_unavailable
   rescue ArgumentError, TypeError, KeyError, JSON::ParserError
     head :bad_request
   end
@@ -79,36 +97,43 @@ class CountsExportsController < ApplicationController
       CountsExportJobs.download_attempt(params[:job], owner, **arguments, status: 'unavailable')
       return head :not_found
     end
-    # Open before committing the response, but never buffer the entire export.
-    stream = File.open(file.fetch('path'), 'rb')
+    job = params[:job]
+    authority = HistoryAccess.user
+    first_check = true
+    stream = CountsExportDownload::ResponseBody.new(file.fetch('path')) do |status|
+      if status == 'active'
+        if first_check
+          first_check = false
+          fresh = authority && User.find_by(id: authority.id)
+          next false if authority && (!fresh || fresh.banned?)
+
+          authority = fresh
+          HistoryAccess.set(user: authority) do
+            CountsExportDownload.find(arguments[:file_token], owner) &&
+                CountsExportJobs.stream_active?(job, owner, **arguments)
+          end
+        else
+          CountsExportJobs.stream_active?(job, owner, **arguments)
+        end
+      else
+        HistoryAccess.set(user: authority) do
+          CountsExportJobs.download_attempt(job, owner, **arguments, status:)
+        end
+      end
+    end
+    unless CountsExportJobs.download_attempt(job, owner, **arguments, status: 'streaming')
+      stream.close
+      return head :conflict
+    end
     response.headers['Cache-Control'] = 'private, no-store'
     send_file_headers! filename: "counts-#{ file.fetch('resource') }.#{ file.fetch('format') }",
         type: file['format'] == 'json' ? 'application/json' : 'text/csv',
         disposition: 'attachment'
-    job = params[:job]
-    authority = HistoryAccess.user
-    self.response_body = Enumerator.new do |output|
-      begin
-        first = true
-        while (chunk = stream.read(64 * 1024))
-          output << chunk
-          if first
-            HistoryAccess.set(user: authority) do
-              CountsExportJobs.download_attempt(job, owner, **arguments, status: 'started')
-            end
-            first = false
-          end
-        end
-      rescue IOError, SystemCallError => error
-        Rails.logger.warn("counts_download_stream_failure: #{ error.class }")
-        HistoryAccess.set(user: authority) do
-          CountsExportJobs.download_attempt(job, owner, **arguments, status: 'failed')
-        end
-      ensure
-        stream.close
-      end
-    end
+    response.headers['Last-Modified'] = File.mtime(file.fetch('path')).httpdate
+    response.headers['X-Accel-Buffering'] = 'no'
+    self.response_body = stream
   rescue IOError, SystemCallError
+    stream&.close
     CountsExportJobs.download_attempt(params[:job], owner, **arguments, status: 'failed')
     head :service_unavailable
   end
@@ -138,6 +163,15 @@ class CountsExportsController < ApplicationController
         filename: "counts-#{ exporter.resource }.#{ exporter.format }",
         type: exporter.format == 'json' ? 'application/json' : 'text/csv',
         disposition: 'attachment'
+  rescue CountsAxes::TooManyAxes
+    render json: { error: 'too_many_axes', max_axes: CountsAxes::MAX_AXES }, status: :bad_request
+  rescue CountsAxes::PrivateCondition
+    render json: { error: 'private_condition' }, status: :forbidden
+  rescue CountsTableQuery::TooManyColumns
+    render json: { error: 'too_many_columns', max_columns: CountsTableQuery::MAX_COLUMNS },
+        status: :unprocessable_entity
+  rescue ActiveRecord::StatementTimeout
+    render json: { error: 'counts_query_timeout' }, status: :service_unavailable
   rescue ArgumentError, TypeError, KeyError, JSON::ParserError
     head :bad_request
   end