datapipeline/lib/datapipeline/storage/postgres.rb

require_relative 'storage_plugin'

class DataPipeline
  #
  # options
  #   host, port, database, username, password
  #   query:   sql for read, placeholders $1, $2, ...
  #   params:  values for the placeholders of query
  #   table:   target table for write
  #   truncate: true -> delete all rows of the table before the first write (same transaction)
  #   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
      self.class.require_gem 'pg'
      @options = options

      begin
        @conn = PG.connect(
          host: options[:host],
          port: options[:port] || '5432',
          dbname: options[:database],
          user: options[:username],
          password: options[:password]
        )
      rescue PG::ConnectionBad => e
        raise PG::ConnectionBad, "Error connection to database: #{e.message}"
      end

      @conn.type_map_for_queries = PG::BasicTypeMapForQueries.new(@conn)
      @conn.type_map_for_results = PG::BasicTypeMapForResults.new(@conn) if options[:typed]

      yield if block_given?
    end

    # without block all rows at once, with block row by row through a server side cursor
    def read(&block)
      raise "Query for PostgresStorage missing" unless @options[:query]

      params = @options.fetch(:params, [])
      return query(query: @options[:query], params:) unless block_given?

      each_row(@options[:query], params, &block)
    end

    def query options = {}
      raise "Query for PostgresStorage:query missing" unless options[:query]

      result = @conn.exec_params(options[:query], options.fetch(:params, []))
      return result.to_a unless block_given?

      result.each { |row| yield row }
      nil
    end

    #
    # write rows             -> table from option table:
    # write table_name, rows -> explicit table (old signature)
    #
    # all rows in one transaction, values as bind parameters, nil -> NULL
    #
    def write table_or_data, data = nil
      table_name, data = data.nil? ? [@options[:table], table_or_data] : [table_or_data, data]
      raise "Table for PostgresStorage#write missing" unless table_name

      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]

      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
      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
      size = Integer(@options.fetch(:fetch_size, 1000))
      @conn.transaction do
        @conn.exec_params("declare etl_read no scroll cursor for #{sql.strip.chomp(';')}", params)
        loop do
          result = @conn.exec("fetch #{size} from etl_read")
          break if result.ntuples.zero?

          result.each { |row| yield row }
        end
        @conn.exec('close etl_read')
      end
      nil
    end

    def truncate_once table_name
      return unless @options[:truncate] && !(@truncated ||= {})[table_name]

      @conn.exec("delete from #{quote_table(table_name)}")
      @truncated[table_name] = true
    end

    def insert table_name, row
      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})#{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('.')
    end
  end
end