datapipeline/lib/datapipeline/storage/rest.rb

require 'json'
require_relative 'storage_plugin'

class DataPipeline
  #
  # options
  #   url:          base url (read and write)
  #   headers:      hash of request headers
  #   auth:         value of the Authorization header, e.g. "Bearer #{ENV['API_TOKEN']}"
  #   params:       query parameters for read
  #   root:         path to the array in the response, e.g. 'data.items' (default: the response itself)
  #   timeout:      seconds (default 30)
  #   ssl_verify:   false -> accept any server certificate (own CA, IP address); default true
  #   retries:      repeats on 5xx, 429 and network errors (default 2), wait retry_wait * attempt seconds
  #
  # paging for read (one of)
  #   page_param:   'page' -> page_start (default 1), page_start + 1, ... until a page returns no rows
  #   next:         ->(body) { body['next'] } -> url of the next page or nil
  #   max_pages:    safety stop (default 1000)
  #
  # write
  #   write_url:    default url; or ->(row) { "#{base}/kunden/#{row['id']}" } for one url per row
  #   write_method: :post (default), :put, :patch
  #   batch_size:   nil -> one request per row, n -> json array of n rows per request
  #
  # query(query: 'path', params: {...}) -> GET url + path, rows like read (for lookups)
  #
  # call(:post, 'person/listkunden', json: {...}) -> one request, parsed json back (for rpc style apis)
  #   retries: 0 for requests that create something, so a 5xx is not sent twice
  #   a failed request raises RestStorage::Error with #status (nil for network errors)
  #
  class RestStorage < StoragePlugin
    class Error < StandardError
      attr_reader :status

      def initialize message = nil, status: nil
        super(message)
        @status = status
      end
    end

    class RetryableError < Error; end

    RETRY_STATUS = [429, 500, 502, 503, 504].freeze

    def initialize **options, &block
      self.class.require_gem 'http'
      @options = { timeout: 30, retries: 2, retry_wait: 1, max_pages: 1000, write_method: :post }.merge(options)
      @url = @options.fetch(:url) { raise ArgumentError, 'RestStorage: option url missing' }
      if @options[:batch_size] && @options[:write_url].respond_to?(:call)
        raise ArgumentError, 'RestStorage: write_url per row does not work with batch_size'
      end
      @buffer = []

      yield if block_given?
    end

    def read(&block)
      rows = enum_for(:each_page_row, @url, @options.fetch(:params, {}))
      return rows.to_a unless block

      rows.each(&block)
      nil
    end

    def query options = {}
      path = options[:query].to_s
      url = path.start_with?('http') ? path : "#{@url.chomp('/')}/#{path.delete_prefix('/')}"
      rows = enum_for(:each_page_row, url, options.fetch(:params, {}))
      return rows.to_a unless block_given?

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

    # one request against url + path; json body and/or query params
    def call method, path, json: nil, params: nil, retries: nil
      url = path.to_s.start_with?('http') ? path.to_s : "#{@url.chomp('/')}/#{path.to_s.delete_prefix('/')}"
      args = { json:, params: }.compact
      request(method, url, retries:, **args)
    end

    def write rows
      rows.each { |row| write_row row }
    ensure
      close
    end

    def write_row row
      if @options[:batch_size]
        @buffer << row
        flush if @buffer.size >= @options[:batch_size]
      else
        send_rows row
      end
    end

    def close
      flush
    end

    private

    def flush
      return if @buffer.empty?

      send_rows @buffer
      @buffer = []
    end

    def send_rows payload
      url = @options.fetch(:write_url, @url)
      url = url.call(payload) if url.respond_to?(:call)
      request(@options[:write_method], url, json: payload)
    end

    def each_page_row url, params
      page = @options.fetch(:page_start, 1)
      @options[:max_pages].times do
        query = @options[:page_param] ? params.merge(@options[:page_param] => page) : params
        body = request(:get, url, params: query)
        rows = rows_of(body)
        rows.each { |row| yield row }

        if @options[:next]
          url = @options[:next].call(body)
          params = {}
          break if url.nil? || url.to_s.empty?
        elsif @options[:page_param]
          break if rows.empty?

          page += 1
        else
          break
        end
      end
    end

    def rows_of body
      data = body
      @options[:root].to_s.split('.').each { |key| data = data.fetch(key) } if @options[:root]
      data = [] if data.nil?
      raise Error, "RestStorage: response is not an array (root: #{@options[:root].inspect})" unless data.is_a?(Array)

      data
    end

    def request method, url, retries: nil, **args
      retries ||= @options[:retries]
      args[:ssl_context] = ssl_context if @options[:ssl_verify] == false && url.start_with?('https')
      attempt = 0
      begin
        attempt += 1
        response = client.request(method, url, **args)
        status = response.status.code
        text = response.body.to_s
        raise RetryableError.new("HTTP #{status}", status:) if RETRY_STATUS.include?(status)
        raise Error.new("#{method.upcase} #{url}: HTTP #{status} #{text[0, 200]}", status:) unless response.status.success?

        text.strip.empty? ? nil : JSON.parse(text)
      rescue RetryableError, HTTP::Error, SystemCallError => e
        raise Error.new("#{method.upcase} #{url}: #{e.message}", status: e.respond_to?(:status) ? e.status : nil) if attempt > retries

        sleep(@options[:retry_wait] * attempt)
        retry
      end
    end

    def ssl_context
      require 'openssl'
      OpenSSL::SSL::SSLContext.new.tap { |ctx| ctx.verify_mode = OpenSSL::SSL::VERIFY_NONE }
    end

    def client
      headers = { 'Accept' => 'application/json' }.merge(@options.fetch(:headers, {}))
      headers['Authorization'] = @options[:auth] if @options[:auth]
      HTTP.timeout(@options[:timeout]).headers(headers)
    end
  end
end