Betriebssicherheit: Sperre gegen parallele Läufe, Upsert, Batch-Schreiben für Postgres, Zeichensatz für CSV

lib/datapipeline.rb

@@ -6,6 +6,7 @@ require_relative 'datapipeline/storage/storage_plugin'
 require_relative 'datapipeline/validation'
 require_relative 'datapipeline/layout'
 require_relative 'datapipeline/watermark'
+require_relative 'datapipeline/lock'
 
 class DataPipeline
   # plugins are loaded on first use, their gems too: pg (postgres), sqlite3 (sqlite), roo/caxlsx (excel),
@@ -58,6 +59,7 @@ class DataPipeline
     limit: nil,
     preview: 10,
     on_error: :abort,
+    lock: nil,
     log_level: Logger::WARN
   }.freeze
 
@@ -204,6 +206,11 @@ class DataPipeline
     @on_finish << block
   end
 
+  # lock 'data/.cudo.lock' -> held during run; a second run at the same time aborts with LockError
+  def lock file
+    @settings[:lock] = file
+  end
+
   # see DataPipeline::Watermark
   def watermark name, column:, file:, initial: nil
     @watermark = Watermark.new(name, column:, file:, initial:)
@@ -303,6 +310,7 @@ class DataPipeline
     @logger.info "start#{" #{@settings[:version]}" if @settings[:version]}#{@settings[:dry_run] ? ' (dry run)' : ''}"
     @preview = []
     raise ArgumentError, 'read storage missing' unless @source
+    (@lock = Lock.new(@settings[:lock])).acquire if @settings[:lock]
     raise ArgumentError, 'write storage missing' if @targets.empty?
 
     before = @transformations.select { |t| EDGE_BEFORE.include?(t[:kind]) }
@@ -342,6 +350,7 @@ class DataPipeline
   def finish
     @stats.finished_at = Time.now
     [@source, *@targets.map { |t| t[:storage] }, @reject_target].compact.uniq.each(&:close)
+    @lock&.release
     self.class.runs << @stats
     print_preview if @stats.dry_run && @stats.error.nil?
     @logger.info "finish #{@stats}"

lib/datapipeline/lock.rb

@@ -0,0 +1,57 @@
+require 'fileutils'
+
+class DataPipeline
+  # another run holds the lock
+  class LockError < StandardError; end
+
+  #
+  # exclusive lock on a file, against overlapping runs (cron)
+  #
+  #   lock 'data/.cudo.lock'            -> in a pipeline: held from the start to the end of run
+  #   etl run cudo.rb --lock=data/.lock -> for the whole script
+  #
+  # the lock is an flock on the file, released by the operating system when the process ends,
+  # also after a crash; the file stays and holds the pid of the last holder
+  #
+  class Lock
+    attr_reader :file
+
+    def initialize file
+      @file = file.to_s
+    end
+
+    def acquire
+      return self if @io
+
+      FileUtils.mkdir_p(File.dirname(@file))
+      io = File.open(@file, File::RDWR | File::CREAT, 0o644)
+      unless io.flock(File::LOCK_EX | File::LOCK_NB)
+        holder = io.read.strip
+        io.close
+        raise LockError, "#{@file} is locked#{" by pid #{holder}" unless holder.empty?}: another run is still active"
+      end
+      io.truncate(0)
+      io.write(Process.pid.to_s)
+      io.flush
+      @io = io
+      self
+    end
+
+    def release
+      return unless @io
+
+      @io.flock(File::LOCK_UN)
+      @io.close
+      @io = nil
+    end
+
+    def locked? = !@io.nil?
+
+    def synchronize
+      acquire
+      yield
+    ensure
+      release
+    end
+  end
+end

lib/datapipeline/storage/csv.rb

@@ -11,6 +11,10 @@ class DataPipeline
   #             without columns the values are written in hash order
   #   layout:   DataPipeline::Layout, field order per record type (several record types in one file),
   #             unknown fields raise; no header line
+  #   encoding: encoding of the file, e.g. 'ISO-8859-1' or 'Windows-1252' (default UTF-8);
+  #             read as UTF-8 strings, written transcoded from UTF-8; a UTF-8 BOM is skipped when reading
+  #   replace:  '?' -> characters the encoding cannot hold are replaced when writing (default: error)
+  #   bom:      true -> write a UTF-8 BOM (Excel)
   #   counter:  print a dot per written row
   #   verbose:  print settings
   #
@@ -44,11 +48,11 @@ class DataPipeline
 
     def read(&block)
       if block_given?
-        CSV.foreach(@file_name, **csv_options) { |row| yield to_hash(row) }
+        CSV.foreach(@file_name, **read_options) { |row| yield to_hash(row) }
         return
       end
 
-      CSV.read(@file_name, **csv_options).map { |row| to_hash(row) }
+      CSV.read(@file_name, **read_options).map { |row| to_hash(row) }
     end
 
     def write data
@@ -93,10 +97,30 @@ class DataPipeline
       @options[:columns]
     end
 
+    def utf8?
+      @options[:encoding].nil? || Encoding.find(@options[:encoding]) == Encoding::UTF_8
+    end
+
+    # csv: { encoding: ... } is passed unchanged (older scripts)
+    def read_options
+      return csv_options if csv_options.key?(:encoding)
+
+      csv_options.merge(encoding: utf8? ? 'bom|utf-8' : "#{@options[:encoding]}:utf-8")
+    end
+
+    def write_options
+      return csv_options if csv_options.key?(:encoding) || utf8?
+
+      options = csv_options.merge(encoding: @options[:encoding])
+      options.merge!(undef: :replace, invalid: :replace, replace: @options[:replace]) if @options[:replace]
+      options
+    end
+
     def open_file
       return if @file
 
-      @file = CSV.open(@file_name, 'w', **csv_options)
+      @file = CSV.open(@file_name, 'w', **write_options)
+      @file.to_io.write("\uFEFF") if @options[:bom] && utf8?
       @header_written = false
       return unless columns && csv_options[:headers] && !@options[:layout]
 

lib/datapipeline/storage/postgres.rb

@@ -11,6 +11,9 @@ class DataPipeline
   #   typed:   true -> values come back as ruby types (Integer, Float, Date, ...)
   #            false (default) -> all values as String, NULL as nil
   #   fetch_size: rows per fetch when reading with a block (default 1000)
+  #   upsert:  %w[belnr] -> insert ... on conflict (belnr) do update; needs a unique index on these columns
+  #   batch_size: write_row collects rows and writes them in one transaction per batch (default 1000);
+  #            :all -> one transaction for the whole stream, committed by finalize, rolled back by close
   #
   class PostgresStorage < StoragePlugin
     def initialize **options, &block
@@ -65,25 +68,88 @@ class DataPipeline
       table_name, data = data.nil? ? [@options[:table], table_or_data] : [table_or_data, data]
       raise "Table for PostgresStorage#write missing" unless table_name
 
-      @conn.transaction do
+      transaction do
         truncate_once table_name
         data.each { |row| insert table_name, row }
       end
     end
 
+    # rows are collected and written per batch_size in one transaction
     def write_row row
       raise "Table for PostgresStorage#write_row missing" unless @options[:table]
 
-      truncate_once @options[:table]
-      insert @options[:table], row
+      if batch_size == :all
+        begin_stream
+        truncate_once @options[:table]
+        insert @options[:table], row
+      else
+        (@pending ||= []) << row
+        flush if @pending.size >= batch_size
+      end
     end
 
+    # end of a stream: remaining rows, commit of batch_size: :all
+    def finalize
+      flush
+      if @stream_open
+        @conn.exec('commit')
+        @stream_open = false
+      end
+      close
+    end
+
+    # pending rows are written; an open stream with batch_size: :all is rolled back (run aborted)
+    # an error while flushing does not hide the error that aborted the run
     def close
-      @conn.close if @conn && !@conn.finished?
+      return unless @conn && !@conn.finished?
+
+      aborted_by = $!
+      begin
+        flush
+      rescue StandardError
+        raise unless aborted_by
+      end
+      @conn.exec('rollback') if @stream_open
+      @stream_open = false
+      @conn.close
     end
 
     private
 
+    def batch_size
+      size = @options.fetch(:batch_size, 1000)
+      size == :all ? :all : Integer(size)
+    end
+
+    def flush
+      return if @pending.nil? || @pending.empty?
+
+      rows = @pending
+      @pending = []
+      transaction do
+        truncate_once @options[:table]
+        rows.each { |row| insert @options[:table], row }
+      end
+    end
+
+    def begin_stream
+      return if @stream_open
+
+      raise 'PostgresStorage: batch_size: :all needs its own connection (no open transaction)' unless idle?
+
+      @conn.exec('begin')
+      @stream_open = true
+    end
+
+    # inside a transaction (e.g. a cursor of the same connection) no own transaction is opened
+    def transaction(&block)
+      idle? ? @conn.transaction(&block) : yield
+    end
+
+    def idle?
+      @conn.transaction_status == PG::PQTRANS_IDLE
+    end
+
     # rows are fetched in batches inside a transaction; other queries on the
     # same connection (lookups, subqueries) are possible while reading
     def each_row sql, params
@@ -112,11 +178,24 @@ class DataPipeline
       columns = row.keys.map { |col| @conn.quote_ident(col.to_s) }.join(',')
       placeholders = (1..row.size).map { |i| "$#{i}" }.join(',')
       @conn.exec_params(
-        "insert into #{quote_table(table_name)} (#{columns}) values (#{placeholders})",
+        "insert into #{quote_table(table_name)} (#{columns}) values (#{placeholders})#{upsert_clause(row)}",
         row.values
       )
     end
 
+    # on conflict (keys) do update set <other columns> = excluded.<column>
+    def upsert_clause row
+      keys = Array(@options[:upsert]).map(&:to_s)
+      return '' if keys.empty?
+
+      others = row.keys.map(&:to_s) - keys
+      target = keys.map { |key| @conn.quote_ident(key) }.join(',')
+      return " on conflict (#{target}) do nothing" if others.empty?
+
+      updates = others.map { |col| "#{@conn.quote_ident(col)} = excluded.#{@conn.quote_ident(col)}" }.join(', ')
+      " on conflict (#{target}) do update set #{updates}"
+    end
+
     # schema.table -> "schema"."table"
     def quote_table name
       name.to_s.split('.').map { |part| @conn.quote_ident(part) }.join('.')

lib/datapipeline/storage/sqlite.rb

@@ -10,6 +10,8 @@ class DataPipeline
   #   create:   true -> create the table on the first write, columns from the first row
   #             (Integer -> integer, Float -> real, else text)
   #   truncate: true -> delete all rows of the table before the first write
+  #   upsert:   %w[belnr] -> insert ... on conflict (belnr) do update; with create: true the unique index
+  #             is created too, otherwise the table needs one on these columns
   #
   # rows come back as Hash<String, value> with sqlite types (Integer, Float, String, nil).
   # the database file is opened on first use: a dry run creates no file.
@@ -86,6 +88,7 @@ class DataPipeline
       if @options[:create]
         columns = row.map { |key, value| "#{quote(key)} #{column_type(value)}" }.join(', ')
         db.execute("create table if not exists #{quote_table(table_name)} (#{columns})")
+        create_upsert_index table_name if upsert_keys.any?
       end
       db.execute("delete from #{quote_table(table_name)}") if @options[:truncate]
       @prepared[table_name] = true
@@ -108,7 +111,29 @@ class DataPipeline
       columns = row.keys.map { |col| quote(col) }.join(',')
       placeholders = (['?'] * row.size).join(',')
       values = row.values.map { |v| v.is_a?(Date) || v.is_a?(Time) ? v.iso8601 : v }
-      db.execute("insert into #{quote_table(table_name)} (#{columns}) values (#{placeholders})", values)
+      db.execute("insert into #{quote_table(table_name)} (#{columns}) values (#{placeholders})#{upsert_clause(row)}", values)
+    end
+
+    def upsert_keys
+      Array(@options[:upsert]).map(&:to_s)
+    end
+
+    # on conflict (keys) do update set <other columns> = excluded.<column>
+    def upsert_clause row
+      return '' if upsert_keys.empty?
+
+      others = row.keys.map(&:to_s) - upsert_keys
+      target = upsert_keys.map { |key| quote(key) }.join(',')
+      return " on conflict (#{target}) do nothing" if others.empty?
+
+      " on conflict (#{target}) do update set #{others.map { |col| "#{quote(col)} = excluded.#{quote(col)}" }.join(', ')}"
+    end
+
+    # schema.table -> index schema.table_upsert on table (keys)
+    def create_upsert_index table_name
+      *schema, table = table_name.to_s.split('.')
+      index = [*schema, "#{table}_upsert"].map { |part| quote(part) }.join('.')
+      db.execute("create unique index if not exists #{index} on #{quote(table)} (#{upsert_keys.map { |k| quote(k) }.join(',')})")
     end
 
     def quote name

lib/datapipeline/version.rb

@@ -1,5 +1,5 @@
 class DataPipeline
-  VERSION = '2.0.0'.freeze
+  VERSION = '2.1.0'.freeze
 
   # kept for scripts that print the version (rake info:info); the history is in CHANGELOG.md
   class Version