bricolage-streamingload 0.3.0 → 0.4.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA1:
-  metadata.gz: d81ff86cb9addccb7ca9db4d240218679b1f72f9
-  data.tar.gz: f2ec045c994f1c6b619695f74a13aebcc9318722
+  metadata.gz: 848b815669c6580505119917a72a4e97833064c2
+  data.tar.gz: 1809410699822e2a60a21407cad1b5814551adcd
 SHA512:
-  metadata.gz: a78c3b8f35f8d10cbc3da301667ea811874ccadb7653b02388990d63afc2183ee3d0c4e357dccc7799721e1dcd8e33b88e9b3d66147903c32ff109595bced0f6
-  data.tar.gz: a3632f13d3ea039aa690deca646659e0d89ad636b2af55abf5f7011f54cfecddcef53d57152de513bcbaf5c03671af7a512a10781e7846efc60af9c752b8c364
+  metadata.gz: 84ecbe1e548953cc1f4889eb07e9289101e1ade28e89370dae39462b61d193521ceda25f59d8a0b0f1760b8555cf254abab50cf7ed3d11e70d81c422af2f8b82
+  data.tar.gz: b367bb4faa24e9755bee5dbeaaaf937a317826e8e562466e3264a19563bbfcc877ddba9801ffda4875d21e39feddfcc8e87df81a6433b9cd1c3a2cabcefc0754

data/lib/bricolage/nulllogger.rb

@@ -0,0 +1,20 @@
+require 'logger'
+
+module Bricolage
+  # FIXME: should be defined in the Bricolage package
+  class NullLogger
+    def debug(*args) end
+    def debug?() false end
+    def info(*args) end
+    def info?() false end
+    def warn(*args) end
+    def warn?() false end
+    def error(*args) end
+    def error?() false end
+    def exception(*args) end
+    def with_elapsed_time(*args) yield end
+    def elapsed_time(*args) yield end
+    def level() Logger::ERROR end
+    def level=(l) l end
+  end
+end

data/lib/bricolage/sqsdatasource.rb

@@ -1,5 +1,4 @@
 require 'bricolage/datasource'
-require 'bricolage/sqswrapper'
 require 'securerandom'
 require 'aws-sdk'
 require 'json'
@@ -28,6 +27,10 @@ module Bricolage
     attr_reader :access_key_id
     attr_reader :secret_access_key
 
+    attr_reader :visibility_timeout
+    attr_reader :max_number_of_messages
+    attr_reader :wait_time_seconds
+
     def client
       @client ||= begin
         c = @noop ? DummySQSClient.new : Aws::SQS::Client.new(region: @region, access_key_id: @access_key_id, secret_access_key: @secret_access_key)
@@ -39,22 +42,18 @@ module Bricolage
     # High-Level Polling Interface
     #
 
-    def main_handler_loop(handlers:, message_class:)
+    def handle_messages(handler:, message_class:)
       trap_signals
-
-      n_zero = 0
-      until terminating?
-        insert_handler_wait(n_zero)
-        n_msg = handle_messages(handlers: handlers, message_class: message_class)
-        if n_msg == 0
-          n_zero += 1
-        else
-          n_zero = 0
+      polling_loop do
+        result = poll or next true
+        msgs = message_class.for_sqs_result(result)
+        msgs.each do |msg|
+          handler.handle(msg)
         end
-        delete_message_buffer.flush
+        handler.after_message_batch
+        break if terminating?
+        msgs.empty?
       end
-      delete_message_buffer.flush_force
-      logger.info "shutdown gracefully"
     end
 
     def trap_signals
@@ -66,6 +65,7 @@ module Bricolage
         initiate_terminate
       }
     end
+    private :trap_signals
 
     def initiate_terminate
       # No I/O allowed in this method
@@ -76,36 +76,72 @@ module Bricolage
       @terminating
     end
 
-    def insert_handler_wait(n_zero)
-      sec = 2 ** [n_zero, 6].min   # max 64s
-      logger.info "queue wait: sleep #{sec}" if n_zero > 0
-      sleep sec
-    end
-
-    def handle_messages(handlers:, message_class:)
-      n_msg = foreach_message(message_class) do |msg|
-        logger.debug "handling message: #{msg.inspect}" if logger.debug?
-        mid = "handle_#{msg.message_type}"
-        # just ignore unknown event to make app migration easy
-        if handlers.respond_to?(mid, true)
-          handlers.__send__(mid, msg)
+    def polling_loop
+      n_failure = 0
+      while true
+        failed = yield
+        if failed
+          n_failure += 1
         else
-          logger.error "unknown SQS message type: #{msg.message_type.inspect} (message-id: #{msg.message_id})"
+          n_failure = 0
         end
+        insert_handler_wait(n_failure)
       end
-      n_msg
     end
+    private :polling_loop
 
-    def foreach_message(message_class, &block)
+    def insert_handler_wait(n_failure)
+      sec = 2 ** [n_failure, 6].min   # max 64s
+      logger.info "queue wait: sleep #{sec}" if n_failure > 0
+      sleep sec
+    end
+    private :insert_handler_wait
+
+    def poll
       result = receive_messages()
       unless result and result.successful?
         logger.error "ReceiveMessage failed: #{result ? result.error.message : '(result=nil)'}"
         return nil
       end
-      logger.info "receive #{result.messages.size} messages" unless result.messages.empty?
-      msgs = message_class.for_sqs_result(result)
-      msgs.each(&block)
-      msgs.size
+      logger.info "receive #{result.messages.size} messages"
+      result
+    end
+
+    class MessageHandler
+      # abstract logger()
+
+      def handle(msg)
+        logger.debug "handling message: #{msg.inspect}" if logger.debug?
+        if handleable?(msg)
+          call_handler_method(msg)
+        else
+          handle_unknown(msg)
+        end
+      end
+
+      def handleable?(msg)
+        respond_to?(handler_method(msg), true)
+      end
+
+      def call_handler_method(msg)
+        __send__(handler_method(msg), msg)
+      end
+
+      def handler_method(msg)
+        "handle_#{msg.message_type}".intern
+      end
+
+      # Unknown message handler.
+      # Feel free to override this method.
+      def handle_unknown(msg)
+        # just ignore unknown message to make app migration easy
+        logger.error "unknown message type: #{msg.message_type.inspect} (message-id: #{msg.message_id})"
+      end
+
+      # Called after each message batch (ReceiveMessage) is processed.
+      # Override this method in subclasses on demand.
+      def after_message_batch
+      end
     end
 
     #
@@ -122,6 +158,18 @@ module Bricolage
       result
     end
 
+    def put(msg)
+      send_message(msg)
+    end
+
+    def send_message(msg)
+      client.send_message(
+        queue_url: @url,
+        message_body: { 'Records' => [msg.body] }.to_json,
+        delay_seconds: msg.delay_seconds
+      )
+    end
+
     def delete_message(msg)
       client.delete_message(
         queue_url: @url,
@@ -133,20 +181,18 @@ module Bricolage
       delete_message_buffer.put(msg)
     end
 
-    def delete_message_buffer
-      @delete_message_buffer ||= DeleteMessageBuffer.new(client, @url, logger)
+    def process_async_delete(now = Time.now)
+      delete_message_buffer.flush(now)
     end
 
-    def put(msg)
-      send_message(msg)
+    def process_async_delete_force
+      delete_message_buffer.flush_force
     end
 
-    def send_message(msg)
-      client.send_message(
-        queue_url: @url,
-        message_body: { 'Records' => [msg.body] }.to_json,
-        delay_seconds: msg.delay_seconds
-      )
+    private
+
+    def delete_message_buffer
+      @delete_message_buffer ||= DeleteMessageBuffer.new(client, @url, logger)
     end
 
     class DeleteMessageBuffer
@@ -256,6 +302,36 @@ module Bricolage
   end # class SQSDataSource
 
 
+  class SQSClientWrapper
+
+    def initialize(sqs, logger:)
+      @sqs = sqs
+      @logger = logger
+    end
+
+    def receive_message(**args)
+      @logger.debug "receive_message(#{args.inspect})"
+      @sqs.receive_message(**args)
+    end
+
+    def send_message(**args)
+      @logger.debug "send_message(#{args.inspect})"
+      @sqs.send_message(**args)
+    end
+
+    def delete_message(**args)
+      @logger.debug "delete_message(#{args.inspect})"
+      @sqs.delete_message(**args)
+    end
+
+    def delete_message_batch(**args)
+      @logger.debug "delete_message_batch(#{args.inspect})"
+      @sqs.delete_message_batch(**args)
+    end
+
+  end   # class SQSClientWrapper
+
+
   class SQSMessage
 
     SQS_EVENT_SOURCE = 'bricolage:system'

data/lib/bricolage/sqsmock.rb

@@ -0,0 +1,194 @@
+require 'bricolage/sqsdatasource'
+require 'bricolage/nulllogger'
+require 'json'
+
+module Bricolage
+
+  def SQSDataSource.new_mock(**args)
+    SQSDataSource.new(
+      url: 'http://sqs/000000000000/queue-name',
+      access_key_id: 'access_key_id_1',
+      secret_access_key: 'secret_access_key_1',
+      visibility_timeout: 30
+    ).tap {|ds|
+      logger = NullLogger.new
+      #logger = Bricolage::Logger.default
+      ds.__send__(:initialize_base, 'name', nil, logger)
+      ds.instance_variable_set(:@client, SQSMock::Client.new(**args))
+    }
+  end
+
+  module SQSMock
+
+    class Client
+      def initialize(queue: [], receive_message: nil, send_message: nil, delete_message: nil, delete_message_batch: nil)
+        @queue = queue   # [[record]]
+        @call_history = []
+
+        @receive_message = receive_message || lambda {|**args|
+          msgs = @queue.shift or break ReceiveMessageResponse.successful([])
+          ReceiveMessageResponse.successful(msgs)
+        }
+
+        @send_message = send_message || lambda {|**args|
+          SendMessageResponse.successful
+        }
+
+        @delete_message = delete_message || lambda {|**args|
+          Response.successful
+        }
+
+        @delete_message_batch = delete_message_batch || lambda {|queue_url:, entries:|
+          # Returns success for all requests by default.
+          DeleteMessageBatchResponse.new.tap {|res|
+            entries.each do |ent|
+              res.add_success_for(ent)
+            end
+          }
+        }
+      end
+
+      # Free free to modify this array contents
+      attr_reader :call_history
+
+      def self.def_mock_method(name)
+        define_method(name) {|**args|
+          @call_history.push CallHistory.new(name.intern, args)
+          instance_variable_get("@#{name}").(**args)
+        }
+      end
+
+      def_mock_method :receive_message
+      def_mock_method :send_message
+      def_mock_method :delete_message
+      def_mock_method :delete_message_batch
+    end
+
+    CallHistory = Struct.new(:name, :args)
+
+    # success/failure only result
+    class Response
+      def Response.successful
+        new(successful: true)
+      end
+
+      def initialize(successful:)
+        @successful = successful
+      end
+
+      def successful?
+        @successful
+      end
+    end
+
+    class ReceiveMessageResponse < Response
+      def ReceiveMessageResponse.successful(msgs)
+        new(successful: true, messages: msgs)
+      end
+
+      def initialize(successful:, messages:)
+        super(successful: successful)
+        @messages = messages
+      end
+
+      attr_reader :messages
+    end
+
+    class SendMessageResponse < Response
+      def SendMessageResponse.successful
+        new(successful: true, message_id: "sqs-sent-message-id-#{Message.new_seq}")
+      end
+
+      def initialize(successful:, message_id:)
+        super(successful: successful)
+        @message_id = message_id
+      end
+
+      attr_reader :message_id
+    end
+
+    class DeleteMessageBatchResponse
+      def initialize(successful: [], failed: [])
+        @successful = successful
+        @failed = failed
+      end
+
+      attr_reader :successful
+      attr_reader :failed
+
+      Success = Struct.new(:id)
+      Failure = Struct.new(:id, :sender_fault, :code, :message)
+
+      def add_success_for(ent)
+        @successful.push Success.new(ent[:id])
+      end
+
+      def add_failure_for(ent)
+        @failed.push Failure.new(ent[:id], true, '400', 'some reason')
+      end
+    end
+
+    class Message
+      def Message.s3_object_created_event(url)
+        raise "is not a S3 URL: #{url.inspect}" unless %r<\As3://\w> =~ url
+        bucket, key = url.sub(%r<s3://>, '').split('/', 2)
+        with_body({
+          eventVersion: '2.0',
+          eventSource: 'aws:s3',
+          awsRegion: 'ap-northeast-1',
+          eventTime: Time.now.iso8601,
+          eventName: 'ObjectCreated:Put',
+          s3: {
+            s3SchemaVersion: '1.0',
+            configurationId: 'TestConfig',
+            bucket: {
+              name: bucket,
+              arn: "arn:aws:s3:::#{bucket}"
+            },
+            object: {
+              key: key,
+              size: 1024
+            }
+          }
+        })
+      end
+
+      @seq = 0
+
+      def Message.new_seq
+        @seq += 1
+        @seq
+      end
+
+      def Message.with_body(body)
+        seq = new_seq
+        new(
+          message_id: "sqs-message-id-#{seq}",
+          receipt_handle: "sqs-receipt-handle-#{seq}",
+          body: body
+        )
+      end
+
+      def initialize(message_id: nil, receipt_handle: nil, body: nil)
+        @message_id = message_id
+        @receipt_handle = receipt_handle
+        @body = body
+        @body_json = { Records: [body] }.to_json
+      end
+
+      attr_reader :message_id
+      attr_reader :receipt_handle
+
+      def body
+        @body_json
+      end
+
+      # for debug
+      def body_object
+        @body
+      end
+    end
+
+  end   # module SQSMock
+
+end   # module Bricolage

data/lib/bricolage/streamingload/dispatcher.rb

@@ -1,3 +1,4 @@
+require 'bricolage/context'
 require 'bricolage/exception'
 require 'bricolage/version'
 require 'bricolage/sqsdatasource'
@@ -15,7 +16,7 @@ module Bricolage
 
   module StreamingLoad
 
-    class Dispatcher
+    class Dispatcher < SQSDataSource::MessageHandler
 
       def Dispatcher.main
         opts = DispatcherOptions.new(ARGV)
@@ -54,7 +55,6 @@ module Bricolage
 
         Process.daemon(true) if opts.daemon?
         create_pid_file opts.pid_file_path if opts.pid_file_path
-        dispatcher.set_dispatch_timer
         dispatcher.event_loop
       end
 
@@ -82,10 +82,24 @@ module Bricolage
         @dispatch_interval = dispatch_interval
         @dispatch_message_id = nil
         @logger = logger
+        @checkpoint_requested = false
       end
 
+      attr_reader :logger
+
       def event_loop
-        @event_queue.main_handler_loop(handlers: self, message_class: Event)
+        set_dispatch_timer
+        @event_queue.handle_messages(handler: self, message_class: Event)
+        @event_queue.process_async_delete_force
+        logger.info "shutdown gracefully"
+      end
+
+      # override
+      def after_message_batch
+        @event_queue.process_async_delete
+        if @checkpoint_requested
+          create_checkpoint
+        end
       end
 
       def handle_shutdown(e)
@@ -94,6 +108,29 @@ module Bricolage
         @event_queue.delete_message(e)
       end
 
+      def handle_checkpoint(e)
+        # Delay creating CHECKPOINT after the current message batch,
+        # because any other extra events are already received.
+        @checkpoint_requested = true
+        # Delete this event immediately
+        @event_queue.delete_message(e)
+      end
+
+      def create_checkpoint
+        logger.info "*** Creating checkpoint requested ***"
+        logger.info "Force-flushing all objects..."
+        flush_all_tasks_immediately
+        logger.info "All objects flushed; shutting down..."
+        @event_queue.initiate_terminate
+      end
+
+      def flush_all_tasks_immediately
+        tasks = @object_buffer.flush_tasks_force
+        tasks.each do |task|
+          @task_queue.put task
+        end
+      end
+
       def handle_data(e)
         unless e.created?
           @event_queue.delete_message_async(e)

data/lib/bricolage/streamingload/event.rb

@@ -11,6 +11,7 @@ module Bricolage
         when rec['eventName'] == 'shutdown' then ShutdownEvent
         when rec['eventName'] == 'dispatch' then DispatchEvent
         when rec['eventName'] == 'flush' then FlushEvent
+        when rec['eventName'] == 'checkpoint' then CheckPointEvent
         when rec['eventSource'] == 'aws:s3'
           S3ObjectEvent
         else
@@ -41,7 +42,26 @@ module Bricolage
 
       alias message_type name
 
-      def init_message
+      def init_message(dummy: nil)
+      end
+
+    end
+
+
+    # Flushes all tables and shutdown
+    class CheckPointEvent < Event
+
+      def CheckPointEvent.create
+        super name: 'checkpoint'
+      end
+
+      def CheckPointEvent.parse_sqs_record(msg, rec)
+        {}
+      end
+
+      alias message_type name
+
+      def init_message(dummy: nil)
       end
 
     end
@@ -75,6 +95,7 @@ module Bricolage
 
     end
 
+
     class DispatchEvent < Event
 
       def DispatchEvent.create(delay_seconds:)
@@ -83,8 +104,9 @@ module Bricolage
 
       alias message_type name
 
-      def init_message(dummy)
+      def init_message(dummy: nil)
       end
+
     end
 
 

data/lib/bricolage/streamingload/loaderservice.rb

@@ -11,7 +11,7 @@ module Bricolage
 
   module StreamingLoad
 
-    class LoaderService
+    class LoaderService < SQSDataSource::MessageHandler
 
       def LoaderService.main
         opts = LoaderServiceOptions.new(ARGV)
@@ -76,7 +76,8 @@ module Bricolage
       end
 
       def event_loop
-        @task_queue.main_handler_loop(handlers: self, message_class: Task)
+        @task_queue.handle_messages(handler: self, message_class: Task)
+        @logger.info "shutdown gracefully"
       end
 
       def execute_task_by_id(task_id)
@@ -87,6 +88,7 @@ module Bricolage
         @ctl_ds.open {|conn| LoadTask.load(conn, task_id, force: force) }
       end
 
+      # message handler
       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)

data/lib/bricolage/streamingload/objectbuffer.rb

@@ -39,6 +39,7 @@ module Bricolage
 
     end
 
+
     class ObjectBuffer
 
       include SQLUtils
@@ -55,7 +56,7 @@ module Bricolage
       end
 
       def flush_tasks
-        task_ids  = []
+        task_ids = nil
         @ctl_ds.open {|conn|
           conn.transaction {|txn|
             task_ids = insert_tasks(conn)
@@ -65,145 +66,192 @@ module Bricolage
         return task_ids.map {|id| LoadTask.create(task_id: id) }
       end
 
+      # Flushes all objects of all tables immediately with no
+      # additional conditions, to create "stream checkpoint".
+      def flush_tasks_force
+        task_ids  = []
+        @ctl_ds.open {|conn|
+          conn.transaction {|txn|
+            # insert_task_object_mappings may not consume all saved objects
+            # (e.g. there are too many objects for one table), we must create
+            # tasks repeatedly until there are no unassigned objects.
+            until (ids = insert_tasks_force(conn)).empty?
+              insert_task_object_mappings(conn)
+              task_ids.concat ids
+            end
+          }
+        }
+        return task_ids.map {|id| LoadTask.create(task_id: id) }
+      end
+
       private
 
       def insert_object(conn, obj)
-        #HACK - suppress log per object
-        log_level = @logger.level
-        @logger.level = Logger::ERROR
-        conn.update(<<-EndSQL)
-            insert into strload_objects
-                (object_url
-                , object_size
-                , data_source_id
-                , message_id
-                , event_time
+        suppress_sql_logging {
+          conn.update(<<-EndSQL)
+              insert into strload_objects
+                  ( object_url
+                  , object_size
+                  , data_source_id
+                  , message_id
+                  , event_time
+                  , submit_time
+                  )
+              select
+                  #{s obj.url}
+                  , #{obj.size}
+                  , #{s obj.data_source_id}
+                  , #{s obj.message_id}
+                  , '#{obj.event_time}' AT TIME ZONE 'JST'
+                  , current_timestamp
+              from
+                  strload_tables
+              where
+                  data_source_id = #{s obj.data_source_id}
+              ;
+          EndSQL
+        }
+      end
+
+      def insert_tasks_force(conn)
+        insert_tasks(conn, force: true)
+      end
+
+      def insert_tasks(conn, force: false)
+        task_ids = conn.query_values(<<-EndSQL)
+            insert into strload_tasks
+                ( task_class
+                , schema_name
+                , table_name
                 , submit_time
                 )
             select
-                #{s obj.url}
-                , #{obj.size}
-                , #{s obj.data_source_id}
-                , #{s obj.message_id}
-                , '#{obj.event_time}' AT TIME ZONE 'JST'
+                'streaming_load_v3'
+                , tbl.schema_name
+                , tbl.table_name
                 , current_timestamp
             from
-                strload_tables
+                strload_tables tbl
+
+                -- number of objects not assigned to a task for each schema_name.table_name (> 0)
+                inner join (
+                    select
+                        data_source_id
+                        , count(*) as object_count
+                    from
+                        (
+                            select
+                                min(object_id) as object_id
+                                , object_url
+                            from
+                                strload_objects
+                            group by
+                                object_url
+                        ) uniq_objects
+                        inner join strload_objects using (object_id)
+                        left outer join strload_task_objects using (object_id)
+                    where
+                        task_id is null -- not assigned to a task
+                    group by
+                        data_source_id
+                ) obj
+                using (data_source_id)
+
+                -- preceeding task's submit time
+                left outer join (
+                    select
+                        schema_name
+                        , table_name
+                        , max(submit_time) as latest_submit_time
+                    from
+                        strload_tasks
+                    group by
+                        schema_name, table_name
+                ) task
+                using (schema_name, table_name)
             where
-                data_source_id = #{s obj.data_source_id}
+                not tbl.disabled -- not disabled
+                and (
+                    #{force ? "true or" : ""}                                                      -- Creates tasks with no conditions if forced
+                    obj.object_count > tbl.load_batch_size                                         -- batch_size exceeded?
+                    or extract(epoch from current_timestamp - latest_submit_time) > load_interval  -- load_interval exceeded?
+                    or latest_submit_time is null                                                  -- no previous tasks?
+                )
+            returning task_id
             ;
         EndSQL
-        @logger.level = log_level
-      end
 
-      def insert_tasks(conn)
-        vals = conn.query_values(<<-EndSQL)
-          insert into
-              strload_tasks (task_class, schema_name, table_name, submit_time)
-          select
-              'streaming_load_v3'
-              , tbl.schema_name
-              , tbl.table_name
-              , current_timestamp
-          from
-              strload_tables tbl
-              inner join (
-                  select
-                      data_source_id
-                      , count(*) as object_count
-                  from (
-                      select
-                          min(object_id) as object_id
-                          , object_url
-                      from
-                          strload_objects
-                      group by
-                          object_url
-                      ) uniq_objects
-                      inner join strload_objects
-                          using(object_id)
-                      left outer join strload_task_objects
-                          using(object_id)
-                  where
-                      task_id is null -- not assigned to a task
-                  group by
-                      data_source_id
-                  ) obj -- number of objects not assigned to a task per schema_name.table_name (won't return zero)
-                  using (data_source_id)
-              left outer join (
-                  select
-                      schema_name
-                      , table_name
-                      , max(submit_time) as latest_submit_time
-                  from
-                      strload_tasks
-                  group by
-                      schema_name, table_name
-                  ) task -- preceeding task's submit time
-                  using(schema_name, table_name)
-          where
-              not tbl.disabled -- not disabled
-              and (
-                obj.object_count > tbl.load_batch_size -- batch_size exceeded?
-                or extract(epoch from current_timestamp - latest_submit_time) > load_interval -- load_interval exceeded?
-                or latest_submit_time is null -- no last task
-              )
-          returning task_id
-          ;
-        EndSQL
-        @logger.info "Number of task created: #{vals.size}"
-        vals
+        @logger.info "Number of task created: #{task_ids.size}"
+        task_ids
       end
 
       def insert_task_object_mappings(conn)
         conn.update(<<-EndSQL)
-          insert into
-              strload_task_objects
-          select
-              task_id
-              , object_id
-          from (
-              select
-                  row_number() over(partition by task.task_id order by obj.object_id) as object_count
-                  , task.task_id
-                  , obj.object_id
-                  , load_batch_size
-              from (
-                  select
-                      min(object_id) as object_id
-                      , object_url
-                      , data_source_id
-                  from
-                      strload_objects
-                  group by
-                      2, 3
-                  ) obj
-                  inner join (
-                      select
-                          min(task_id) as task_id -- oldest task
-                          , tbl.data_source_id
-                          , max(load_batch_size) as load_batch_size
-                      from
-                          strload_tasks
-                          inner join strload_tables tbl
-                              using(schema_name, table_name)
-                      where
-                          task_id not in (select distinct task_id from strload_task_objects) -- no assigned objects
-                      group by
-                          2
-                      ) task -- tasks without objects
-                      using(data_source_id)
-                  left outer join strload_task_objects task_obj
-                      using(object_id)
-              where
-                  task_obj.object_id is null -- not assigned to a task
-              ) as t
-          where
-              object_count <= load_batch_size -- limit number of objects assigned to single task
-          ;
+            insert into strload_task_objects
+                ( task_id
+                , object_id
+                )
+            select
+                task_id
+                , object_id
+            from (
+                select
+                    row_number() over (partition by task.task_id order by obj.object_id) as object_count
+                    , task.task_id
+                    , obj.object_id
+                    , load_batch_size
+                from
+                    (
+                        select
+                            data_source_id
+                            , object_url
+                            , min(object_id) as object_id
+                        from
+                            strload_objects
+                        group by
+                            1, 2
+                    ) obj
+
+                    -- tasks without objects
+                    inner join (
+                        select
+                            tbl.data_source_id
+                            , min(task_id) as task_id   -- pick up oldest task
+                            , max(load_batch_size) as load_batch_size
+                        from
+                            strload_tasks
+                            inner join strload_tables tbl
+                            using (schema_name, table_name)
+                        where
+                            -- unassigned objects
+                            task_id not in (select task_id from strload_task_objects)
+                        group by
+                            1
+                    ) task
+                    using (data_source_id)
+
+                    left outer join strload_task_objects task_obj
+                    using (object_id)
+                where
+                    task_obj.object_id is null   -- unassigned to a task
+                ) as t
+            where
+                object_count <= load_batch_size   -- limit number of objects assigned to single task
+            ;
         EndSQL
       end
 
+      def suppress_sql_logging
+        # CLUDGE
+        orig = @logger.level
+        begin
+          @logger.level = Logger::ERROR
+          yield
+        ensure
+          @logger.level = orig
+        end
+      end
+
     end
 
   end

data/lib/bricolage/streamingload/version.rb

@@ -1,5 +1,5 @@
 module Bricolage
   module StreamingLoad
-    VERSION = '0.3.0'
+    VERSION = '0.4.0'
   end
 end

data/test/streamingload/test_dispatcher.rb

@@ -0,0 +1,111 @@
+require 'test/unit'
+require 'bricolage/context'
+require 'bricolage/sqsdatasource'
+require 'bricolage/sqsmock'
+require 'bricolage/streamingload/dispatcher'
+
+module Bricolage
+  module StreamingLoad
+
+    class TestDispatcher < Test::Unit::TestCase
+
+      test "checkpoint event" do
+        ctx = Context.for_application('.', environment: 'test', logger: NullLogger.new)
+        ctl_ds = ctx.get_data_source('sql', 'dwhctl')
+
+        event_queue = SQSDataSource.new_mock(queue: [
+          # 1st ReceiveMessage
+          [
+            SQSMock::Message.s3_object_created_event('s3://test-bucket/testschema.desttable/datafile-0001.json.gz'),
+            SQSMock::Message.s3_object_created_event('s3://test-bucket/testschema.desttable/datafile-0002.json.gz'),
+            SQSMock::Message.s3_object_created_event('s3://test-bucket/testschema.desttable/datafile-0003.json.gz'),
+            SQSMock::Message.s3_object_created_event('s3://test-bucket/testschema.desttable/datafile-0004.json.gz'),
+            SQSMock::Message.s3_object_created_event('s3://test-bucket/testschema.desttable/datafile-0005.json.gz')
+          ],
+          # 2nd ReceiveMessage
+          [
+            SQSMock::Message.s3_object_created_event('s3://test-bucket/testschema.desttable/datafile-0006.json.gz'),
+            SQSMock::Message.s3_object_created_event('s3://test-bucket/testschema.desttable/datafile-0007.json.gz'),
+            SQSMock::Message.s3_object_created_event('s3://test-bucket/testschema.desttable/datafile-0008.json.gz'),
+            SQSMock::Message.s3_object_created_event('s3://test-bucket/testschema.desttable/datafile-0009.json.gz'),
+            SQSMock::Message.new(body: {eventSource: 'bricolage:system', eventName: 'checkpoint'}),
+            SQSMock::Message.s3_object_created_event('s3://test-bucket/testschema.desttable/datafile-0010.json.gz')
+          ]
+        ])
+
+        task_queue = SQSDataSource.new_mock
+
+        object_buffer = ObjectBuffer.new(
+          control_data_source: ctl_ds,
+          logger: ctx.logger
+        )
+
+        url_patterns = URLPatterns.for_config([
+          {
+            "url" => %r<\As3://test-bucket/testschema\.desttable/datafile-\d{4}\.json\.gz>.source,
+            "schema" => 'testschema',
+            "table" => 'desttable'
+          }
+        ])
+
+        dispatcher = Dispatcher.new(
+          event_queue: event_queue,
+          task_queue: task_queue,
+          object_buffer: object_buffer,
+          url_patterns: url_patterns,
+          dispatch_interval: 600,
+          logger: ctx.logger
+        )
+
+        # FIXME: database cleaner
+        ctl_ds.open {|conn|
+          conn.update("truncate strload_tables")
+          conn.update("truncate strload_objects")
+          conn.update("truncate strload_task_objects")
+          conn.update("truncate strload_tasks")
+          conn.update("insert into strload_tables values ('testschema', 'desttable', 'testschema.desttable', 100, 1800, false)")
+        }
+        dispatcher.event_loop
+
+        # Event Queue Call Sequence
+        hst = event_queue.client.call_history
+        assert_equal :send_message, hst[0].name      # start flush timer
+        assert_equal :receive_message, hst[1].name
+        assert_equal :delete_message_batch, hst[2].name
+        assert_equal :receive_message, hst[3].name
+        assert_equal :delete_message, hst[4].name    # delete checkpoint
+        assert_equal :delete_message_batch, hst[5].name
+
+        # Task Queue Call Sequence
+        hst = task_queue.client.call_history
+        assert_equal :send_message, hst[0].name
+        assert(/streaming_load_v3/ =~ hst[0].args[:message_body])
+        task_id = JSON.load(hst[0].args[:message_body])['Records'][0]['taskId'].to_i
+        assert_not_equal 0, task_id
+
+        # Object Buffer
+        assert_equal [], unassigned_objects(ctl_ds)
+        task = ctl_ds.open {|conn| LoadTask.load(conn, task_id) }
+        assert_equal 'testschema', task.schema
+        assert_equal 'desttable', task.table
+        assert_equal 10, task.object_urls.size
+      end
+
+      def unassigned_objects(ctl_ds)
+        ctl_ds.open {|conn|
+          conn.query_values(<<-EndSQL)
+              select
+                  object_url
+              from
+                  strload_objects
+              where
+                  object_id not in (select object_id from strload_task_objects)
+              ;
+          EndSQL
+        }
+      end
+
+    end
+
+  end
+end

data/test/test_sqsdatasource.rb

@@ -1,111 +1,54 @@
 require 'test/unit'
 require 'bricolage/streamingload/event'
+require 'bricolage/sqsmock'
 require 'bricolage/logger'
 
 module Bricolage
 
   class TestSQSDataSource < Test::Unit::TestCase
 
-    def new_sqs_ds(mock_client = nil)
-      SQSDataSource.new(
-        url: 'http://sqs/000000000000/queue-name',
-        access_key_id: 'access_key_id_1',
-        secret_access_key: 'secret_access_key_1',
-        visibility_timeout: 30
-      ).tap {|ds|
-        logger = NullLogger.new
-        #logger = Bricolage::Logger.default
-        ds.__send__(:initialize_base, 'name', nil, logger)
-        ds.instance_variable_set(:@client, mock_client) if mock_client
-      }
-    end
-
-    class MockSQSClient
-      def initialize(&block)
-        @handler = block
-      end
-
-      def delete_message_batch(**args)
-        @handler.call(args)
-      end
-    end
-
-    class NullLogger
-      def debug(*args) end
-      def info(*args) end
-      def warn(*args) end
-      def error(*args) end
-      def exception(*args) end
-      def with_elapsed_time(*args) yield end
-      def elapsed_time(*args) yield end
-    end
-
-    def sqs_message(seq)
-      MockSQSMessage.new("message_id_#{seq}", "receipt_handle_#{seq}")
-    end
-
-    MockSQSMessage = Struct.new(:message_id, :receipt_handle)
-
-    class MockSQSResponse
-      def initialize(successful: [], failed: [])
-        @successful = successful
-        @failed = failed
-      end
-
-      attr_reader :successful
-      attr_reader :failed
-
-      Success = Struct.new(:id)
-      Failure = Struct.new(:id, :sender_fault, :code, :message)
-
-      def add_success_for(ent)
-        @successful.push Success.new(ent[:id])
-      end
-
-      def add_failure_for(ent)
-        @failed.push Failure.new(ent[:id], true, '400', 'some reason')
-      end
-    end
-
     test "#delete_message_async" do
-      messages = [sqs_message(0), sqs_message(1), sqs_message(2)]
-      mock = MockSQSClient.new {|args|
-        entries = args[:entries]
-        if entries.size == 3
-          # first time
-          assert_equal messages[0].receipt_handle, entries[0][:receipt_handle]
-          assert_equal messages[1].receipt_handle, entries[1][:receipt_handle]
-          assert_equal messages[2].receipt_handle, entries[2][:receipt_handle]
-          MockSQSResponse.new.tap {|res|
-            res.add_success_for(entries[0])
-            res.add_failure_for(entries[1])
-            res.add_success_for(entries[2])
-          }
-        else
-          # second time
-          MockSQSResponse.new.tap {|res|
-            res.add_success_for(entries[0])
-          }
-        end
+      messages = (0..2).map {|seq|
+        SQSMock::Message.new(message_id: "message_id_#{seq}", receipt_handle: "receipt_handle_#{seq}")
       }
-      ds = new_sqs_ds(mock)
+      ds = SQSDataSource.new_mock(
+        delete_message_batch: -> (queue_url:, entries:) {
+          if entries.size == 3
+            # first time
+            assert_equal messages[0].receipt_handle, entries[0][:receipt_handle]
+            assert_equal messages[1].receipt_handle, entries[1][:receipt_handle]
+            assert_equal messages[2].receipt_handle, entries[2][:receipt_handle]
+            SQSMock::DeleteMessageBatchResponse.new.tap {|res|
+              res.add_success_for(entries[0])
+              res.add_failure_for(entries[1])
+              res.add_success_for(entries[2])
+            }
+          else
+            # second time
+            SQSMock::DeleteMessageBatchResponse.new.tap {|res|
+              res.add_success_for(entries[0])
+            }
+          end
+        }
+      )
+
       ds.delete_message_async(messages[0])
       ds.delete_message_async(messages[1])
       ds.delete_message_async(messages[2])
 
       # first flush
       flush_time = Time.now
-      ds.delete_message_buffer.flush(flush_time)
-      assert_equal 1, ds.delete_message_buffer.size
-      bufent = ds.delete_message_buffer.instance_variable_get(:@buf).values.first
+      ds.process_async_delete(flush_time)
+      delete_buf = ds.__send__(:delete_message_buffer)
+      bufent = delete_buf.instance_variable_get(:@buf).values.first
       assert_equal 'receipt_handle_1', bufent.message.receipt_handle
       assert_equal 1, bufent.n_failure
       assert_false bufent.issuable?(flush_time)
       assert_true bufent.issuable?(flush_time + 180)
 
       # second flush
-      ds.delete_message_buffer.flush(flush_time + 180)
-      assert_true ds.delete_message_buffer.empty?
+      ds.process_async_delete(flush_time + 180)
+      assert_true delete_buf.empty?
     end
 
   end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: bricolage-streamingload
 version: !ruby/object:Gem::Version
-  version: 0.3.0
+  version: 0.4.0
 platform: ruby
 authors:
 - Minero Aoki
@@ -107,9 +107,10 @@ files:
 - README.md
 - bin/bricolage-streaming-dispatcher
 - bin/bricolage-streaming-loader
+- lib/bricolage/nulllogger.rb
 - lib/bricolage/snsdatasource.rb
 - lib/bricolage/sqsdatasource.rb
-- lib/bricolage/sqswrapper.rb
+- lib/bricolage/sqsmock.rb
 - lib/bricolage/streamingload/alertinglogger.rb
 - lib/bricolage/streamingload/dispatcher.rb
 - lib/bricolage/streamingload/event.rb
@@ -122,6 +123,7 @@ files:
 - lib/bricolage/streamingload/urlpatterns.rb
 - lib/bricolage/streamingload/version.rb
 - test/all.rb
+- test/streamingload/test_dispatcher.rb
 - test/streamingload/test_event.rb
 - test/test_sqsdatasource.rb
 homepage: https://github.com/aamine/bricolage-streamingload

data/lib/bricolage/sqswrapper.rb

@@ -1,77 +0,0 @@
-require 'json'
-
-module Bricolage
-
-  class SQSClientWrapper
-    def initialize(sqs, logger:)
-      @sqs = sqs
-      @logger = logger
-    end
-
-    def receive_message(**args)
-      @logger.debug "receive_message(#{args.inspect})"
-      @sqs.receive_message(**args)
-    end
-
-    def send_message(**args)
-      @logger.debug "send_message(#{args.inspect})"
-      @sqs.send_message(**args)
-    end
-
-    def delete_message(**args)
-      @logger.debug "delete_message(#{args.inspect})"
-      @sqs.delete_message(**args)
-    end
-
-    def delete_message_batch(**args)
-      @logger.debug "delete_message_batch(#{args.inspect})"
-      @sqs.delete_message_batch(**args)
-    end
-  end
-
-
-  class DummySQSClient
-    def initialize(queue = [])
-      @queue = queue
-    end
-
-    def receive_message(**args)
-      msg_recs = @queue.shift or return EMPTY_RESULT
-      msgs = msg_recs.map {|recs| Message.new({'Records' => recs}.to_json) }
-      Result.new(true, msgs)
-    end
-
-    def send_message(**args)
-      SUCCESS_RESULT
-    end
-
-    def delete_message(**args)
-      SUCCESS_RESULT
-    end
-
-    class Result
-      def initialize(successful, messages = nil)
-        @successful = successful
-        @messages = messages
-      end
-
-      def successful?
-        @successful
-      end
-
-      attr_reader :messages
-    end
-
-    SUCCESS_RESULT = Result.new(true)
-    EMPTY_RESULT = Result.new(true, [])
-
-    class Message
-      def initialize(body)
-        @body = body
-      end
-
-      attr_reader :body
-    end
-  end
-
-end   # module Bricolage