datapipeline 2.0.0: Gem aus dem Projekt etl herausgelöst

test/plugins/postgres_storage_test.rb

@@ -0,0 +1,133 @@
+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