bricolage 5.28.1 → 6.0.0beta1

data/lib/bricolage/psqldatasource.rb

@@ -15,6 +15,8 @@ module Bricolage
 
     include CommandUtils
 
+    DEFAULT_RETRY_LIMIT = 3
+
     def initialize(
         host: 'localhost',
         port: 5439,
@@ -24,7 +26,9 @@ module Bricolage
         pgpass: nil,
         encoding: nil,
         psql: 'psql',
-        sql_log_level: Logger::INFO,
+        sql_log_level: nil,
+        query_sql_log_level: nil,
+        update_sql_log_level: nil,
         tmpdir: Dir.tmpdir)
       @host = host
       @port = port
@@ -34,8 +38,11 @@ module Bricolage
       @pgpass = pgpass
       @encoding = encoding
       @psql = psql
-      @sql_log_level = Logger.intern_severity(sql_log_level)
+      @sql_log_level = Logger.intern_severity(sql_log_level || Logger::DEBUG)
+      @query_sql_log_level = Logger.intern_severity(query_sql_log_level || sql_log_level || Logger::DEBUG)
+      @update_sql_log_level = Logger.intern_severity(update_sql_log_level || sql_log_level || Logger::INFO)
       @tmpdir = tmpdir
+      @connection_pool = []
       raise ParameterError, "missing psql host" unless @host
       raise ParameterError, "missing psql port" unless @port
       raise ParameterError, "missing psql database" unless @database
@@ -51,6 +58,8 @@ module Bricolage
     attr_reader :user
 
     attr_reader :sql_log_level
+    attr_reader :query_sql_log_level
+    attr_reader :update_sql_log_level
 
     def new_task
       PSQLTask.new(self)
@@ -71,7 +80,15 @@ module Bricolage
             '--no-password',
             *options,
             env: get_psql_env
-        msg = LogLocator.slice_last_stderr(/^psql:.*?:\d+: ERROR: (.*)/, 1) unless st.success?
+        unless st.success?
+          begin
+            msg = LogLocator.slice_last_stderr(/^psql:.*?:\d+: ERROR: (.*)/, 1)
+          rescue IOError => ex
+            # slice_last_stderr may fail if stderr is not a file
+            logger.error ex.message
+            msg = nil
+          end
+        end
         JobResult.for_process_status(st, msg)
       }
     end
@@ -101,7 +118,48 @@ module Bricolage
     end
 
     def open(&block)
-      PostgresConnection.open_data_source(self, &block)
+      retries = (ENV['BRICOLAGE_OPEN_RETRY_LIMIT'] || DEFAULT_RETRY_LIMIT).to_i
+      begin
+        conn = PostgresConnection.open_data_source(self)
+        conn.execute_query('select 1'){}
+      rescue PG::ConnectionBad, PG::UnableToSend => ex
+        retries -= 1
+        if retries >= 0
+          logger.warn "Retry PG connection for execute query: #{ex.message}"
+          sleep 1
+          retry
+        end
+      end
+      if block_given?
+        yield conn
+      else
+        return conn
+      end
+    end
+
+    def open_shared_connection
+      raise ParameterError, 'open_shared_connection require block' unless block_given?
+      conn = nil
+      if @connection_pool.empty?
+        conn = open
+      else
+        begin
+          conn = @connection_pool.shift
+          conn.execute_query('select 1'){}
+        rescue
+          conn.close
+          conn = open
+        end
+      end
+
+      yield conn
+    ensure
+      @connection_pool.push(conn)
+    end
+
+    def clear_connection_pool
+      @connection_pool.map(&:close)
+      @connection_pool = []
     end
 
     def query_batch(query, batch_size = 5000, &block)
@@ -320,7 +378,6 @@ module Bricolage
       unless src_ds.redshift_loader_source?
         raise ParameterError, "input data source does not support redshift as bulk loading source: #{src_ds.name}"
       end
-      opts.provide_defaults(src_ds)
       buf = StringIO.new
       buf.puts "copy #{dest_table}"
       buf.puts "from '#{src_ds.url(src_path)}'"
@@ -340,12 +397,23 @@ module Bricolage
       when 'csv'
         %q(delimiter ',')
       when 'json'
-        jsonpath ? "json \'#{src_ds.url(jsonpath)}\'" : %q(json 'auto')
+        "json '#{json_param(jsonpath)}'"
       else
         raise ParameterError, "unsupported format: #{fmt}"
       end
     end
 
+    def json_param(jsonpath)
+      case jsonpath
+      when nil
+        'auto'
+      when %r{\As3://}
+        jsonpath
+      else
+        src_ds.url(jsonpath)
+      end
+    end
+
     def unload(stmt, dest_ds, dest_path, format, opts)
       exec unload_statement(stmt, dest_ds, dest_path, format, opts)
     end
@@ -376,7 +444,7 @@ module Bricolage
     end
 
     def format_query(query)
-      query.gsub(/^--.*/, '').strip.gsub(/[ \t]*\n[ \t]*/, ' ').gsub("'", "\\\\'")
+      query.gsub(/^--.*/, '').strip.gsub(/[ \t]*\n[ \t]*/, ' ').gsub(/\\/,"\\\\\\\\").gsub("'", "\\\\'")
     end
   end
 

data/lib/bricolage/sqlutils.rb

@@ -11,7 +11,7 @@ module Bricolage
     alias s sql_string_literal
 
     def escape_sql_string(s)
-      s.gsub(/'/, "''")
+      s.gsub(/'/, "''").gsub(/\\/, '\\\\')
     end
 
     def sql_timestamp_literal(time)
@@ -24,6 +24,48 @@ module Bricolage
       time.strftime('%Y-%m-%d %H:%M:%S')
     end
 
+    def compile_set_expr(values_hash)
+      columns = values_hash.keys.map(&:to_s).join(', ')
+      values = values_hash.values.map{|v| convert_value(v) }.join(', ')
+      return columns, values
+    end
+
+    def convert_value(value)
+      if value == :now
+        'now()'
+      elsif value.nil?
+        "null"
+      elsif value == true or value == false
+        "#{value.to_s}"
+      elsif value.instance_of?(Integer) or value.instance_of?(Float)
+        "#{value.to_s}"
+      elsif value.instance_of?(String) or value.instance_of?(Pathname)
+        "#{s(value.to_s)}"
+      else
+        raise "invalid type for 'value' argument in JobExecution#convert_value: #{value} is #{value.class}"
+      end
+    end
+
+    def compile_where_expr(conds_hash)
+      conds_hash.map{|k,v| convert_cond(k,v) }.join(' and ')
+    end
+
+    def convert_cond(column, cond)
+      if cond.nil?
+        "#{column} is null"
+      elsif cond.instance_of?(Array) # not support subquery
+        in_clause = cond.map{|c| convert_cond(column, c)}.join(' or ')
+        "(#{in_clause})"
+      elsif cond == true or cond == false
+        "#{column} is #{cond.to_s}"
+      elsif cond.instance_of?(Integer) or cond.instance_of?(Float)
+        "#{column} = #{cond}"
+      elsif cond.instance_of?(String) or cond.instance_of?(Pathname)
+        "#{column} = #{s(cond.to_s)}"
+      else
+        raise "invalid type for 'cond' argument in JobExecution#convert_cond: #{cond} is #{cond.class}"
+      end
+    end
   end
 
 end

data/lib/bricolage/taskqueue.rb

@@ -1,11 +1,17 @@
 require 'bricolage/jobnet'
+require 'bricolage/sqlutils'
 require 'bricolage/exception'
+require 'bricolage/dao/job'
+require 'bricolage/dao/jobnet'
+require 'bricolage/dao/jobexecution'
 require 'fileutils'
 require 'pathname'
+require 'pg'
 
 module Bricolage
 
-  class TaskQueue
+  class MemoryTaskQueue
+
     def initialize
       @queue = []
     end
@@ -18,77 +24,90 @@ module Bricolage
       @queue.size
     end
 
-    def queued?
-      not empty?
-    end
-
     def each(&block)
       @queue.each(&block)
     end
 
     def consume_each
-      lock
-      save
-      while task = self.next
-        yield task
-        deq
+      while job = @queue.first
+        result = yield job
+        break unless result.success?
+        @queue.shift
       end
-    ensure
-      unlock
     end
 
-    def enq(task)
-      @queue.push task
+    def restore_jobnet(jobnet)
     end
 
-    def next
-      @queue.first
+    def enqueue_jobnet(jobnet)
+      jobnet.sequential_jobs.each do |job|
+        @queue.push job
+      end
     end
 
-    def deq
-      task = @queue.shift
-      save
-      task
+    def locked?(jobnet)
+      false
     end
 
-    def save
+    def unlock_help(jobnet)
+      raise "[BUG] this message must not be shown"
     end
 
-    def restore
+    def cancel_jobnet(jobnet, message)
+      @queue.clear
     end
 
-    def locked?
-      false
+  end
+
+
+  class FileTaskQueue
+
+    def initialize(path:)
+      @path = path
+      @queue = []
     end
 
-    def lock
+    def empty?
+      @queue.empty?
     end
 
-    def unlock
+    def size
+      @queue.size
     end
 
-    def unlock_help
-      "[MUST NOT HAPPEN] this message must not be shown"
+    def restore_jobnet(jobnet)
+      return unless File.exist?(@path)
+      File.foreach(@path) do |line|
+        @queue.push Task.deserialize(line)
+      end
     end
-  end
 
-  class FileTaskQueue < TaskQueue
-    def FileTaskQueue.restore_if_exist(path)
-      q = new(path)
-      q.restore if q.queued?
-      q
+    def enqueue_jobnet(jobnet)
+      jobnet.sequential_jobs.each do |ref|
+        @queue.push Task.new(ref)
+      end
     end
 
-    def initialize(path)
-      super()
-      @path = path
+    def each
+      @queue.each do |task|
+        yield task.job
+      end
     end
 
-    def queued?
-      @path.exist?
+    def consume_each
+      lock
+      save
+      while task = @queue.first
+        task_result = yield task.job
+        break unless task_result.success?
+        @queue.shift
+        save
+      end
+    ensure
+      unlock
     end
 
-    def save
+    private def save
       if empty?
         @path.unlink if @path.exist?
         return
@@ -97,7 +116,7 @@ module Bricolage
       tmpname = "#{@path}.tmp.#{Process.pid}"
       begin
         File.open(tmpname, 'w') {|f|
-          each do |task|
+          @queue.each do |task|
             f.puts task.serialize
           end
         }
@@ -107,48 +126,187 @@ module Bricolage
       end
     end
 
-    def restore
-      File.foreach(@path) do |line|
-        enq JobTask.deserialize(line)
-      end
-    end
-
-    def locked?
+    def locked?(jobnet)
       lock_file_path.exist?
     end
 
-    def lock
-      FileUtils.touch lock_file_path
+    private def lock
+      FileUtils.touch(lock_file_path)
     end
 
-    def unlock
-      FileUtils.rm_f lock_file_path
+    private def unlock
+      FileUtils.rm_f(lock_file_path)
     end
 
-    def lock_file_path
+    private def lock_file_path
       Pathname.new("#{@path}.LOCK")
     end
 
-    def unlock_help
+    def unlock_help(jobnet)
       "remove the file: #{lock_file_path}"
     end
+
+    def cancel_jobnet(jobnet, message)
+      unlock
+      FileUtils.rm_f(@path)
+      @queue.clear
+    end
+
+    class Task
+      def initialize(job)
+        @job = job
+      end
+
+      attr_reader :job
+
+      def serialize
+        [@job].join("\t")
+      end
+
+      def Task.deserialize(str)
+        job, * = str.strip.split("\t")
+        new(JobNet::Ref.parse(job))
+      end
+    end
+
   end
 
-  class JobTask
-    def initialize(job)
-      @job = job
+
+  class DatabaseTaskQueue
+
+    def initialize(datasource:, executor_id:, enable_lock: false)
+      @ds = datasource
+      @executor_id = executor_id
+      @enable_lock = enable_lock
+
+      @queue = []
+      @jobnet_dao = DAO::JobNet.new(@ds)
+      @job_dao = DAO::Job.new(@ds)
+      @jobexecution_dao = DAO::JobExecution.new(@ds)
+      @jobnet = nil
+    end
+
+    def empty?
+      @queue.empty?
+    end
+
+    def size
+      @queue.size
+    end
+
+    private def find_or_create_jobnet(ref)
+      @jobnet_rec ||= @jobnet_dao.find_or_create(ref)
+    end
+
+    def restore_jobnet(jobnet)
+      raise "jobnet is already bound to queue" if @jobnet
+
+      job_executions = @jobexecution_dao.enqueued_jobs(jobnet.ref)
+      unless job_executions.empty?
+        job_executions.each do |job_execution|
+          @queue.push Task.for_job_execution(job_execution)
+        end
+        @jobnet = jobnet
+      end
+    end
+
+    def enqueue_jobnet(jobnet)
+      raise "jobnet is already bound to queue" if @jobnet
+
+      jobnet_rec = find_or_create_jobnet(jobnet.ref)
+      jobnet.sequential_jobs.each_with_index do |job_ref, index|
+        job = @job_dao.find_or_create(jobnet_rec.id, job_ref)
+        job_execution = @jobexecution_dao.enqueue_job(job, index + 1)
+        @queue.push Task.for_job_execution(job_execution)
+      end
+      @jobnet = jobnet
+    end
+
+    def each
+      @queue.each do |task|
+        yield task.job
+      end
     end
 
-    attr_reader :job
+    def consume_each
+      raise "jobnet is not bound to queue" unless @jobnet
+
+      jobnet_rec = find_or_create_jobnet(@jobnet.ref)
+      @jobnet_dao.lock(jobnet_rec.id, @executor_id) if @enable_lock
+      while task = @queue.first
+        @job_dao.lock(task.job_id, @executor_id) if @enable_lock
+        begin
+          @jobexecution_dao.transition_to_running(task.job_execution_id)
+
+          # Note: fork(2) breaks current connections,
+          # we must close current connections before fork.
+          # (psql datasource forks process)
+          @ds.clear_connection_pool
+
+          job_completed = false
+          begin
+            task_result = yield task.job
 
-    def serialize
-      [@job].join("\t")
+            if task_result.success?
+              @jobexecution_dao.transition_to_succeeded(task.job_execution_id)
+              job_completed = true
+              @queue.shift
+            else
+              @jobexecution_dao.transition_to_failed(task.job_execution_id, task_result.message)
+              job_completed = true
+              break
+            end
+          ensure
+            unless job_completed
+              begin
+                @jobexecution_dao.transition_to_failed(task.job_execution_id, 'unexpected error')
+              rescue => ex
+                $stderr.puts "warning: could not write job state: #{ex.class}: #{ex.message} (this error is ignored)"
+              end
+            end
+          end
+        ensure
+          @job_dao.unlock(task.job_id, @executor_id) if @enable_lock
+        end
+      end
+    ensure
+      @jobnet_dao.unlock(jobnet_rec.id, @executor_id) if @enable_lock
     end
 
-    def JobTask.deserialize(str)
-      job, * = str.strip.split("\t")
-      new(JobNet::Ref.parse(job))
+    def locked?(jobnet)
+      @jobnet_dao.locked?(jobnet.ref)
     end
+
+    def unlock_help(jobnet)
+      jobnet_rec = find_or_create_jobnet(jobnet.ref)
+      locked_jobs = @job_dao.locked_jobs(jobnet_rec.id)
+      "clear executor_id of the jobnet (id: #{jobnet_rec.id}) and/or the jobs (id: #{locked_jobs.map(&:id).join(', ')})"
+    end
+
+    def cancel_jobnet(jobnet, message)
+      @jobexecution_dao.cancel_jobnet(jobnet.ref, message)
+      jobnet_rec = find_or_create_jobnet(jobnet.ref)
+      @jobnet_dao.clear_lock(jobnet_rec.id)
+      @job_dao.clear_lock_all(jobnet_rec.id)
+    end
+
+    class Task
+      def Task.for_job_execution(exec)
+        job_ref = JobNet::JobRef.new(exec.subsystem, exec.job_name, JobNet::Location.dummy)
+        new(job_ref, exec)
+      end
+
+      def initialize(job_ref, job_execution)
+        @job = job_ref
+        @job_id = job_execution.job_id
+        @job_execution_id = job_execution.job_execution_id
+      end
+
+      attr_reader :job
+      attr_reader :job_id
+      attr_reader :job_execution_id
+    end
+
   end
 
 end