datapipeline/test/plugins/postgres_storage_test.rb

require_relative '../test_helper'
require_relative '../support/storage_contract'

class PostgresStorageTest < PipelineTest
  include StorageContract

  def with_contract_env
    with_pg_database do |db|
      pg_exec db, 'create table rows (id text, code text, name text, note text)'
      yield db
    end
  end

  def contract_storage db, purpose
    purpose == :write ? DataPipeline::PostgresStorage.new(**db, table: 'rows') : DataPipeline::PostgresStorage.new(**db, query: 'select * from rows order by id')
  end

  def with_table
    with_pg_database do |db|
      pg_exec db, 'create table t (id int, name text, note text); create schema s; create table s."Mixed Case" ("Col A" text)'
      storage = DataPipeline::PostgresStorage.new(**db)
      yield db, storage
    ensure
      storage&.close
    end
  end

  def test_write_rows_with_table_option_quotes_and_null
    with_table do |db, pg|
      DataPipeline::PostgresStorage.new(**db, table: 't').tap do |w|
        w.write([{ 'id' => 1, 'name' => "O'Brien", 'note' => nil }, { id: 2, name: "x'); drop table t; --" }])
        w.close
      end
      assert_equal [{ 'id' => '1', 'name' => "O'Brien", 'note' => nil },
                    { 'id' => '2', 'name' => "x'); drop table t; --", 'note' => nil }],
                   pg.query(query: 'select id, name, note from t order by id')
    end
  end

  def test_write_old_signature_and_schema_quoting
    with_table do |_db, pg|
      pg.write('t', [{ 'id' => 3 }])
      pg.write('s.Mixed Case', [{ 'Col A' => 'v' }])
      assert_equal '1', pg.query(query: 'select count(*) from t').first['count']
      assert_equal [{ 'Col A' => 'v' }], pg.query(query: 'select * from s."Mixed Case"')
    end
  end

  def test_truncate_before_first_write
    with_table do |db, pg|
      pg.write('t', [{ 'id' => 1 }, { 'id' => 2 }])
      writer = DataPipeline::PostgresStorage.new(**db, table: 't', truncate: true)
      writer.write([{ 'id' => 3 }])
      writer.write_row({ 'id' => 4 })
      writer.close
      assert_equal %w[3 4], pg.query(query: 'select id from t order by id').map { |r| r['id'] }
    end
  end

  def test_write_is_one_transaction
    with_table do |_db, pg|
      assert_raises(PG::Error) { pg.write('t', [{ 'id' => 10 }, { 'id' => 'kein int' }]) }
      assert_equal '0', pg.query(query: 'select count(*) from t').first['count']
    end
  end

  def test_write_without_table
    with_table do |_db, pg|
      assert_includes assert_raises(RuntimeError) { pg.write([{ 'id' => 1 }]) }.message, 'Table'
    end
  end

  def test_query_and_read_with_params
    with_table do |db, pg|
      pg.write('t', [{ 'id' => 1, 'name' => 'a' }, { 'id' => 2, 'name' => 'b' }, { 'id' => 3 }])
      assert_equal [{ 'name' => 'a' }], pg.query(query: 'select name from t where id = $1', params: [1])

      reader = DataPipeline::PostgresStorage.new(**db)
      reader.set(query: 'select id from t where id > $1 order by id', params: [1])
      assert_equal [{ 'id' => '2' }, { 'id' => '3' }], reader.read
      reader.close
    end
  end

  def test_typed_and_default_string_values
    with_table do |db, pg|
      typed = DataPipeline::PostgresStorage.new(**db, typed: true)
      row = typed.query(query: 'select 1::int as i, 1.5::float8 as f, current_date as d').first
      typed.close
      assert_equal 1, row['i']
      assert_in_delta 1.5, row['f']
      assert_kind_of Date, row['d']
      assert_equal '1', pg.query(query: 'select 1::int as i').first['i']
    end
  end

  def test_cursor_read_in_batches
    with_pg_database do |db|
      pg_exec db, 'create table k as select i as id from generate_series(1, 2500) i'
      reader = DataPipeline::PostgresStorage.new(**db, query: "select id from k where id > $1 order by id;\n", params: [500],
                                         fetch_size: 100)
      count = 0
      reader.read { count += 1 }
      reader.close
      assert_equal 2000, count
    end
  end

  def test_cursor_rollback_on_error_keeps_connection_usable
    with_pg_database do |db|
      pg_exec db, 'create table k as select 1 as id'
      pg = DataPipeline::PostgresStorage.new(**db, query: 'select id from k')
      assert_raises(RuntimeError) { pg.read { raise 'x' } }
      assert_equal [{ 'x' => '1' }], pg.query(query: 'select 1 as x')
      pg.close
    end
  end

  def test_close_twice
    with_pg_database do |db|
      pg = DataPipeline::PostgresStorage.new(**db)
      pg.close
      pg.close
      assert pg.instance_variable_get(:@conn).finished?
    end
  end

  def test_connection_error_keeps_cause
    skip 'no local postgres' unless PgHelper.available?
    error = assert_raises(PG::ConnectionBad) { DataPipeline::PostgresStorage.new(host: 'localhost', port: 1, database: 'x') }
    assert_includes error.message, 'Error connection to database:'
  end
end