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'
require_relative 'datapipeline/connections'

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
  #   totals: { 'column' => { read: Rational, written: Rational } } of the columns named by totals
  Stats = Struct.new(:name, :version, :read, :written, :output, :rejected, :errors, :warnings, :started_at,
                     :finished_at, :dry_run, :error, :totals, keyword_init: true) do
    def duration = finished_at && started_at ? (finished_at - started_at).round(3) : nil
    def ok? = error.nil? && rejected.zero?

    # stats.total('dmbtr') -> sum after all steps, stats.total('dmbtr', :read) -> sum of the source rows
    def total column, side = :written
      totals.to_h.dig(column.to_s, side)
    end

    def to_s
      "read=#{read} written=#{written} rejected=#{rejected}#{warnings.to_i.positive? ? " warnings=#{warnings}" : ''}" \
        "#{totals.to_h.map { |col, sum| " #{col}=#{DataPipeline.decimal(sum[:read])}/#{DataPipeline.decimal(sum[:written])}" }.join} " \
        "#{duration}s#{dry_run ? ' (dry run)' : ''}" \
        "#{error ? " ABORTED: #{error.class}: #{error.message}" : ''}"
    end

    # one line of the run history
    def to_h
      { name:, version:, started_at: started_at&.iso8601(3), finished_at: finished_at&.iso8601(3), duration:,
        read:, written:, output:, rejected:, warnings:, dry_run: dry_run ? true : false, ok: ok?,
        error: error && "#{error.class}: #{error.message}",
        totals: totals.to_h.transform_values { |sum| sum.transform_values { |v| DataPipeline.decimal(v) } } }
    end
  end

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

  # raised when a reconcile check fails
  class ReconcileError < StandardError; end

  ON_ERROR    = %i[abort skip collect].freeze
  MAX_ERRORS  = 100

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

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

  SETTINGS_LOCK = Mutex.new
  @defaults = DEFAULT_SETTINGS
  @plugins = {}

  #
  # 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

  # see DataPipeline::Connections
  def self.connections_file = Connections.file
  def self.connections_file=(path)
    Connections.file = path
  end

  # entries of a run history file (see history), oldest first; name: only the runs of this pipeline
  def self.history file, name: nil
    return [] unless File.exist?(file)

    entries = File.foreach(file).filter_map { |line| JSON.parse(line) unless line.strip.empty? }
    name ? entries.select { |entry| entry['name'] == name.to_s } : entries
  end

  # exact decimal text of a sum: 1234.5 -> '1234.5', 1/3 -> '0.333333'
  def self.decimal value
    return value.to_s unless value.is_a?(Rational)
    return value.to_i.to_s if value.denominator == 1

    places = (1..10).find { |n| (value * 10**n).denominator == 1 } || 6
    format("%.#{places}f", value)
  end

  # DataPipeline.new(verbose: true) { ... } overrides the defaults for this pipeline only
  def initialize(**settings, &block)
    @settings = DataPipeline.settings.merge(settings)
    @steps = []
    @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

  # 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

  # name of the pipeline in the stats and the run history; etl sets the script name as default
  def name value
    @settings[:name] = value
  end

  # history 'data/runs.ndjson' -> one json line per run (also aborted and dry runs), see DataPipeline.history
  def history file
    @settings[:history] = file
  end

  #
  # totals 'dmbtr', 'menge'   -> sums of these columns over the source rows and over the rows after all steps
  #                              (stats.totals, stats.total('dmbtr')); empty and non numeric values are left out
  #
  def totals *columns
    @totals = columns.flatten.map(&:to_s)
  end

  #
  # reconcile { |stats| 'Summe weicht ab' if stats.total('dmbtr', :read) != stats.total('dmbtr') }
  #
  # runs after all steps, before writing (streaming: after the last row); a message (or array of messages)
  # aborts the run with ReconcileError, nil means ok
  #
  def reconcile(&block)
    raise ArgumentError, 'reconcile: block missing' unless block

    (@reconcile ||= []) << block
  end

  # connection(:pds, query: '...') -> container of the type in the connections file, options of the call win
  def connection name, **options, &block
    settings = Connections.fetch(name)
    type = settings.fetch(:type) { raise ArgumentError, "connection #{name.inspect}: type missing" }
    container(type, **settings.except(:type), **options, &block)
  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
  #
  # container :postgres, connection: :pds, table: 'x' -> options from the connections file, the call wins
  #
  def container type, **options, &block
    klass = DataPipeline.plugin(type)
    raise "Unsupported storage type or class missing: #{type.to_s.capitalize}Storage" unless klass

    if (name = options.delete(:connection))
      settings = Connections.fetch(name)
      if settings[:type] && settings[:type].to_s != type.to_s
        raise ArgumentError, "connection #{name.inspect} is #{settings[:type]}, not #{type}"
      end

      options = settings.except(:type).merge(options)
    end
    klass.new(**options, &block)
  end

  # map { |row, ctx| ... }    -> hash, array of hashes or nil
  # condition: ->(row) { ... } on map and filter: the step only for matching rows, the others are dropped
  def map(options = {}, &block)
    add_step :map, options, &block
  end

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

  #
  # short steps, each one a map step
  #
  #   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
  # 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
    @steps << { 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 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(name: @settings[:name], version: @settings[:version], read: 0, written: 0, output: 0, rejected: 0,
                       errors: [], warnings: 0, started_at: Time.now, dry_run: @settings[:dry_run],
                       totals: @totals&.to_h { |col| [col, { read: 0r, written: 0r }] })
    @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 = @steps.select { |t| t[:kind] == :header }
    after  = @steps.select { |t| t[:kind] == :trailer }
    body   = @steps - before - after

    @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
    if options.key?(:condition) && !options[:condition].respond_to?(:call)
      raise ArgumentError, "#{kind}: condition must be a lambda ->(row) { ... }"
    end

    @steps << { 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

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

  # one json line per run, appended under an flock (parallel runs of other pipelines)
  def append_history
    return unless @settings[:history]

    FileUtils.mkdir_p(File.dirname(@settings[:history]))
    File.open(@settings[:history], 'a') do |file|
      file.flock(File::LOCK_EX)
      file.puts JSON.generate(@stats.to_h)
    end
  rescue StandardError => e
    @logger.error "history: #{e.class}: #{e.message}"
  end

  def add_total side, row
    @stats.totals&.each do |col, sum|
      value = row.key?(col) ? row[col] : row[col.to_sym]
      number = to_number(value)
      sum[side] += number if number
    end
  end

  def to_number value
    case value
    when Float then value.finite? ? value.to_s.to_r : nil
    when Numeric then value.to_r
    when String then value.strip.match?(Layout::NUMBER) ? value.strip.to_r : nil
    end
  end

  # messages of the reconcile blocks abort the run
  def check_reconcile
    messages = Array(@reconcile).flat_map { |block| Array(block.call(@stats)).compact }
    raise ReconcileError, messages.join('; ') if messages.any?
  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
    data.each { |row| add_total(:read, row) } if @totals
    dump 'Source#data', data

    rows = []
    counters = Array.new(body.size, 0)
    data.each { |row| push(row, body, counters, data.length) { |out| rows << out } }
    dump 'Transformation#data', rows
    @stats.written = rows.length
    rows.each { |row| add_total(:written, row) } if @totals
    check_reconcile

    # header is computed after the body, so it knows the counters
    ctx = Context.new(read: data.length, written: rows.length)
    new_data = []
    before.each { |t| add_rows new_data, t[:block].call(ctx) }
    new_data.concat rows
    after.each { |t| add_rows new_data, t[:block].call(ctx) }

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

    before.each { |t| emit.call t[:block].call(Context.new) }

    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)
        add_total(:read, row) if @totals
        push(row, body, counters, nil) do |out|
          written += 1
          add_total(:written, out) if @totals
          emit.call out
        end
      end
    rescue LimitReached
      @logger.info "limit #{@settings[:limit]} reached"
    end
    @stats.read = read
    @stats.written = written
    check_reconcile

    ctx = Context.new(read:, written:)
    after.each { |t| emit.call t[:block].call(ctx) }
    @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

  # 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}##{@steps.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:))
    end
  end

  # condition: ->(row) { ... } on map and filter
  def condition? step, row
    condition = step[:options][:condition]
    condition.nil? || condition.call(row)
  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'