#434

backend/app/services/counts_export_download.rb

@@ -13,15 +13,17 @@ class CountsExportDownload
   class TimedOut < StandardError; end
   class StorageFull < StandardError; end
 
-  def self.directory = Rails.root.join('tmp', 'counts_exports')
+  def self.directory
+    Rails.root.join(ENV.fetch('COUNTS_EXPORT_DIRECTORY', 'tmp/counts_exports'))
+  end
   def self.verifier = Rails.application.message_verifier('counts_export_download')
 
-  def self.prepare(exporter, owner)
+  def self.prepare(exporter, owner, progress: nil, check: nil, identifier: nil)
     FileUtils.mkdir_p(directory, mode: 0o700)
     slot = acquire_slot
 
     begin
-      generate(exporter, owner)
+      generate(exporter, owner, progress:, check:, identifier:)
     ensure
       slot.close
     end
@@ -38,7 +40,7 @@ class CountsExportDownload
     raise Busy
   end
 
-  def self.generate(exporter, owner)
+  def self.generate(exporter, owner, progress: nil, check: nil, identifier: nil)
     FileUtils.mkdir_p(directory, mode: 0o700)
     Dir.glob(directory.join('counts-*.export')).each do |path|
       File.unlink(path) if File.mtime(path) < Time.current - EXPIRY
@@ -53,14 +55,16 @@ class CountsExportDownload
     reserved = MAX_CONCURRENT_EXPORTS * MAX_FILE_BYTES
     raise StorageFull if stored + reserved > MAX_STORED_BYTES
 
-    file = Tempfile.create(['counts-', '.part'], directory)
+    prefix = identifier ? "counts-#{ identifier }-" : 'counts-'
+    file = Tempfile.create([prefix, '.part'], directory)
     complete = false
     bytes = 0
     deadline = Process.clock_gettime(Process::CLOCK_MONOTONIC) + MAX_EXPORT_SECONDS
     path = file.path
     begin
       file.binmode
-      exporter.body.each do |chunk|
+      exporter.body(progress:).each do |chunk|
+        check&.call
         raise TimedOut if Process.clock_gettime(Process::CLOCK_MONOTONIC) > deadline
 
         bytes += chunk.bytesize
@@ -68,6 +72,8 @@ class CountsExportDownload
 
         file.write(chunk)
       end
+      file.flush
+      file.fsync
       file.close
       finished_path = path.sub(/\.part\z/, '.export')
       File.rename(path, finished_path)