Add batch command for BYO snapshots (#57)

Author: trevor-e

lib/commands/upload/snapshots/snapshots.rb

@@ -6,6 +6,8 @@
 require 'async/barrier'
 require 'async/semaphore'
 require 'async/http/internet/instance'
+require 'zip'
+require 'tempfile'
 
 module EmergeCLI
   module Commands
@@ -33,6 +35,8 @@ class Snapshots < EmergeCLI::Commands::GlobalOptions
 
         option :profile, type: :boolean, default: false, desc: 'Enable performance profiling metrics'
 
+        option :batch, type: :boolean, default: false, desc: 'Upload images in batch using zip file'
+
         argument :image_paths, type: :array, required: false, desc: 'Paths to folders containing images'
 
         def initialize(network: nil, git_info_provider: nil)
@@ -178,6 +182,86 @@ def create_run
         def upload_images(run_id, concurrency, image_files, client)
           Logger.info 'Uploading images...'
 
+          if @options[:batch]
+            batch_upload_images(run_id, image_files, client)
+          else
+            individual_upload_images(run_id, concurrency, image_files, client)
+          end
+        end
+
+        def batch_upload_images(run_id, image_files, client)
+          Logger.info 'Preparing batch upload...'
+
+          metadata_barrier = Async::Barrier.new
+          metadata_semaphore = Async::Semaphore.new(10, parent: metadata_barrier)
+
+          image_metadata = {
+            manifestVersion: 1,
+            images: {},
+            errors: []
+          }
+
+          @profiler.measure('process_image_metadata') do
+            image_files.each do |image_path|
+              metadata_semaphore.async do
+                file_info = client.parse_file_info(image_path)
+
+                dimensions = @profiler.measure('chunky_png_processing') do
+                  datastream = ChunkyPNG::Datastream.from_file(image_path)
+                  {
+                    width: datastream.header_chunk.width,
+                    height: datastream.header_chunk.height
+                  }
+                end
+
+                metadata = {
+                  fileName: file_info[:file_name],
+                  groupName: file_info[:group_name],
+                  displayName: file_info[:variant_name],
+                  width: dimensions[:width],
+                  height: dimensions[:height]
+                }
+
+                image_name = File.basename(image_path, '.*')
+                image_metadata[:images][image_name] = metadata
+              end
+            end
+
+            metadata_barrier.wait
+          end
+
+          Tempfile.create(['snapshot_batch', '.zip']) do |zip_file|
+            @profiler.measure('create_zip_file') do
+              Zip::File.open(zip_file.path, Zip::File::CREATE) do |zipfile|
+                zipfile.get_output_stream('manifest.json') { |f| f.write(JSON.generate(image_metadata)) }
+
+                image_files.each do |image_path|
+                  image_name = File.basename(image_path)
+                  zipfile.add(image_name, image_path)
+                end
+              end
+            end
+
+            upload_url = @profiler.measure('create_batch_upload_url') do
+              response = @network.post(path: '/v1/snapshots/run/batch-image', body: { run_id: run_id })
+              JSON.parse(response.read).fetch('zip_url')
+            end
+
+            Logger.info 'Uploading images...'
+            Logger.debug "Uploading batch zip file to #{upload_url}"
+            @profiler.measure('upload_batch_zip') do
+              @network.put(
+                path: upload_url,
+                headers: { 'Content-Type' => 'application/zip' },
+                body: File.read(zip_file.path)
+              )
+            end
+          end
+        ensure
+          metadata_barrier&.stop
+        end
+
+        def individual_upload_images(run_id, concurrency, image_files, client)
           post_image_barrier = Async::Barrier.new
           post_image_semaphore = Async::Semaphore.new(concurrency, parent: post_image_barrier)