datapipeline/test/plugins/sqlite_storage_test.rb

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

class SqliteStorageTest < PipelineTest
  include StorageContract

  def contract_storage dir, purpose
    db = "#{dir}/test.sqlite3"
    unless File.exist?(db)
      DataPipeline::SqliteStorage.new(database: db).tap do |s|
        s.query(query: 'create table rows (id text, code text, name text, note text)')
        s.close
      end
    end
    purpose == :write ? DataPipeline::SqliteStorage.new(database: db, table: 'rows') : DataPipeline::SqliteStorage.new(database: db, query: 'select * from rows order by id')
  end

  def memory_db
    DataPipeline::SqliteStorage.new(database: ':memory:').tap do |s|
      s.query(query: 'create table t (id integer, name text, betrag real, datum text)')
    end
  end

  def test_types_null_and_quotes
    db = memory_db
    db.write('t', [{ 'id' => 1, 'name' => "O'Brien", 'betrag' => 1.5, 'datum' => Date.new(2025, 11, 3) }, { id: 2, name: nil }])
    assert_equal [{ 'id' => 1, 'name' => "O'Brien", 'betrag' => 1.5, 'datum' => '2025-11-03' },
                  { 'id' => 2, 'name' => nil, 'betrag' => nil, 'datum' => nil }],
                 db.query(query: 'select * from t order by id')
    db.close
  end

  def test_positional_and_named_params
    db = memory_db
    db.write('t', [{ 'id' => 1, 'name' => 'a' }, { 'id' => 2, 'name' => 'b' }])
    assert_equal [{ 'name' => 'b' }], db.query(query: 'select name from t where id = ?', params: [2])
    assert_equal [{ 'name' => 'a' }], db.query(query: 'select name from t where id = :id', params: { id: 1 })
    db.close
  end

  def test_write_is_one_transaction
    db = memory_db
    db.query(query: 'create table u (id integer not null)')
    assert_raises(SQLite3::Exception) { db.write('u', [{ 'id' => 1 }, { 'id' => nil }]) }
    assert_equal 0, db.query(query: 'select count(*) as n from u').first['n']
    db.close
  end

  def test_sql_injection_in_values_and_identifiers
    db = memory_db
    db.query(query: 'create table "Mixed ""Name""" ("Col A" text)')
    db.write('Mixed "Name"', [{ 'Col A' => "x'); drop table t; --" }])
    assert_equal [{ 'Col A' => "x'); drop table t; --" }], db.query(query: 'select * from "Mixed ""Name"""')
    assert_equal 0, db.query(query: 'select count(*) as n from t').first['n']
    db.close
  end

  def test_create_truncate_and_lazy_open
    with_tmpdir do |dir|
      path = "#{dir}/neu.sqlite3"
      storage = DataPipeline::SqliteStorage.new(database: path, table: 'k', create: true, truncate: true)
      refute File.exist?(path), 'opened only on first use'
      storage.write([{ 'id' => 1, 'betrag' => 1.5, 'name' => 'a' }, { 'id' => 2, 'betrag' => 2.0, 'name' => 'b' }])
      storage.close

      again = DataPipeline::SqliteStorage.new(database: path, table: 'k', create: true, truncate: true)
      again.write_row({ 'id' => 3, 'betrag' => 3.0, 'name' => 'c' })
      assert_equal [{ 'id' => 3, 'betrag' => 3.0, 'name' => 'c' }], again.query(query: 'select * from k')
      assert_equal %w[integer real text], again.query(query: 'pragma table_info(k)').map { |c| c['type'].downcase }
      again.close
    end
  end

  def test_missing_options
    assert_raises(ArgumentError) { DataPipeline::SqliteStorage.new }
    db = memory_db
    assert_raises(RuntimeError) { db.write([{ 'id' => 1 }]) }
    assert_raises(RuntimeError) { db.read }
    db.close
  end

  def test_short_name_and_lookup
    with_tmpdir do |dir|
      path = "#{dir}/l.sqlite3"
      names = nil
      DataPipeline.new do
        db = container :sqlite3, database: path
        db.query(query: 'create table land (code text, name text)')
        db.write('land', [{ 'code' => 'DE', 'name' => 'Deutschland' }])
        names = lookup db, query: 'select * from land where code like ?', params: ['%'], key: 'code', unique: true
      end
      assert_equal 'Deutschland', names['DE']['name']
    end
  end

  def test_upsert_with_created_unique_index
    with_tmpdir do |dir|
      writer = DataPipeline::SqliteStorage.new(database: "#{dir}/u.sqlite3", table: 'u', create: true, upsert: %w[belnr])
      writer.write([{ 'belnr' => 'A', 'betrag' => 1.0 }, { 'belnr' => 'B', 'betrag' => 2.0 }])
      writer.write([{ 'belnr' => 'A', 'betrag' => 5.0 }])
      writer.write_row({ 'belnr' => 'B' })
      assert_equal [['A', 5.0], ['B', 2.0]], writer.query(query: 'select belnr, betrag from u order by belnr').map(&:values)
      writer.close
    end
  end

  def count(file) = DataPipeline::SqliteStorage.new(database: file).then { |db| db.query(query: 'select count(*) as n from b').first['n'].tap { db.close } }

  def test_write_row_in_batches_and_query_sees_pending_rows
    with_tmpdir do |dir|
      file = "#{dir}/b.sqlite3"
      writer = DataPipeline::SqliteStorage.new(database: file, table: 'b', create: true, batch_size: 2)
      3.times { |i| writer.write_row({ 'id' => i }) }
      assert_equal 2, count(file) # third row still pending
      assert_equal 3, writer.query(query: 'select count(*) as n from b').first['n']
      writer.write_row({ 'id' => 3 })
      writer.finalize
      assert_equal 4, count(file)
    end
  end

  def test_close_writes_pending_rows
    with_tmpdir do |dir|
      writer = DataPipeline::SqliteStorage.new(database: "#{dir}/b.sqlite3", table: 'b', create: true)
      writer.write_row({ 'id' => 1 })
      writer.close
      assert_equal 1, count("#{dir}/b.sqlite3")
    end
  end

  def test_batch_size_all_is_one_transaction_for_the_stream
    with_tmpdir do |dir|
      file = "#{dir}/b.sqlite3"
      DataPipeline::SqliteStorage.new(database: file, table: 'b', create: true).tap { |w| w.write([{ 'id' => 99 }]); w.close }
      aborted = DataPipeline::SqliteStorage.new(database: file, table: 'b', truncate: true, batch_size: :all)
      2.times { |i| aborted.write_row({ 'id' => i }) }
      aborted.close # without finalize: rolled back, truncate too
      assert_equal 1, count(file)

      done = DataPipeline::SqliteStorage.new(database: file, table: 'b', truncate: true, batch_size: :all)
      2.times { |i| done.write_row({ 'id' => i }) }
      done.finalize
      assert_equal 2, count(file)
    end
  end

  def test_streamed_pipeline_aborted_with_batch_size_all_writes_nothing
    with_tmpdir do |dir|
      file = "#{dir}/b.sqlite3"
      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::SqliteStorage.new(database: file, table: 'b', create: true, batch_size: :all)
        end.run
      end
      refute DataPipeline::SqliteStorage.new(database: file).query(query: "select name from sqlite_master where name = 'b'").any?
    end
  end
end