bricolage-streamingload 0.1.0

data/lib/bricolage/streamingload/manifest.rb

@@ -0,0 +1,62 @@
+module Bricolage
+
+  module StreamingLoad
+
+    class ManifestFile
+
+      def ManifestFile.create(ds, job_id:, object_urls:, logger:, noop: false, &block)
+        manifest = new(ds, job_id, object_urls, logger: logger, noop: noop)
+        manifest.create_temporary(&block)
+      end
+
+      def initialize(ds, job_id, object_urls, logger:, noop: false)
+        @ds = ds
+        @job_id = job_id
+        @object_urls = object_urls
+        @logger = logger
+        @noop = noop
+      end
+
+      def credential_string
+        @ds.credential_string
+      end
+
+      def name
+        @name ||= "manifest-#{@job_id}.json"
+      end
+
+      def url
+        @url ||= @ds.url(name)
+      end
+
+      def content
+        @content ||= begin
+          ents = @object_urls.map {|url|
+            { "url" => url, "mandatory" => true }
+          }
+          obj = { "entries" => ents }
+          JSON.pretty_generate(obj)
+        end
+      end
+
+      def put
+        @logger.info "s3: put: #{url}"
+        @ds.object(name).put(body: content) unless @noop
+      end
+
+      def delete
+        @logger.info "s3: delete: #{url}"
+        @ds.object(name).delete unless @noop
+      end
+
+      def create_temporary
+        put
+        yield self
+        delete
+      end
+
+    end
+
+  end
+
+end

data/lib/bricolage/streamingload/objectbuffer.rb

@@ -0,0 +1,211 @@
+require 'bricolage/streamingload/task'
+require 'bricolage/streamingload/loaderparams'
+require 'bricolage/sqlutils'
+require 'json'
+require 'securerandom'
+require 'forwardable'
+
+module Bricolage
+
+  module StreamingLoad
+
+    class LoadableObject
+
+      extend Forwardable
+
+      def initialize(event, components)
+        @event = event
+        @components = components
+      end
+
+      attr_reader :event
+
+      def_delegator '@event', :url
+      def_delegator '@event', :size
+      def_delegator '@event', :message_id
+      def_delegator '@event', :receipt_handle
+      def_delegator '@components', :schema_name
+      def_delegator '@components', :table_name
+
+      def data_source_id
+        "#{schema_name}.#{table_name}"
+      end
+
+      alias qualified_name data_source_id
+
+      def event_time
+        @event.time
+      end
+
+    end
+
+    class ObjectBuffer
+
+      include SQLUtils
+
+      def initialize(control_data_source:, logger:)
+        @ctl_ds = control_data_source
+        @logger = logger
+      end
+
+      def put(obj)
+        @ctl_ds.open {|conn|
+          insert_object(conn, obj)
+        }
+      end
+
+      def flush_tasks
+        task_ids  = []
+        @ctl_ds.open {|conn|
+          conn.transaction {|txn|
+            task_ids = insert_tasks(conn)
+            insert_task_object_mappings(conn) unless task_ids.empty?
+          }
+        }
+        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
+                , 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
+        @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
+      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
+          ;
+        EndSQL
+      end
+
+    end
+
+  end
+
+end

data/lib/bricolage/streamingload/task.rb

@@ -0,0 +1,124 @@
+require 'bricolage/sqsdatasource'
+require 'json'
+
+module Bricolage
+
+  module StreamingLoad
+
+    class Task < SQSMessage
+
+      def Task.get_concrete_class(msg, rec)
+        case
+        when rec['eventName'] == 'streaming_load_v3' then LoadTask
+        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 LoadTask < Task
+
+      def LoadTask.create(task_id:, force: false)
+        super name: 'streaming_load_v3', task_id: task_id, force: force
+      end
+
+      def LoadTask.parse_sqs_record(msg, rec)
+        {
+          task_id: rec['taskId'],
+          force: rec['force'],
+        }
+      end
+
+      def LoadTask.load(conn, task_id, force: false)
+        rec = conn.query_row(<<-EndSQL)
+          select
+              task_class
+              , tbl.schema_name
+              , tbl.table_name
+              , disabled
+          from
+              strload_tasks tsk
+              inner join strload_tables tbl
+                  using(schema_name, table_name)
+          where
+              task_id = #{task_id}
+          ;
+        EndSQL
+        object_urls = conn.query_values(<<-EndSQL)
+          select
+              object_url
+          from
+              strload_task_objects
+              inner join strload_objects
+              using (object_id)
+              inner join strload_tasks
+              using (task_id)
+          where
+              task_id = #{task_id}
+          ;
+        EndSQL
+        return nil unless rec
+        new(
+          name: rec['task_class'],
+          time: nil,
+          source: nil,
+          task_id: task_id,
+          schema: rec['schema_name'],
+          table: rec['table_name'],
+          object_urls: object_urls,
+          disabled: rec['disabled'] == 'f' ? false : true,
+          force: force
+        )
+      end
+
+      alias message_type name
+
+      def init_message(task_id:, schema: nil, table: nil, object_urls: nil, disabled: false, force: false)
+        @id = task_id
+        @force = force
+
+        # Effective only for queue reader process
+        @schema = schema
+        @table = table
+        @object_urls = object_urls
+        @disabled = disabled
+      end
+
+      attr_reader :id, :force
+
+      #
+      # For writer only
+      #
+
+      attr_reader :schema, :table, :object_urls, :disabled
+
+      def qualified_name
+        "#{@schema}.#{@table}"
+      end
+
+      def body
+        obj = super
+        obj['taskId'] = @id
+        obj['schemaName'] = @schema
+        obj['tableName'] = @table
+        obj['objectUrls'] = @object_urls
+        obj['disabled'] = @disabled
+        obj['force'] = @force
+        obj
+      end
+
+    end
+
+  end   # module StreamingLoad
+
+end   # module Bricolage

data/lib/bricolage/streamingload/urlpatterns.rb

@@ -0,0 +1,59 @@
+module Bricolage
+
+  module StreamingLoad
+
+    class URLPatternNotMatched < StandardError; end
+
+
+    class URLPatterns
+
+      def URLPatterns.for_config(configs)
+        new(configs.map {|c|
+          Pattern.new(url: c.fetch('url'), schema: c.fetch('schema'), table: c.fetch('table'))
+        })
+      end
+
+      def initialize(patterns)
+        @patterns = patterns
+      end
+
+      def match(url)
+        @patterns.each do |pat|
+          components = pat.match(url)
+          return components if components
+        end
+        raise URLPatternNotMatched, "no URL pattern matches the object url: #{url.inspect}"
+      end
+
+      class Pattern
+        def initialize(url:, schema:, table:)
+          @url_pattern = /\A#{url}\z/
+          @schema = schema
+          @table = table
+        end
+
+        attr_reader :url_pattern
+        attr_reader :schema
+        attr_reader :table
+
+        def match(url)
+          m = @url_pattern.match(url) or return nil
+          Components.new(get_component(m, @schema), get_component(m, @table))
+        end
+
+        def get_component(m, label)
+          if /\A%/ =~ label
+            m[label[1..-1]]
+          else
+            label
+          end
+        end
+      end
+
+      Components = Struct.new(:schema_name, :table_name)
+
+    end
+
+  end   # module StreamingLoad
+
+end   # module Bricolage