bricolage-streamingload 0.1.0

data/lib/bricolage/streamingload/event.rb

@@ -0,0 +1,139 @@
+require 'bricolage/sqsdatasource'
+
+module Bricolage
+
+  module StreamingLoad
+
+    class Event < SQSMessage
+
+      def Event.get_concrete_class(msg, rec)
+        case
+        when rec['eventName'] == 'shutdown' then ShutdownEvent
+        when rec['eventName'] == 'dispatch' then DispatchEvent
+        when rec['eventName'] == 'flush' then FlushEvent
+        when rec['eventSource'] == 'aws:s3'
+          S3ObjectEvent
+        else
+          raise "[FATAL] unknown SQS message record: eventSource=#{rec['eventSource']} event=#{rec['eventName']} message_id=#{msg.message_id}"
+        end
+      end
+
+      def message_type
+        raise "#{self.class}\#message_type must be implemented"
+      end
+
+      def data?
+        false
+      end
+
+    end
+
+
+    class ShutdownEvent < Event
+
+      def ShutdownEvent.create
+        super name: 'shutdown'
+      end
+
+      def ShutdownEvent.parse_sqs_record(msg, rec)
+        {}
+      end
+
+      alias message_type name
+
+      def init_message
+      end
+
+    end
+
+
+    class FlushEvent < Event
+
+      def FlushEvent.create(delay_seconds:, table_name:)
+        super name: 'flush', delay_seconds: delay_seconds, table_name: table_name
+      end
+
+      def FlushEvent.parse_sqs_record(msg, rec)
+        {
+          table_name: rec['tableName']
+        }
+      end
+
+      alias message_type name
+
+      def init_message(table_name:)
+        @table_name = table_name
+      end
+
+      attr_reader :table_name
+
+      def body
+        obj = super
+        obj['tableName'] = @table_name
+        obj
+      end
+
+    end
+
+    class DispatchEvent < Event
+
+      def DispatchEvent.create(delay_seconds:)
+        super name: 'dispatch', delay_seconds: delay_seconds
+      end
+
+      alias message_type name
+
+      def init_message(dummy)
+      end
+    end
+
+
+    class S3ObjectEvent < Event
+
+      def S3ObjectEvent.parse_sqs_record(msg, rec)
+        {
+          region: rec['awsRegion'],
+          bucket: rec['s3']['bucket']['name'],
+          key: rec['s3']['object']['key'],
+          size: rec['s3']['object']['size']
+        }
+      end
+
+      def message_type
+        'data'
+      end
+
+      def init_message(region:, bucket:, key:, size:)
+        @region = region
+        @bucket = bucket
+        @key = key
+        @size = size
+      end
+
+      attr_reader :region
+      attr_reader :bucket
+      attr_reader :key
+      attr_reader :size
+
+      def url
+        "s3://#{@bucket}/#{@key}"
+      end
+
+      # override
+      def data?
+        true
+      end
+
+      def created?
+        !!(/\AObjectCreated:(?!Copy)/ =~ @name)
+      end
+
+      def loadable_object(url_patterns)
+        LoadableObject.new(self, url_patterns.match(url))
+      end
+
+    end
+
+  end
+
+end

data/lib/bricolage/streamingload/loader.rb

@@ -0,0 +1,144 @@
+require 'bricolage/streamingload/loaderparams'
+require 'bricolage/streamingload/manifest'
+require 'bricolage/sqlutils'
+require 'socket'
+require 'json'
+
+module Bricolage
+
+  module StreamingLoad
+
+    class Loader
+
+      include SQLUtils
+
+      def Loader.load_from_file(ctx, ctl_ds, task, logger:)
+        params = LoaderParams.load(ctx, task)
+        new(ctl_ds, params, logger: logger)
+      end
+
+      def initialize(ctl_ds, params, logger:)
+        @ctl_ds = ctl_ds
+        @params = params
+        @logger = logger
+        @process_id = "#{Socket.gethostname}-#{$$}"
+      end
+
+      def execute
+        @job_id = assign_task
+        return unless @job_id # task already executed by other loader
+        @params.ds.open {|conn|
+          @connection = conn
+          do_load
+        }
+      end
+
+      def assign_task
+        @ctl_ds.open {|conn|
+          job_id = conn.query_value(<<-EndSQL)
+            insert into strload_jobs
+                ( task_id
+                , process_id
+                , status
+                , start_time
+                )
+            select
+                task_id
+                , #{s @process_id}
+                , 'running'
+                , current_timestamp
+            from
+                strload_tasks
+            where
+                task_id = #{@params.task_id}
+                and (task_id not in (select task_id from strload_jobs) or #{@params.force})
+            returning job_id
+            ;
+          EndSQL
+          return job_id
+        }
+      end
+
+      def do_load
+        ManifestFile.create(
+          @params.ctl_bucket,
+          job_id: @job_id,
+          object_urls: @params.object_urls,
+          logger: @logger
+        ) {|manifest|
+          if @params.enable_work_table?
+            prepare_work_table @params.work_table
+            load_objects @params.work_table, manifest, @params.load_options_string
+            @connection.transaction {
+              commit_work_table @params
+              commit_job_result
+            }
+          else
+            @connection.transaction {
+              load_objects @params.dest_table, manifest, @params.load_options_string
+              commit_job_result
+            }
+          end
+        }
+      rescue JobFailure => ex
+        write_job_error 'failure', ex.message
+        raise
+      rescue Exception => ex
+        write_job_error 'error', ex.message
+        raise
+      end
+
+      def prepare_work_table(work_table)
+        @connection.execute("truncate #{work_table}")
+      end
+
+      def load_objects(dest_table, manifest, options)
+        @connection.execute(<<-EndSQL.strip.gsub(/\s+/, ' '))
+            copy #{dest_table}
+            from #{s manifest.url}
+            credentials #{s manifest.credential_string}
+            manifest
+            statupdate false
+            compupdate false
+            #{options}
+            ;
+        EndSQL
+        @logger.info "load succeeded: #{manifest.url}"
+      end
+
+      def commit_work_table(params)
+        @connection.execute(params.sql_source)
+        # keep work table records for later tracking
+      end
+
+      def commit_job_result
+        @end_time = Time.now
+        write_job_result 'success', ''
+      end
+
+      MAX_MESSAGE_LENGTH = 1000
+
+      def write_job_error(status, message)
+        @end_time = Time.now
+        write_job_result status, message.lines.first.strip[0, MAX_MESSAGE_LENGTH]
+      end
+
+      def write_job_result(status, message)
+        @ctl_ds.open {|conn|
+          conn.execute(<<-EndSQL)
+            update
+                strload_jobs
+            set
+                (status, finish_time, message) = (#{s status}, current_timestamp, #{s message})
+            where
+                job_id = #{@job_id}
+            ;
+          EndSQL
+        }
+      end
+
+    end
+
+  end
+
+end

data/lib/bricolage/streamingload/loaderparams.rb

@@ -0,0 +1,153 @@
+require 'bricolage/rubyjobclass'
+require 'bricolage/psqldatasource'
+
+module Bricolage
+
+  module StreamingLoad
+
+    class LoaderParams
+
+      def LoaderParams.load(ctx, task)
+        job = load_job(ctx, task)
+        schema = resolve_schema(ctx, task.schema)
+        job.provide_default 'dest-table', "#{schema}.#{task.table}"
+        #job.provide_sql_file_by_job_id   # FIXME: provide only when exist
+        job.compile
+        new(task, job)
+      end
+
+      def LoaderParams.load_job(ctx, task)
+        if job_file = find_job_file(ctx, task.schema, task.table)
+          ctx.logger.debug "using .job file: #{job_file}"
+          Job.load_file(job_file, ctx.subsystem(task.schema))
+        else
+          ctx.logger.debug "using default job parameters (no .job file)"
+          Job.instantiate(task.table, 'streaming_load_v3', ctx).tap {|job|
+            job.bind_parameters({})
+          }
+        end
+      end
+
+      def LoaderParams.find_job_file(ctx, schema, table)
+        paths = Dir.glob("#{ctx.home_path}/#{schema}/#{table}.*")
+        paths.select {|path| File.extname(path) == '.job' }.sort.first
+      end
+
+      def LoaderParams.resolve_schema(ctx, schema)
+        ctx.global_variables["#{schema}_schema"] || schema
+      end
+      private_class_method :resolve_schema
+
+      def initialize(task, job)
+        @task = task
+        @job = job
+        @params = job.params
+      end
+
+      def task_id
+        @task.id
+      end
+
+      def task_id
+        @task.id
+      end
+
+      def schema
+        @task.schema
+      end
+
+      def table
+        @task.table
+      end
+
+      def force
+        @task.force
+      end
+
+      def object_urls
+        @task.object_urls
+      end
+
+      def ds
+        @params['redshift-ds']
+      end
+
+      def ctl_bucket
+        @params['ctl-ds']
+      end
+
+      def enable_work_table?
+        !!@params['work-table']
+      end
+
+      def work_table
+        @params['work-table']
+      end
+
+      def dest_table
+        @params['dest-table']
+      end
+
+      def load_options_string
+        @params['load-options'].to_s
+      end
+
+      def sql_source
+        sql = @params['sql-file']
+        sql ? sql.source : "insert into #{dest_table} select * from #{work_table};"
+      end
+
+    end
+
+
+    class LoaderJob < RubyJobClass
+
+      job_class_id 'streaming_load_v3'
+
+      def self.parameters(params)
+        params.add DestTableParam.new(optional: false)
+        params.add DestTableParam.new('work-table', optional: true)
+        params.add KeyValuePairsParam.new('load-options', 'OPTIONS', 'Loader options.',
+            optional: true, default: DEFAULT_LOAD_OPTIONS,
+            value_handler: lambda {|value, ctx, vars| PSQLLoadOptions.parse(value) })
+        params.add SQLFileParam.new('sql-file', 'PATH', 'SQL to insert rows from the work table to the target table.', optional: true)
+        params.add DataSourceParam.new('sql', 'redshift-ds', 'Target data source.')
+        params.add DataSourceParam.new('s3', 'ctl-ds', 'Manifest file data source.')
+      end
+
+      def self.default_load_options
+      end
+
+      # Use loosen options by default
+      default_options = [
+        ['json', 'auto'],
+        ['gzip', true],
+        ['timeformat', 'auto'],
+        ['dateformat', 'auto'],
+        ['acceptanydate', true],
+        ['acceptinvchars', ' '],
+        ['truncatecolumns', true],
+        ['trimblanks', true]
+      ]
+      opts = default_options.map {|name, value| PSQLLoadOptions::Option.new(name, value) }
+      DEFAULT_LOAD_OPTIONS = PSQLLoadOptions.new(opts)
+
+      def self.declarations(params)
+        Bricolage::Declarations.new(
+          'dest_table' => nil,
+        )
+      end
+
+      def initialize(params)
+        @params = params
+      end
+
+      def bind(ctx, vars)
+        @params['sql-file'].bind(ctx, vars) if @params['sql-file']
+      end
+
+    end
+
+  end
+
+end

data/lib/bricolage/streamingload/loaderservice.rb

@@ -0,0 +1,163 @@
+require 'bricolage/sqsdatasource'
+require 'bricolage/streamingload/task'
+require 'bricolage/streamingload/loader'
+require 'bricolage/logger'
+require 'bricolage/exception'
+require 'bricolage/version'
+require 'optparse'
+
+module Bricolage
+
+  module StreamingLoad
+
+    class LoaderService
+
+      def LoaderService.main
+        opts = LoaderServiceOptions.new(ARGV)
+        opts.parse
+        unless opts.rest_arguments.size == 1
+          $stderr.puts opts.usage
+          exit 1
+        end
+        config_path, * = opts.rest_arguments
+        config = YAML.load(File.read(config_path))
+        logger = opts.log_file_path ? new_logger(opts.log_file_path, config) : nil
+        ctx = Context.for_application('.', environment: opts.environment, logger: logger)
+        redshift_ds = ctx.get_data_source('sql', config.fetch('redshift-ds'))
+        task_queue = ctx.get_data_source('sqs', config.fetch('task-queue-ds'))
+
+        service = new(
+          context: ctx,
+          control_data_source: ctx.get_data_source('sql', config.fetch('ctl-postgres-ds')),
+          data_source: redshift_ds,
+          task_queue: task_queue,
+          logger: ctx.logger
+        )
+
+        if opts.task_id
+          # Single task mode
+          service.execute_task_by_id opts.task_id
+        else
+          # Server mode
+          Process.daemon(true) if opts.daemon?
+          create_pid_file opts.pid_file_path if opts.pid_file_path
+          service.event_loop
+        end
+      end
+
+      def LoaderService.new_logger(path, config)
+        Logger.new(
+          device: path,
+          rotation_period: config.fetch('log-rotation-period', 'daily'),
+          rotation_size: config.fetch('log-rotation-size', nil)
+        )
+      end
+
+      def LoaderService.create_pid_file(path)
+        File.open(path, 'w') {|f|
+          f.puts $$
+        }
+      rescue
+        # ignore
+      end
+
+      def initialize(context:, control_data_source:, data_source:, task_queue:, logger:)
+        @ctx = context
+        @ctl_ds = control_data_source
+        @ds = data_source
+        @task_queue = task_queue
+        @logger = logger
+      end
+
+      def event_loop
+        @task_queue.main_handler_loop(handlers: self, message_class: Task)
+      end
+
+      def execute_task_by_id(task_id)
+        execute_task load_task(task_id)
+      end
+
+      def load_task(task_id, force: true)
+        @ctl_ds.open {|conn| LoadTask.load(conn, task_id, force: force) }
+      end
+
+      def handle_streaming_load_v3(task)
+        # 1. Load task detail from table
+        # 2. Skip disabled (sqs message should not have disabled state since it will never be exectuted)
+        # 3. Try execute
+        #   - Skip if the task has already been executed AND force = false
+        loadtask = load_task(task.id, force: task.force)
+        return if loadtask.disabled # skip if disabled, but don't delete sqs msg
+        execute_task(loadtask)
+        @task_queue.delete_message(task)
+      end
+
+      def execute_task(task)
+        @logger.info "handling load task: table=#{task.qualified_name} task_id=#{task.id}"
+        loader = Loader.load_from_file(@ctx, @ctl_ds, task, logger: @logger)
+        loader.execute
+      end
+
+    end
+
+    class LoaderServiceOptions
+
+      def initialize(argv)
+        @argv = argv
+        @task_id = nil
+        @daemon = false
+        @log_file_path = nil
+        @pid_file_path = nil
+        @rest_arguments = nil
+
+        @opts = opts = OptionParser.new("Usage: #{$0} CONFIG_PATH")
+        opts.on('--task-id=ID', 'Execute oneshot load task (implicitly disables daemon mode).') {|task_id|
+          @task_id = task_id
+        }
+        opts.on('-e', '--environment=NAME', "Sets execution environment [default: #{Context::DEFAULT_ENV}]") {|env|
+          @environment = env
+        }
+        opts.on('--daemon', 'Becomes daemon in server mode.') {
+          @daemon = true
+        }
+        opts.on('--log-file=PATH', 'Log file path') {|path|
+          @log_file_path = path
+        }
+        opts.on('--pid-file=PATH', 'Creates PID file.') {|path|
+          @pid_file_path = path
+        }
+        opts.on('--help', 'Prints this message and quit.') {
+          puts opts.help
+          exit 0
+        }
+        opts.on('--version', 'Prints version and quit.') {
+          puts "#{File.basename($0)} version #{VERSION}"
+          exit 0
+        }
+      end
+
+      def usage
+        @opts.help
+      end
+
+      def parse
+        @opts.parse!(@argv)
+        @rest_arguments = @argv.dup
+      rescue OptionParser::ParseError => err
+        raise OptionError, err.message
+      end
+
+      attr_reader :rest_arguments, :environment, :log_file_path
+      attr_reader :task_id
+
+      def daemon?
+        @daemon
+      end
+
+      attr_reader :pid_file_path
+
+    end
+
+  end
+
+end