#434

backend/app/controllers/counts_exports_controller.rb

@@ -18,7 +18,7 @@ class CountsExportsController < ApplicationController
     render json: { error: 'private_condition' }, status: :forbidden
   rescue CountsTableQuery::TooManyColumns
     render json: { error: 'too_many_columns', max_columns: CountsTableQuery::MAX_COLUMNS },
-        status: :unprocessable_entity
+           status: :unprocessable_entity
   rescue ActiveRecord::StatementTimeout
     render json: { error: 'counts_query_timeout' }, status: :service_unavailable
   rescue ArgumentError, TypeError, KeyError, JSON::ParserError
@@ -86,68 +86,6 @@ class CountsExportsController < ApplicationController
               disposition: 'attachment'
   end
 
-  def confirmed_download
-    owner = session[:counts_export_owner]
-    arguments = { file_token: params[:token], attempt: params[:attempt] }
-    return head :not_found unless arguments[:attempt].present? &&
-        CountsExportJobs.download_attempt(params[:job], owner, **arguments)
-
-    file = CountsExportDownload.find(params[:token], owner)
-    unless file
-      CountsExportJobs.download_attempt(params[:job], owner, **arguments, status: 'unavailable')
-      return head :not_found
-    end
-    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'
-    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
-  private :confirmed_download
-
-  def job_state(cancel:)
-    result = CountsExportJobs.status(params[:token], session[:counts_export_owner], cancel:)
-    return head :not_found unless result
-
-    response.headers['Cache-Control'] = 'private, no-store'
-    render json: result
-  end
-  private :job_state
-
   def index
     locale = resolve_locale!
     payload = request.request_parameters.except('authenticity_token')
@@ -175,4 +113,73 @@ class CountsExportsController < ApplicationController
   rescue ArgumentError, TypeError, KeyError, JSON::ParserError
     head :bad_request
   end
+
+  private
+
+  def confirmed_download
+    owner = session[:counts_export_owner]
+    arguments = { file_token: params[:token], attempt: params[:attempt] }
+    if (!(arguments[:attempt].present?) ||
+        !(CountsExportJobs.download_attempt(params[:job], owner, **arguments)))
+      return head :not_found
+    end
+
+    file = CountsExportDownload.find(params[:token], owner)
+    unless file
+      CountsExportJobs.download_attempt(params[:job], owner, **arguments, status: 'unavailable')
+      return head :not_found
+    end
+    job = params[:job]
+    authority = HistoryAccess.user
+    first_check = true
+    state = lambda 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')
+      return head :conflict
+    end
+    response.headers['Cache-Control'] = 'private, no-store'
+    response.headers['Last-Modified'] = File.mtime(file.fetch('path')).httpdate
+    response.headers['X-Accel-Buffering'] = 'no'
+    send_file file.fetch('path'),
+        filename: "counts-#{ file.fetch('resource') }.#{ file.fetch('format') }",
+        type: file['format'] == 'json' ? 'application/json' : 'text/csv',
+        disposition: 'attachment'
+    # Keep Rails' bounded FileBody reader, but hide to_path from Rack::Sendfile:
+    # offloading would bypass the GET acknowledgement and cancellation checks.
+    stream = CountsExportDownload::ResponseBody.new(response.stream, &state)
+    self.response_body = stream
+  rescue IOError, SystemCallError
+    stream&.close
+    CountsExportJobs.download_attempt(params[:job], owner, **arguments, status: 'failed')
+    raise if response.committed?
+
+    head :service_unavailable
+  end
+
+  def job_state cancel:
+    result = CountsExportJobs.status(params[:token], session[:counts_export_owner], cancel:)
+    return head :not_found unless result
+
+    response.headers['Cache-Control'] = 'private, no-store'
+    render json: result
+  end
 end

backend/app/controllers/counts_tables_controller.rb

@@ -25,6 +25,11 @@ class CountsTablesController < ApplicationController
   rescue CountsTableQuery::TooManyColumns
     render json: { error: 'too_many_columns', max_columns: CountsTableQuery::MAX_COLUMNS },
            status: :unprocessable_entity
+  rescue CountsPivotTotals::TimedOut
+    Rails.logger.warn('counts_query_timeout: pivot_totals')
+    render json: { error: 'counts_query_timeout', stage: 'totals',
+                   timeout_seconds: CountsTableQuery::SQL_TIMEOUT_SECONDS },
+           status: :service_unavailable
   rescue ActiveRecord::StatementTimeout => error
     Rails.logger.warn("counts_query_timeout: #{ error.class }")
     render json: { error: 'counts_query_timeout' }, status: :service_unavailable