s3_data_packer 0.2.0

data/lib/s3_data_packer/configuration.rb

@@ -0,0 +1,99 @@
+module S3DataPacker
+  class Configuration
+    # Standard logger to output information
+    attr_accessor :logger
+
+    # How many threads to run for reading and processing items. This needs to be
+    # balanced out with the speed at which item keys are gathered to prevent
+    # emptying the queue too early.
+    attr_accessor :thread_count
+
+    # Time in seconds to let a thread sleep when there's no pending items in queue.
+    attr_accessor :thread_sleep_time
+
+    # Time in seconds for thread to wait when locked.
+    attr_accessor :thread_lock_wait_time
+
+    # Maximum number of items to maintain in queue to not overflow while workers
+    # process items.
+    attr_accessor :max_queue_size
+
+    # Time in seconds to wait when the queue reached max_queue_size to keep adding
+    # new items.
+    attr_accessor :max_queue_wait
+
+    # Directory to keep working files. Make sure you have permissions on the path
+    # set. If the path does not exist, the packer will try to create it before using
+    # it.
+    attr_accessor :workdir
+
+    # Whether to keep or delete the finalized batch file. Set to false if you want to
+    # keep the output files in the workdir.
+    attr_accessor :cleanup_batch
+
+    # Whether to compress the final batch file or not. If set to true, the output file
+    # will be compressed with GZip, and the uncompressed file will be removed.
+    attr_accessor :compress_batch
+
+    # Number of items of the final batch size
+    attr_accessor :batch_size
+
+    attr_accessor :s3_api_key
+    attr_accessor :s3_api_secret
+    attr_accessor :s3_region
+
+    # String prefix to include in output filenames for the batches
+    attr_accessor :output_filename_prefix
+
+    # String suffix to include in output filenames for the batches
+    attr_accessor :output_filename_suffix
+
+    # Desired pattern to construct output filenames
+    attr_accessor :output_filename_pattern
+
+    # Splitter character to concat different values to generate a filename
+    attr_accessor :output_filename_splitter
+
+    def initialize
+      @thread_count = 2
+      @thread_sleep_time = 1
+      @thread_lock_wait_time = 1
+      @max_queue_size = 10000
+      @max_queue_wait = 5
+      @batch_size = 100000
+      @workdir = 'tmp/s3_data_packer'
+      @cleanup_batch = true
+      @compress_batch = true
+      @output_filename_prefix = nil
+      @output_filename_suffix = 'batch'
+      @output_filename_pattern = %i[timecode_int suffix]
+      @output_filename_splitter = '_'
+    end
+
+    def compress_batch?
+      compress_batch == true
+    end
+
+    def cleanup_batch?
+      cleanup_batch == true
+    end
+
+    def s3_credentials?
+      s3_api_key && s3_api_secret
+    end
+
+    def default_s3_credentials
+      return nil unless s3_credentials?
+
+      Aws::Credentials.new(s3_api_key, s3_api_secret)
+    end
+
+    def filename_generator_defaults
+      { prefix: output_filename_prefix,
+        suffix: output_filename_suffix,
+        pattern: output_filename_pattern,
+        splitter: output_filename_splitter }
+    end
+
+  end
+end

data/lib/s3_data_packer/filename_generator.rb

@@ -0,0 +1,45 @@
+module S3DataPacker
+ class FilenameGenerator
+   attr_reader :pattern, :number, :splitter
+
+   def initialize opts={}
+     @number = opts[:start_at] || 1
+     @prefix = opts[:prefix] || default_options[:prefix]
+     @suffix = opts[:suffix] || default_options[:suffix]
+     @pattern = opts[:pattern] || default_options[:pattern]
+     @splitter = opts[:splitter] || default_options[:splitter]
+     validate_pattern!
+   end
+
+   def timecode_int; Time.now.to_i; end
+   def timecode_dec; Time.now.to_f; end
+   def number; @number; end
+   def timestamp; Time.now.strftime('%Y%m%d%H%M%s'); end
+   def datestamp; Time.now.strftime('%Y%m%d'); end
+   def prefix; @prefix; end
+   def suffix; @suffix; end
+
+   def generate!
+     name = pattern.map{ |key| send(key) }
+     name.delete_if { |value| value.nil? || value == '' }
+     name = name.map(&:to_s).join(splitter)
+     @number += 1
+     name
+   end
+
+   private
+
+   def default_options
+     @default_options ||= S3DataPacker.config.filename_generator_defaults
+   end
+
+   def validate_pattern!
+     valid = %i[timecode_int timecode_dec number timestamp datestamp prefix suffix]
+     pattern.each do |item|
+       raise ArgumentError, "Invalid pattern key, has to be a symbol" unless Symbol === item
+       raise ArgumentError, "Invalid pattern key #{item}. Allowed: #{valid.join(', ')}" unless valid.include?(item)
+     end
+   end
+     
+ end
+end

data/lib/s3_data_packer/json_batch.rb

@@ -0,0 +1,93 @@
+module S3DataPacker
+  class JSONBatch
+    attr_reader :delimitter, :batch, :item_count
+
+    def initialize opts = {}
+      @delimitter = "\r\n"
+      @workdir = opts[:workdir]
+      @filename_generator = opts[:filename_generator]
+      @pre_processor = opts[:pre_processor] # Should be a proc
+      @size = opts[:size]
+      @item_count = 0
+      init_workdir!
+    end
+
+    def size
+      @size ||= S3DataPacker.config.batch_size
+    end
+
+    def workdir
+      @workdir ||= S3DataPacker.config.workdir
+    end
+
+    def filename_generator
+      @filename_generator ||= S3DataPacker::FilenameGenerator.new
+    end
+
+    def full?
+      item_count >= size
+    end
+
+    def generate_filename
+      name = filename_generator.generate!
+      "#{workdir}/#{name}.json"
+    end
+
+    def new_file!
+      close! if @batch
+      @batch = File.open(generate_filename, 'w')
+    end
+
+    def append_data! data
+      digested = pre_proccess_data(data)
+      batch << "#{digested}#{delimitter}"
+      @item_count += 1
+    end
+
+    def path
+      batch.path
+    end
+
+    def close!
+      batch.close
+    end
+
+    def delete!
+      close! if !@batch.closed?
+      File.delete(path) if File.exist?(path)
+      reset!
+    end
+
+    def finalize!
+      close! if !batch.closed?
+      final_path = batch.path
+      final_path = compress! if S3DataPacker.config.compress_batch?
+      reset!
+      final_path
+    end
+
+    private
+
+    def pre_proccess_data(raw_data)
+      # Transformations here, return string for this one
+      return @pre_processor.call(raw_data) if @pre_processor
+      raw_data
+    end
+
+    def init_workdir!
+      Dir.mkdir(workdir) unless Dir.exist?(workdir)
+    end
+
+    def compress!
+      new_path = "#{batch.path}.gz"
+      `gzip #{batch.path}`
+      new_path
+    end
+
+    def reset!
+      @batch = nil
+      @item_count = 0
+    end
+
+  end
+end

data/lib/s3_data_packer/packer.rb

@@ -0,0 +1,105 @@
+module S3DataPacker
+  class Packer
+    module Error
+      class DeadWorkers < StandardError; end
+    end
+
+    attr_reader :source, :target, :output
+
+    def initialize opts = {}
+      @source = opts[:source]
+      @target = opts[:target]
+      @output = opts[:output] || S3DataPacker::JSONBatch.new
+    end
+
+    def summary
+      @summary ||= S3DataPacker::Summary.new
+    end
+
+    def logger
+      @logger ||= S3DataPacker.logger
+    end
+
+    def workers
+      @workers ||= S3DataPacker::ThreadSet.new
+    end
+
+    def pack!
+      log "Packing data from #{source.name} to #{target.name} ..."
+      boot_workers!
+
+      @start_time = Time.now
+      begin
+        each_item { |item| workers.queue.add!(item) }
+        finalize_processing!
+      rescue Exception => e
+        log "Unexpected error, killing threads", :error
+        raise e
+      ensure
+        workers.kill!
+      end
+    end
+
+    def process_item(data)
+      output.append_data! data
+      summary.count_processed
+      if output.full?
+        flush_batch!
+        output.new_file!
+      end
+    end
+
+    private
+
+    def finalize_processing!
+      log "No more items found to enqueue, signaling workers to finish"
+      workers.finish!
+      workers.wait!
+      workers.kill!
+      log "Pushing last open batch #{output.path}"
+      flush_batch!
+      summary.set_time(@start_time, Time.now)
+      log "Finished\n#{summary.flush!}"
+    end
+
+    def each_item &block
+      source.each do |item|
+        if workers.dead?
+          log "Workers diead", :error
+          raise Error::DeadWorkers, 'Workers died'
+        end
+        summary.count_item
+        yield item
+      end
+    end
+
+    def flush_batch!
+      summary.count_batch
+      final_filename = output.finalize!
+      send_file!(final_filename)
+    end
+
+    def send_file!(file)
+      target.save_file file
+    end
+
+    def boot_workers!
+      output.new_file!
+      workers.spawn_threads! do |item|
+        data = source.fetch(item)
+        workers.lock.synchronize { process_item(data) }
+        post_process_item(item)
+      end
+    end
+
+    def post_process_item(item)
+      # Do something with the key after processed
+      nil
+    end
+
+    def log(message, level = :info)
+      logger.send level, "Main process: #{message}"
+    end
+
+  end
+end

data/lib/s3_data_packer/queue.rb

@@ -0,0 +1,46 @@
+module S3DataPacker
+  class Queue
+    attr_reader :total_items, :items
+
+    def initialize opts = {}
+      @max_items = opts[:max_items]
+      @wait_time = opts[:wait_time]
+      @total_items = 0
+    end
+
+    def max_items
+      @max_items ||= S3DataPacker.config.max_queue_size
+    end
+
+    def wait_time
+      @wait_time ||= S3DataPacker.config.max_queue_wait
+    end
+
+    def items
+      @items ||= []
+    end
+
+    def add!(item)
+      items << item
+      @total_items += 1
+      if size >= max_items
+        S3DataPacker.logger.info "Queue full, pausing"
+        sleep(wait_time)
+        S3DataPacker.logger.info "Resuming queue"
+      end
+    end
+
+    def fetch!
+      items.shift
+    end
+
+    def size
+      items.size
+    end
+
+    def reset!
+      @items = []
+      @total_items = 0
+    end
+  end
+end

data/lib/s3_data_packer/sources/object.rb

@@ -0,0 +1,28 @@
+module S3DataPacker
+  module Sources
+    class Object
+
+      def initialize(object:, each_method: :each, fetch_method: :fetch, name_method: :name)
+        @object = object
+        @each_method = each_method
+        @fetch_method = fetch_method
+        @name_method = name_method
+      end
+
+      def name
+        @object.send(@name_method)
+      end
+
+      def each &block
+        @object.send(@each_method) do |item|
+          yield item
+        end
+      end
+
+      def fetch(item)
+        @object.send(@fetch_method, item)
+      end
+
+    end
+  end
+end

data/lib/s3_data_packer/sources/s3_bucket.rb

@@ -0,0 +1,21 @@
+module S3DataPacker
+  module Sources
+    
+    class S3Bucket < S3DataPacker::Bucket
+      def name
+        "s3://#{bucket_name}/#{path}"
+      end
+
+      def each(&block)
+        each_key do |key|
+          yield key
+        end
+      end
+
+      def fetch(key)
+        download(key)
+      end
+    end
+
+  end
+end

data/lib/s3_data_packer/summary.rb

@@ -0,0 +1,59 @@
+module S3DataPacker
+  class Summary
+    def stats
+      @stats ||= {}
+    end
+
+    def count_item
+      stats[:total_items] ||= 0
+      stats[:total_items] += 1
+    end
+
+    def count_processed
+      stats[:processed] ||= 0
+      stats[:processed] += 1
+    end
+
+    def count_batch
+      stats[:batches] ||= 0
+      stats[:batches] += 1
+    end
+
+    def set_time start_time, end_time
+      stats[:elapsed] = "#{(end_time.to_i - start_time.to_i)} seconds"
+    end
+
+    def total_items
+      stats[:total_items] || 0
+    end
+
+    def processed
+      stats[:processed] || 0
+    end
+
+    def batches
+      stats[:batches] || 0
+    end
+
+    def elapsed
+      stats[:elapsed]
+    end
+
+    def flush!
+      output = [
+        "Summary:",
+        "Total Items: #{stats[:total_items]}",
+        "Processed Items: #{stats[:processed]}",
+        "Batches: #{stats[:batches]}",
+        "Elapsed: #{stats[:elapsed]}"
+      ].join("\n")
+      reset!
+      output
+    end
+
+    def reset!
+      @stats = {}
+    end
+
+  end
+end

data/lib/s3_data_packer/targets/object.rb

@@ -0,0 +1,21 @@
+module S3DataPacker
+  module Targets
+    class Object
+
+      def initialize(object:, name_method: :each, save_file_method: :save_file)
+        @object = object
+        @name_method = name_method
+        @save_file_method = save_file_method
+      end
+
+      def name
+        @object.send(@name_method)
+      end
+
+      def save_file(filepath)
+        @object.send(@save_file_method, filepath)
+      end
+
+    end
+  end
+end

data/lib/s3_data_packer/targets/s3_bucket.rb

@@ -0,0 +1,16 @@
+module S3DataPacker
+  module Targets
+
+    class S3Bucket < S3DataPacker::Bucket
+      def name
+        "s3://#{bucket_name}/#{path}"
+      end
+
+      def save_file(filepath)
+        upload(filepath)
+        File.delete(filepath) if S3DataPacker.config.cleanup_batch?
+      end
+    end
+
+  end
+end

data/lib/s3_data_packer/thread_set.rb

@@ -0,0 +1,98 @@
+module S3DataPacker
+  class ThreadSet
+    attr_reader :lock, :workers, :queue
+
+    def initialize opts ={}
+      @lock = Mutex.new
+      @workers = []
+      @finish = false
+      @queue = S3DataPacker::Queue.new
+    end
+
+    def wait_time
+      @wait_time ||= S3DataPacker.config.thread_sleep_time
+    end
+
+    def lock_wait_time
+      @lock_wait_time ||= S3DataPacker.config.thread_lock_wait_time
+    end
+
+    def thread_count
+      @thread_count ||= S3DataPacker.config.thread_count
+    end
+
+    def dead?
+      workers.map(&:status).uniq == [nil] || workers.map(&:status).uniq == [false]
+    end
+
+    def kill!
+      log 'All', "Killing #{workers.size} workers"
+      workers.map(&:kill)
+    end
+
+    def reset!
+      return unless dead?
+      @finish = false
+      @workers = []
+    end
+
+    def finish!
+      @finish = true
+    end
+
+    def finished?
+      @finish == true && queue.size == 0
+    end
+
+    def log id, message, level = :info
+      logger.send level, "Thread #{id}: #{message}"
+    end
+
+    def wait!
+      workers.map(&:join)
+    end
+
+    def spawn_thread! id, &block
+      @workers << Thread.new do
+        log id, "Started!"
+        loop do
+          if finished?
+            log id, "Finish signal up and no more work to pull - Exiting"
+            break
+          end
+          item = queue.fetch!
+          if item
+            log id, "Processing item #{item}", :debug
+            begin
+              yield item
+            rescue ThreadError
+              log id, "Locked, retry in #{lock_wait_time}", :warn
+              sleep(lock_wait_time)
+              retry
+            end
+          else
+            log id, "No more work found, sleeping for #{wait_time}"
+            sleep(wait_time)
+          end
+        rescue Exception => e
+          log id, 'Unexpected error!'
+          raise e
+        end
+      end
+    end
+
+    def spawn_threads! &block
+      logger.info "Spawning #{thread_count} threads"
+      thread_count.times do |id|
+        spawn_thread!(id, &block)
+      end
+    end
+
+    private
+
+    def logger
+      @logger ||= S3DataPacker.logger
+    end
+
+  end
+end

data/lib/s3_data_packer/version.rb

@@ -0,0 +1,3 @@
+module S3DataPacker
+  VERSION = "0.2.0"
+end

data/lib/s3_data_packer.rb

@@ -0,0 +1,41 @@
+require 'csv'
+require 'json'
+require 'logger'
+require 'mime/types/full'
+require 'aws-sdk-s3'
+
+require "s3_data_packer/version"
+require 's3_data_packer/configuration'
+require 's3_data_packer/packer'
+require 's3_data_packer/queue'
+require 's3_data_packer/thread_set'
+require 's3_data_packer/summary'
+require 's3_data_packer/json_batch'
+require 's3_data_packer/bucket'
+require 's3_data_packer/filename_generator'
+
+require 's3_data_packer/sources/object'
+require 's3_data_packer/sources/s3_bucket'
+
+require 's3_data_packer/targets/s3_bucket'
+require 's3_data_packer/targets/object'
+
+module S3DataPacker
+  class << self
+    attr_reader :configuration
+
+    def configuration
+      @configuration ||= Configuration.new
+    end
+
+    alias config configuration
+
+    def configure
+      yield configuration
+    end
+
+    def logger
+      @logger ||= config.logger || Logger.new('log/s3_data_packer.log')
+    end
+  end
+end