214 lines
8.1 KiB
Ruby
214 lines
8.1 KiB
Ruby
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
|
|
|
|
def test_upsert_updates_existing_rows
|
|
with_pg_database do |db|
|
|
pg_exec db, 'create table u (belnr text primary key, betrag numeric, text text)'
|
|
writer = DataPipeline::PostgresStorage.new(**db, table: 'u', upsert: %w[belnr])
|
|
writer.write([{ 'belnr' => 'A', 'betrag' => 1, 'text' => 'alt' }, { 'belnr' => 'B', 'betrag' => 2, 'text' => 'b' }])
|
|
writer.write([{ 'belnr' => 'A', 'betrag' => 5, 'text' => 'neu' }])
|
|
writer.write_row({ 'belnr' => 'B' })
|
|
writer.write_row({ 'belnr' => 'C', 'betrag' => 3 })
|
|
writer.close
|
|
assert_equal [%w[A 5 neu], %w[B 2 b], ['C', '3', nil]],
|
|
pg_exec(db, 'select belnr, betrag::int::text as betrag, text from u order by belnr').map(&:values)
|
|
end
|
|
end
|
|
|
|
def test_write_row_in_batches
|
|
with_pg_database do |db|
|
|
pg_exec db, 'create table b (id int)'
|
|
writer = DataPipeline::PostgresStorage.new(**db, table: 'b', batch_size: 2)
|
|
3.times { |i| writer.write_row({ 'id' => i }) }
|
|
assert_equal '2', pg_exec(db, 'select count(*) from b').first['count'] # third row still pending
|
|
writer.finalize
|
|
assert_equal '3', pg_exec(db, 'select count(*) from b').first['count']
|
|
end
|
|
end
|
|
|
|
def test_close_writes_pending_rows
|
|
with_pg_database do |db|
|
|
pg_exec db, 'create table b (id int)'
|
|
writer = DataPipeline::PostgresStorage.new(**db, table: 'b')
|
|
writer.write_row({ 'id' => 1 })
|
|
writer.close
|
|
assert_equal '1', pg_exec(db, 'select count(*) from b').first['count']
|
|
end
|
|
end
|
|
|
|
def test_batch_size_all_is_one_transaction_for_the_stream
|
|
with_pg_database do |db|
|
|
pg_exec db, 'create table b (id int); insert into b values (99)'
|
|
aborted = DataPipeline::PostgresStorage.new(**db, table: 'b', batch_size: :all, truncate: true)
|
|
2.times { |i| aborted.write_row({ 'id' => i }) }
|
|
aborted.close # without finalize: rolled back, truncate too
|
|
assert_equal ['99'], pg_exec(db, 'select id from b').map { |r| r['id'] }
|
|
|
|
done = DataPipeline::PostgresStorage.new(**db, table: 'b', batch_size: :all, truncate: true)
|
|
2.times { |i| done.write_row({ 'id' => i }) }
|
|
done.finalize
|
|
assert_equal %w[0 1], pg_exec(db, 'select id from b order by id').map { |r| r['id'] }
|
|
end
|
|
end
|
|
|
|
def test_streamed_pipeline_aborted_with_batch_size_all_writes_nothing
|
|
with_pg_database do |db|
|
|
pg_exec db, 'create table b (id int)'
|
|
assert_raises(RuntimeError) do
|
|
DataPipeline.new do
|
|
logger PipelineHelper::QUIET
|
|
streamed
|
|
read MemoryStorage.new(rows: Array.new(5) { |i| { 'id' => i } })
|
|
map { |row| row['id'] == 3 ? raise('kaputt') : row }
|
|
write DataPipeline::PostgresStorage.new(**db, table: 'b', batch_size: :all)
|
|
end.run
|
|
end
|
|
assert_equal '0', pg_exec(db, 'select count(*) from b').first['count']
|
|
end
|
|
end
|
|
|
|
def test_streamed_copy_on_the_same_connection_as_the_cursor
|
|
with_pg_database do |db|
|
|
pg_exec db, 'create table k as select i as id from generate_series(1, 25) i; create table c (id int)'
|
|
pg = DataPipeline::PostgresStorage.new(**db, query: 'select id from k', table: 'c', fetch_size: 10, batch_size: 7)
|
|
DataPipeline.new do
|
|
logger PipelineHelper::QUIET
|
|
streamed
|
|
read pg
|
|
write pg
|
|
end.run
|
|
assert_equal '25', pg_exec(db, 'select count(*) from c').first['count']
|
|
end
|
|
end
|
|
end
|