datapipeline/lib/datapipeline.rb

require 'json'
require 'logger'

require_relative 'datapipeline/version'
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),
  # http (rest, slack)
  autoload :CsvStorage,      File.expand_path('datapipeline/storage/csv', __dir__)
  autoload :PostgresStorage, File.expand_path('datapipeline/storage/postgres', __dir__)
  autoload :JsonStorage,     File.expand_path('datapipeline/storage/json', __dir__)
  autoload :SqliteStorage,   File.expand_path('datapipeline/storage/sqlite', __dir__)
  autoload :ExcelStorage,    File.expand_path('datapipeline/storage/excel', __dir__)
  autoload :RestStorage,     File.expand_path('datapipeline/storage/rest', __dir__)
  autoload :FixedStorage,    File.expand_path('datapipeline/storage/fixed', __dir__)
  autoload :Slack,           File.expand_path('datapipeline/notification/slack', __dir__)

  # passed to header, trailer, map and filter
  #   read:    number of source rows (nil while streaming, until the trailer)
  #   written: number of body rows after all steps (nil while the body is processed)
  #   index:   index of the current row within the current step
  Context = Struct.new(:read, :written, :index, keyword_init: true)

  # result of a run, returned by run and passed to on_finish
  #   read: source rows, written: body rows, output: rows to the target incl. header/trailer,
  #   rejected: rows dropped by on_error :skip/:collect, errors: first 100 of them,
  #   warnings: rows that failed a validate step with warn: true (kept), error: exception of an aborted run
  Stats = Struct.new(:version, :read, :written, :output, :rejected, :errors, :warnings, :started_at, :finished_at,
                     :dry_run, :error, keyword_init: true) do
    def duration = finished_at && started_at ? (finished_at - started_at).round(3) : nil
    def ok? = error.nil? && rejected.zero?

    def to_s
      "read=#{read} written=#{written} rejected=#{rejected}#{warnings.to_i.positive? ? " warnings=#{warnings}" : ''} " \
        "#{duration}s#{dry_run ? ' (dry run)' : ''}" \
        "#{error ? " ABORTED: #{error.class}: #{error.message}" : ''}"
    end
  end

  # stops reading at the limit, not an error
  class LimitReached < StandardError; end

  EDGE_BEFORE = %i[prepend header].freeze
  EDGE_AFTER  = %i[append trailer].freeze
  CHAINING    = %i[map filter validate].freeze
  ON_ERROR    = %i[abort skip collect].freeze
  MAX_ERRORS  = 100

  # removed in 3.0, logged once per process as a warning
  DEPRECATIONS = {
    transform: 'transform is deprecated and will be removed in 3.0: use map/filter, ' \
               'header/trailer instead of processing: :prepend/:append',
    parallel: 'parallel mode (several transform steps without chained) is deprecated and will be removed in 3.0: ' \
              'set chained or use map',
    condition: 'condition: as string (eval) is deprecated and will be removed in 3.0: use a lambda ->(row) { ... }'
  }.freeze

  # value conversions for cast, empty values become nil
  CASTS = Validation::TYPES.merge(string: ->(v) { v.to_s }).freeze

  DEFAULT_SETTINGS = {
    version: nil,
    streamed: false,
    chained: false,
    verbose: false,
    dry_run: false,
    limit: nil,
    preview: 10,
    on_error: :abort,
    lock: nil,
    log_level: Logger::WARN
  }.freeze

  SETTINGS_LOCK = Mutex.new
  @defaults = DEFAULT_SETTINGS
  @plugins = {}
  @deprecations_logged = []

  #
  # defaults for every following DataPipeline.new, e.g. from bin/etl
  # the hash is replaced, never changed: a running pipeline keeps its own copy
  #
  def self.settings(options = {})
    SETTINGS_LOCK.synchronize { @defaults = @defaults.merge(options).freeze } if options.any?
    @defaults
  end

  def self.reset_settings!
    SETTINGS_LOCK.synchronize { @defaults = DEFAULT_SETTINGS }
  end

  #
  # container :name finds a registered plugin first, then <Name>Storage by convention
  # (DataPipeline::<Name>Storage of the gem, else a top level <Name>Storage of the application)
  #   DataPipeline.register :mem, MemStorage
  #   DataPipeline.register :sftp, 'SftpStorage', path: '/path/to/sftp_storage'   # loaded on first use
  #
  def self.register type, klass, path: nil
    Object.autoload(klass.to_sym, File.expand_path(path)) if path && klass.is_a?(String)
    SETTINGS_LOCK.synchronize { @plugins = @plugins.merge(type.to_sym => klass).freeze }
  end

  def self.plugin type
    klass = @plugins[type.to_sym]
    klass = Object.const_get(klass) if klass.is_a?(String)
    return klass if klass

    class_name = type.to_s.capitalize + 'Storage'
    return const_get(class_name, false) if const_defined?(class_name, false)

    Object.const_get(class_name) if Object.const_defined?(class_name)
  end

  #
  # parameters from the command line: bin/etl sets them; a script started with ruby directly
  # reads them from ARGV on first use: --name=value -> params[:name], --param key=value -> params[:key]
  #
  def self.params
    @params ||= parse_params(ARGV)
  end

  def self.params= hash
    @params = hash
  end

  def self.parse_params args
    params = {}
    args.each_with_index do |arg, i|
      case arg
      when /\A--param=(\w+)=(.*)\z/ then params[$1.to_sym] = $2
      when '--param' then key, value = args[i + 1].to_s.split('=', 2); params[key.to_sym] = value if value
      when /\A--([a-z][\w-]*)=(.*)\z/ then params[$1.tr('-', '_').to_sym] = $2
      end
    end
    params
  end

  # stats of every run in this process, used by bin/etl for the exit code
  def self.runs
    @runs ||= []
  end

  # true the first time a deprecation is reported in this process
  def self.deprecation_once key
    SETTINGS_LOCK.synchronize do
      next false if @deprecations_logged.include?(key)

      @deprecations_logged << key
      true
    end
  end

  def self.reset_deprecations! = SETTINGS_LOCK.synchronize { @deprecations_logged = [] }

  # DataPipeline.new(verbose: true) { ... } overrides the defaults for this pipeline only
  def initialize(**settings, &block)
    @settings = DataPipeline.settings.merge(settings)
    @transformations = []
    @targets = []
    @on_finish = []
    @logger = Logger.new($stderr, level: @settings[:log_level], progname: 'etl')
    @logger.formatter = proc { |severity, time, _prog, msg| "#{time.strftime('%F %T')} #{severity} #{msg}\n" }
    instance_eval(&block)
  end

  def params
    self.class.params
  end

  def configure(options = {})
    @settings.merge!(options)
  end

  def version value
    @settings[:version] = value
  end

  def verbose flag = true
    @settings[:verbose] = flag
  end

  def streamed flag = true
    @settings[:streamed] = flag
  end

  # each body step gets the output of the previous step instead of the source rows
  def chained flag = true
    @settings[:chained] = flag
  end

  # run everything, write nothing (target, rejects, watermark); prints the first rows
  def dry_run flag = true
    @settings[:dry_run] = flag
  end

  # for steps with side effects (api calls): true in a dry run
  def dry_run?
    @settings[:dry_run] ? true : false
  end

  # process only the first n source rows
  def limit count
    @settings[:limit] = count
  end

  def logger logger = nil
    @logger = logger if logger
    @logger
  end

  #
  # what happens when a body step raises (or a validate step fails) for a row
  #   on_error :abort                 -> stop the run (default)
  #   on_error :skip                  -> drop the row, log a warning, count it
  #   on_error :collect, to: storage  -> like skip, plus the row with _step and _error into storage
  #
  def on_error mode, to: nil
    raise ArgumentError, "on_error: #{mode} unknown, use #{ON_ERROR.join(', ')}" unless ON_ERROR.include?(mode)
    raise ArgumentError, 'on_error :collect needs to: storage' if mode == :collect && to.nil?

    @settings[:on_error] = mode
    @reject_target = to
  end

  # on_finish { |stats| ... } -> after every run, also after an abort (stats.error)
  def on_finish(&block)
    @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:)
  end

  # only a missing plugin is reported as unsupported, errors while loading it
  # (missing gem) or from the plugin itself (connection, options) are passed through
  def container type, **options, &block
    klass = DataPipeline.plugin(type)
    raise "Unsupported storage type or class missing: #{type.to_s.capitalize}Storage" unless klass

    klass.new(**options, &block)
  end

  #
  # transform(options) { |row, index| ... }
  #   processing: :prepend | :normal (default) | :append
  #   condition:  ruby expression as string (row is visible) or lambda { |row| ... }
  #
  def transform(options = {}, &block)
    kind = options.fetch(:processing, :normal)
    kind = :normal unless %i[prepend append].include?(kind)
    add_step kind, options, &block
  end

  # map { |row, ctx| ... }    -> hash, array of hashes or nil; turns chaining on
  def map(options = {}, &block)
    add_step :map, options, &block
  end

  # filter { |row, ctx| ... } -> keeps the row when the block is truthy; turns chaining on
  def filter(options = {}, &block)
    add_step :filter, options, &block
  end

  #
  # short steps, each one a map step (turns chaining on)
  #
  #   rename 'kunnr' => 'kundennr'               -> new key at the same position
  #   select %w[belnr betrag]                    -> only these keys in this order, missing ones nil
  #   cast 'betrag' => :decimal                  -> :integer, :decimal (Float), :date, :string; empty -> nil
  #   default 'land' => 'CH'                     -> for nil or blank values; a lambda gets the row
  #
  def rename mapping
    mapping = mapping.to_h { |from, to| [from.to_s, to.to_s] }
    add_step(:map, { name: 'rename' }) { |row| row.to_h { |key, value| [mapping.fetch(key.to_s, key), value] } }
  end

  def select *keys
    keys = keys.flatten.map(&:to_s)
    add_step(:map, { name: 'select' }) { |row| keys.to_h { |key| [key, value_of(row, key)] } }
  end

  def cast types
    types = types.to_h { |key, type| [key.to_s, type.to_sym] }
    unknown = types.values - CASTS.keys
    raise ArgumentError, "cast: #{unknown.join(', ')} unknown, use #{CASTS.keys.join(', ')}" if unknown.any?

    add_step(:map, { name: 'cast' }) do |row|
      types.each_with_object(row.dup) do |(key, type), out|
        value = value_of(row, key)
        out[key] = blank?(value) ? nil : cast_value(key, value, type)
      end
    end
  end

  def default values
    values = values.to_h { |key, value| [key.to_s, value] }
    add_step(:map, { name: 'default' }) do |row|
      values.each_with_object(row.dup) do |(key, value), out|
        out[key] = value.respond_to?(:call) ? value.call(row) : value if blank?(value_of(row, key))
      end
    end
  end

  # header { |ctx| ... } -> rows in front of the body, called after the body with full counters
  def header(&block)
    add_step :header, {}, &block
  end

  # trailer { |ctx| ... } -> rows after the body
  def trailer(&block)
    add_step :trailer, {}, &block
  end

  #
  # validate 'kostl' => { required: true, max: 10 }, 'shkzg' => { in: %w[S H] }
  # validate layout                                   -> fields of a DataPipeline::Layout
  # validate { |row| 'soll ungleich haben' if ... }   -> nil, message or array of messages
  #
  # a failing row raises ValidationError, handled by on_error; turns chaining on
  # warn: true      -> only a warning in the log and stats.warnings, the row is kept
  # ignore: %w[A B] -> messages for these fields are left out (known gaps)
  #
  def validate(rules = nil, **options, &block)
    warn = options.delete(:warn) || false
    ignore = Array(options.delete(:ignore)).map(&:to_s)
    rules ||= options # validate 'konto' => { ... } arrives as keywords
    @transformations << { kind: :validate, block:, options: { rules:, warn:, ignore: } }
  end

  def read storage, options = {}
    @source = storage
    @source.set options if options.any?
  end

  #
  # write storage                                   -> all rows
  # write storage, only: ->(row) { ... }            -> only matching rows (also header/trailer rows)
  #
  # can be called several times, every target gets its rows
  #
  def write storage, only: nil
    @targets << { storage:, only: }
  end

  def query storage, options = {}
    storage.query options
  end

  #
  # loads a table once, for lookups in map/transform instead of one query per row
  #
  #   lookup pg, query: 'select ...', params: [...], key: 'rechnung_dbid'  -> { key => [rows] }
  #   lookup pg, query: 'select ...', key: 'id', unique: true              -> { key => row }
  #   lookup csv, key: ->(row) { row['land'] }                             -> reads the whole storage
  #
  def lookup storage, key:, unique: false, **options
    rows = options[:query] ? storage.query(options) : storage.read
    key_of = key.respond_to?(:call) ? key : ->(row) { row[key] }
    unique ? rows.to_h { |row| [key_of.call(row), row] } : rows.group_by(&key_of)
  end

  def run
    @stats = Stats.new(version: @settings[:version], read: 0, written: 0, output: 0, rejected: 0, errors: [], warnings: 0,
                       started_at: Time.now, dry_run: @settings[:dry_run])
    @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]) }
    after  = @transformations.select { |t| EDGE_AFTER.include?(t[:kind]) }
    body   = @transformations - before - after
    log_deprecations body

    @settings[:streamed] ? run_streamed(before, body, after) : run_batch(before, body, after)

    @watermark&.save unless @settings[:dry_run]
    @stats
  rescue StandardError => e
    @stats.error = e
    @logger.error "abort: #{e.class}: #{e.message}"
    raise
  ensure
    finish
  end

  private

  def add_step kind, options, &block
    raise ArgumentError, "#{kind}: block missing" unless block

    @transformations << { kind:, block:, options: }
  end

  def value_of row, key
    row.key?(key) ? row[key] : row[key.to_sym]
  end

  def blank? value
    value.nil? || (value.is_a?(String) && value.strip.empty?)
  end

  def cast_value key, value, type
    CASTS.fetch(type).call(value)
  rescue ArgumentError, TypeError, Date::Error
    raise ArgumentError, "#{key}: #{value.inspect} is not #{type}"
  end

  # each deprecated feature once per process, with the logger of the pipeline
  def log_deprecations body
    used = []
    used << :transform if @transformations.any? { |t| %i[prepend normal append].include?(t[:kind]) }
    used << :parallel if !chained? && body.size > 1
    used << :condition if @transformations.any? { |t| t[:options].is_a?(Hash) && t[:options][:condition].is_a?(String) }
    used.each { |key| @logger.warn "deprecated: #{DEPRECATIONS[key]}" if DataPipeline.deprecation_once(key) }
  end

  # without any transformation the source is copied unchanged
  def chained?
    @settings[:chained] || @transformations.empty? ||
      @transformations.any? { |t| CHAINING.include?(t[:kind]) }
  end

  def write_row row
    @targets.each { |t| t[:storage].write_row(row) if t[:only].nil? || t[:only].call(row) }
  end

  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}"
    @on_finish.each do |hook|
      hook.call(@stats)
    rescue StandardError => e
      @logger.error "on_finish: #{e.class}: #{e.message}"
    end
  end

  def source_rows
    data = @source.read
    data = data.first(@settings[:limit]) if @settings[:limit]
    data
  end

  def output rows
    @stats.output += rows.size
    @preview.concat(rows.first(@settings[:preview] - @preview.size)) if @preview.size < @settings[:preview]
    return if @settings[:dry_run]

    yield rows
  end

  def print_preview
    puts "DRY RUN – nothing written: #{@stats}"
    @preview.each { |row| puts JSON.generate(row) }
  end

  # whole source in memory, one target.write at the end
  def run_batch before, body, after
    data = source_rows
    @stats.read = data.length
    data.each { |row| @watermark.track(row) } if @watermark
    dump 'Source#data', data

    # prepend keeps its old timing: called before the body
    results = {}.compare_by_identity
    before.select { |t| t[:kind] == :prepend }.each { |t| results[t] = edge(t, nil, 0) }

    rows = []
    if chained?
      counters = Array.new(body.size, 0)
      data.each { |row| push(row, body, counters, data.length) { |out| rows << out } }
    else
      body.each { |step| rows.concat apply(step, data) }
    end
    dump 'Transformation#data', rows

    ctx = Context.new(read: data.length, written: rows.length)
    before.select { |t| t[:kind] == :header }.each { |t| results[t] = edge(t, ctx, 0) }

    # parallel mode keeps the old count of source rows for append
    count = chained? ? rows.length : data.length
    new_data = []
    before.each { |t| add_rows new_data, results[t] }
    new_data.concat rows
    after.each { |t| add_rows new_data, edge(t, ctx, count) }

    @stats.written = rows.length
    dump 'Target#data', new_data
    output(new_data) do |all|
      @targets.each { |t| t[:storage].write(t[:only] ? all.select(&t[:only]) : all) }
    end
  end

  # row by row from source.read { } to target.write_row, header only knows nothing yet
  def run_streamed before, body, after
    if !chained? && body.size > 1
      raise ArgumentError, 'streamed: several transformations need chained (or map/filter)'
    end

    emit = ->(result) { output(add_rows([], result)) { |rows| rows.each { |row| write_row(row) } } }

    before.each { |t| emit.call edge(t, Context.new, 0) }

    read = written = 0
    counters = Array.new(body.size, 0)
    begin
      @source.read do |row|
        raise LimitReached if @settings[:limit] && read >= @settings[:limit]

        read += 1
        @watermark&.track(row)
        push(row, body, counters, nil) do |out|
          written += 1
          emit.call out
        end
      end
    rescue LimitReached
      @logger.info "limit #{@settings[:limit]} reached"
    end
    @stats.read = read
    @stats.written = written

    ctx = Context.new(read:, written:)
    after.each { |t| emit.call edge(t, ctx, chained? ? written : read) }
    @targets.each { |t| t[:storage].finalize }
  end

  # one row through the chain, depth first; counters hold the index per step
  def push row, body, counters, read, level = 0, &emit
    return emit.call(row) if level == body.size

    index = counters[level]
    counters[level] += 1
    safe_step_rows(body[level], row, index, read).each do |out|
      push(out, body, counters, read, level + 1, &emit)
    end
  end

  # parallel mode: every step reads its own copy of the source rows
  def apply step, rows
    rows.each_with_index.flat_map { |row, index| safe_step_rows(step, row.dup, index, rows.length) }
  end

  # errors of a step for one row go to on_error
  def safe_step_rows step, row, index, read
    step_rows(step, row, index, read)
  rescue StandardError => e
    reject(step, row, e)
    []
  end

  def reject step, row, error
    raise error if @settings[:on_error] == :abort

    name = step_name(step)
    @stats.rejected += 1
    @stats.errors << { step: name, error: error.message, row: } if @stats.errors.size < MAX_ERRORS
    @logger.warn "rejected by #{name}: #{error.message}"
    return if @settings[:dry_run] || @reject_target.nil?

    @reject_target.write_row(row.merge('_step' => name, '_error' => error.message))
  end

  def step_name step
    name = step[:options].is_a?(Hash) && step[:options][:name] || step[:kind]
    "#{name}##{@transformations.index(step) + 1}"
  end

  # rows produced by one step for one input row
  def step_rows step, row, index, read
    return [] unless condition?(step, row)

    case step[:kind]
    when :validate
      rules = step[:options][:rules]
      messages = rules.is_a?(Layout) ? rules.errors(row) : Validation.errors(row, rules)
      messages += Array(step[:block].call(row)).compact if step[:block]
      ignore = step[:options][:ignore]
      messages = messages.reject { |m| ignore.any? { |field| m.start_with?("#{field}:") } } if ignore.any?
      return [row] if messages.empty?
      raise ValidationError, messages.join('; ') unless step[:options][:warn]

      @stats.warnings += 1
      if @stats.warnings <= MAX_ERRORS
        @logger.warn "#{step_name(step)}: #{messages.join('; ')}"
      elsif @stats.warnings == MAX_ERRORS + 1
        @logger.warn "more than #{MAX_ERRORS} warnings, further ones only counted"
      end
      [row]
    when :filter
      step[:block].call(row, Context.new(read:, index:)) ? [row] : []
    when :map
      add_rows [], step[:block].call(row, Context.new(read:, index:))
    else
      add_rows [], step[:block].call(row, index)
    end
  end

  def edge step, ctx, count
    case step[:kind]
    when :prepend then step[:block].call(nil, 0)
    when :append  then step[:block].call(nil, count)
    else step[:block].call(ctx)
    end
  end

  # string conditions are evaluated with the local variable row
  def condition? step, row
    condition = step[:options].fetch(:condition, true)
    return true if condition == true
    return condition.call(row) if condition.respond_to?(:call)

    eval(condition)
  end

  # a hash is one row, an array are several rows, nil is no row
  def add_rows new_data, result
    if result.is_a?(Array)
      new_data.concat(result.compact)
    elsif result
      new_data << result
    end
    new_data
  end

  def dump title, data
    return unless @settings[:verbose]

    puts "#{title}".ljust(28, '-')
    puts JSON.pretty_generate data
    puts "----------------------------"
  end
end

# short names besides the <Name>Storage convention
DataPipeline.register :sqlite3, 'DataPipeline::SqliteStorage'
DataPipeline.register :xlsx, 'DataPipeline::ExcelStorage'