datapipeline/test/core/streaming_test.rb

require_relative '../test_helper'

# streamed mode, lookups, several targets
class StreamingTest < PipelineTest
  ROWS = [{ 'a' => 1 }, { 'a' => 2 }, { 'a' => 3 }].freeze

  CHAIN = proc do
    header { |c| { 'H' => c.written } }
    filter { |r| r['a'] > 1 }
    map { |r, c| [r, r.merge('k' => c.index)] }
    map { |r, c| r.merge('i' => c.index) }
    trailer { |c| { 'T' => c.written, 'R' => c.read } }
  end

  def test_batch_and_stream_same_body
    batch = output(ROWS, &CHAIN)
    stream = output(ROWS, streamed: true, &CHAIN)
    assert_equal batch[1..-2], stream[1..-2]
    assert_equal({ 'H' => 4 }, batch.first)
    assert_equal({ 'H' => nil }, stream.first)
    assert_equal({ 'T' => 4, 'R' => 3 }, stream.last)
  end

  def test_stream_several_steps_one_to_many
    out = output(ROWS, streamed: true) do
      map { |r| [r, r] }
      map { |r, ctx| r.merge('i' => ctx.index) }
      trailer { |ctx| { 'n' => ctx.written } }
    end
    assert_equal({ 'n' => 6 }, out.last)
    assert_equal (0...6).to_a, out.first(6).map { |r| r['i'] }
  end

  def test_stream_condition_and_copy
    assert_equal [{ 'a' => 1 }, { 'a' => 3 }], output(ROWS, streamed: true) { map(condition: ->(r) { r['a'] != 2 }) { |r| r } }
    assert_equal ROWS, output(ROWS, streamed: true)
  end

  def test_stream_csv_to_csv
    with_tmpdir do |dir|
      File.write("#{dir}/in.csv", "a,b\n1,x\n2,y\n3,z\n")
      DataPipeline.new do
        logger PipelineHelper::QUIET
        streamed
        read container(:csv, file: "#{dir}/in.csv")
        filter { |r| r['a'] != '2' }
        trailer { |c| { 'a' => 'T', 'b' => c.written } }
        write container(:csv, file: "#{dir}/out.csv")
      end.run
      assert_equal "a,b\n1,x\n3,z\nT,2\n", File.read("#{dir}/out.csv")
    end
  end

  def test_stream_postgres_with_query_on_same_connection
    with_pg_database do |db|
      pg_exec db, "create table k as select i as id, case when i % 2 = 0 then 'DE' else 'AT' end as land from generate_series(1, 5) i;
                   create table l (land text, name text); insert into l values ('DE','Deutschland'),('AT','Oesterreich')"
      with_tmpdir do |dir|
        DataPipeline.new do
          logger PipelineHelper::QUIET
          streamed
          pg = container :postgres, **db
          read pg, query: 'select id, land from k where id <= $1 order by id', params: [5], fetch_size: 2
          map { |r| r.merge('name' => query(pg, query: 'select name from l where land = $1', params: [r['land']]).first['name']) }
          write container(:csv, file: "#{dir}/o.csv", columns: %w[id name])
        end.run
        assert_equal 6, File.readlines("#{dir}/o.csv").size
        assert_includes File.read("#{dir}/o.csv"), '2,Deutschland'
      end
    end
  end

  def test_lookup_grouped_unique_and_lambda
    with_tmpdir do |dir|
      File.write("#{dir}/lk.csv", "id,grp\n1,a\n2,b\n3,a\n")
      grouped = unique = by_lambda = nil
      DataPipeline.new do
        csv = container(:csv, file: "#{dir}/lk.csv")
        grouped = lookup csv, key: 'grp'
        unique = lookup csv, key: 'id', unique: true
        by_lambda = lookup csv, key: ->(r) { r['grp'].upcase }
      end
      assert_equal 2, grouped['a'].size
      assert_equal 'b', unique['2']['grp']
      assert_equal %w[A B], by_lambda.keys.sort
    end
  end

  def test_lookup_postgres_join_without_n_plus_one
    with_pg_database do |db|
      pg_exec db, "create table k as select i as id, case when i % 2 = 0 then 'DE' else 'AT' end as land from generate_series(1, 3) i;
                   create table l (land text, name text); insert into l values ('DE','Deutschland'),('AT','Oesterreich')"
      out = nil
      DataPipeline.new do
        logger PipelineHelper::QUIET
        pg = container :postgres, **db
        names = lookup pg, query: 'select * from l where land like $1', params: ['%'], key: 'land', unique: true
        read pg, query: 'select id, land from k order by id'
        map { |r| r.merge('name' => names.dig(r['land'], 'name')) }
        write(out = MemoryStorage.new)
      end.run
      assert_equal %w[Oesterreich Deutschland Oesterreich], out.output.map { |r| r['name'] }
    end
  end

  def test_several_targets_with_only
    rows = [{ 'satyp' => 'DOHD', 'n' => 1 }, { 'satyp' => 'DOPS', 'n' => 2 }, { 'satyp' => 'DOPS', 'n' => 3 }]
    [false, true].each do |streamed|
      all = MemoryStorage.new
      heads = MemoryStorage.new
      positions = MemoryStorage.new
      trailers = MemoryStorage.new
      DataPipeline.new(streamed:) do
        logger PipelineHelper::QUIET
        read MemoryStorage.new(rows:)
        header { { 'satyp' => 'FIRC' } }
        map { |r| r }
        trailer { |c| { 'satyp' => 'LARC', 'n' => c.written } }
        write all
        write heads, only: ->(r) { r['satyp'] == 'DOHD' }
        write positions, only: ->(r) { r['satyp'] == 'DOPS' }
        write trailers, only: ->(r) { r['satyp'] == 'LARC' }
      end.run
      assert_equal 5, all.output.size, "streamed=#{streamed}"
      assert_equal [1], heads.output.map { |r| r['n'] }
      assert_equal 2, positions.output.size
      assert_equal [{ 'satyp' => 'LARC', 'n' => 3 }], trailers.output
    end
  end

  def test_memory_stays_flat_when_streaming
    with_tmpdir do |dir|
      File.open("#{dir}/big.csv", 'w') { |f| f.puts 'id,text'; 50_000.times { |i| f.puts "#{i},#{'x' * 80}" } }
      stats = DataPipeline.new(streamed: true) do
        logger PipelineHelper::QUIET
        read container(:csv, file: "#{dir}/big.csv")
        map { |r| r.merge('len' => r['text'].size) }
        write container(:csv, file: "#{dir}/out.csv", columns: %w[id len])
      end.run
      assert_equal 50_000, stats.written
      assert_equal 50_001, File.foreach("#{dir}/out.csv").count
    end
  end
end