/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.beam.sdk.io.jdbc;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertTrue;
import static org.mockito.Matchers.any;
import static org.mockito.Matchers.anyString;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

import java.io.PrintWriter;
import java.io.Serializable;
import java.io.StringWriter;
import java.math.BigDecimal;
import java.net.InetAddress;
import java.nio.charset.Charset;
import java.sql.Array;
import java.sql.Connection;
import java.sql.Date;
import java.sql.JDBCType;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Statement;
import java.sql.Time;
import java.sql.Timestamp;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Calendar;
import java.util.Collections;
import java.util.List;
import java.util.TimeZone;
import javax.sql.DataSource;
import org.apache.beam.sdk.coders.KvCoder;
import org.apache.beam.sdk.coders.SerializableCoder;
import org.apache.beam.sdk.coders.StringUtf8Coder;
import org.apache.beam.sdk.coders.VarIntCoder;
import org.apache.beam.sdk.io.common.DatabaseTestHelper;
import org.apache.beam.sdk.io.common.NetworkTestHelper;
import org.apache.beam.sdk.io.common.TestRow;
import org.apache.beam.sdk.schemas.Schema;
import org.apache.beam.sdk.schemas.transforms.Select;
import org.apache.beam.sdk.testing.ExpectedLogs;
import org.apache.beam.sdk.testing.PAssert;
import org.apache.beam.sdk.testing.TestPipeline;
import org.apache.beam.sdk.transforms.Count;
import org.apache.beam.sdk.transforms.Create;
import org.apache.beam.sdk.transforms.SerializableFunction;
import org.apache.beam.sdk.transforms.Wait;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.Row;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
import org.apache.commons.dbcp2.PoolingDataSource;
import org.apache.derby.drda.NetworkServerControl;
import org.apache.derby.jdbc.ClientDataSource;
import org.joda.time.DateTime;
import org.joda.time.LocalDate;
import org.joda.time.chrono.ISOChronology;
import org.junit.After;
import org.junit.AfterClass;
import org.junit.Assert;
import org.junit.BeforeClass;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.ExpectedException;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/** Test on the JdbcIO. */
@RunWith(JUnit4.class)
public class JdbcIOTest implements Serializable {

  private static final Logger LOG = LoggerFactory.getLogger(JdbcIOTest.class);
  private static final int EXPECTED_ROW_COUNT = 1000;
  private static final String BACKOFF_TABLE = "UT_WRITE_BACKOFF";

  private static NetworkServerControl derbyServer;
  private static ClientDataSource dataSource;

  private static int port;
  private static String readTableName;

  @Rule public final transient TestPipeline pipeline = TestPipeline.create();

  @Rule public final transient ExpectedLogs expectedLogs = ExpectedLogs.none(JdbcIO.class);

  @Rule public transient ExpectedException thrown = ExpectedException.none();

  @BeforeClass
  public static void beforeClass() throws Exception {
    port = NetworkTestHelper.getAvailableLocalPort();
    LOG.info("Starting Derby database on {}", port);

    // by default, derby uses a lock timeout of 60 seconds. In order to speed up the test
    // and detect the lock faster, we decrease this timeout
    System.setProperty("derby.locks.waitTimeout", "2");
    System.setProperty("derby.stream.error.file", "target/derby.log");

    derbyServer = new NetworkServerControl(InetAddress.getByName("localhost"), port);
    StringWriter out = new StringWriter();
    derbyServer.start(new PrintWriter(out));
    boolean started = false;
    int count = 0;
    // Use two different methods to detect when server is started:
    // 1) Check the server stdout for the "started" string
    // 2) wait up to 15 seconds for the derby server to start based on a ping
    // on faster machines and networks, this may return very quick, but on slower
    // networks where the DNS lookups are slow, this may take a little time
    while (!started && count < 30) {
      if (out.toString().contains("started")) {
        started = true;
      } else {
        count++;
        Thread.sleep(500);
        try {
          derbyServer.ping();
          started = true;
        } catch (Throwable t) {
          // ignore, still trying to start
        }
      }
    }

    dataSource = new ClientDataSource();
    dataSource.setCreateDatabase("create");
    dataSource.setDatabaseName("target/beam");
    dataSource.setServerName("localhost");
    dataSource.setPortNumber(port);

    readTableName = DatabaseTestHelper.getTestTableName("UT_READ");

    DatabaseTestHelper.createTable(dataSource, readTableName);
    addInitialData(dataSource, readTableName);
  }

  @AfterClass
  public static void afterClass() throws Exception {
    try {
      DatabaseTestHelper.deleteTable(dataSource, readTableName);
    } finally {
      if (derbyServer != null) {
        derbyServer.shutdown();
      }
    }
  }

  @Test
  public void testDataSourceConfigurationDataSource() throws Exception {
    JdbcIO.DataSourceConfiguration config = JdbcIO.DataSourceConfiguration.create(dataSource);
    try (Connection conn = config.buildDatasource().getConnection()) {
      assertTrue(conn.isValid(0));
    }
  }

  @Test
  public void testDataSourceConfigurationDataSourceWithoutPool() {
    assertTrue(
        JdbcIO.DataSourceConfiguration.create(dataSource).buildDatasource()
            instanceof ClientDataSource);
  }

  @Test
  public void testDataSourceConfigurationDataSourceWithPool() {
    assertTrue(
        JdbcIO.PoolableDataSourceProvider.of(JdbcIO.DataSourceConfiguration.create(dataSource))
                .apply(null)
            instanceof PoolingDataSource);
  }

  @Test
  public void testDataSourceConfigurationDriverAndUrl() throws Exception {
    JdbcIO.DataSourceConfiguration config =
        JdbcIO.DataSourceConfiguration.create(
            "org.apache.derby.jdbc.ClientDriver",
            "jdbc:derby://localhost:" + port + "/target/beam");
    try (Connection conn = config.buildDatasource().getConnection()) {
      assertTrue(conn.isValid(0));
    }
  }

  @Test
  public void testDataSourceConfigurationUsernameAndPassword() throws Exception {
    String username = "sa";
    String password = "sa";
    JdbcIO.DataSourceConfiguration config =
        JdbcIO.DataSourceConfiguration.create(
                "org.apache.derby.jdbc.ClientDriver",
                "jdbc:derby://localhost:" + port + "/target/beam")
            .withUsername(username)
            .withPassword(password);
    try (Connection conn = config.buildDatasource().getConnection()) {
      assertTrue(conn.isValid(0));
    }
  }

  @Test
  public void testDataSourceConfigurationNullPassword() throws Exception {
    String username = "sa";
    String password = null;
    JdbcIO.DataSourceConfiguration config =
        JdbcIO.DataSourceConfiguration.create(
                "org.apache.derby.jdbc.ClientDriver",
                "jdbc:derby://localhost:" + port + "/target/beam")
            .withUsername(username)
            .withPassword(password);
    try (Connection conn = config.buildDatasource().getConnection()) {
      assertTrue(conn.isValid(0));
    }
  }

  @Test
  public void testDataSourceConfigurationNullUsernameAndPassword() throws Exception {
    String username = null;
    String password = null;
    JdbcIO.DataSourceConfiguration config =
        JdbcIO.DataSourceConfiguration.create(
                "org.apache.derby.jdbc.ClientDriver",
                "jdbc:derby://localhost:" + port + "/target/beam")
            .withUsername(username)
            .withPassword(password);
    try (Connection conn = config.buildDatasource().getConnection()) {
      assertTrue(conn.isValid(0));
    }
  }

  /** Create test data that is consistent with that generated by TestRow. */
  private static void addInitialData(DataSource dataSource, String tableName) throws SQLException {
    try (Connection connection = dataSource.getConnection()) {
      connection.setAutoCommit(false);
      try (PreparedStatement preparedStatement =
          connection.prepareStatement(String.format("insert into %s values (?,?)", tableName))) {
        for (int i = 0; i < EXPECTED_ROW_COUNT; i++) {
          preparedStatement.clearParameters();
          preparedStatement.setInt(1, i);
          preparedStatement.setString(2, TestRow.getNameForSeed(i));
          preparedStatement.executeUpdate();
        }
      }
      connection.commit();
    }
  }

  @Test
  public void testRead() {
    PCollection<TestRow> rows =
        pipeline.apply(
            JdbcIO.<TestRow>read()
                .withFetchSize(12)
                .withDataSourceConfiguration(JdbcIO.DataSourceConfiguration.create(dataSource))
                .withQuery("select name,id from " + readTableName)
                .withRowMapper(new JdbcTestHelper.CreateTestRowOfNameAndId())
                .withCoder(SerializableCoder.of(TestRow.class)));

    PAssert.thatSingleton(rows.apply("Count All", Count.globally()))
        .isEqualTo((long) EXPECTED_ROW_COUNT);

    Iterable<TestRow> expectedValues = TestRow.getExpectedValues(0, EXPECTED_ROW_COUNT);
    PAssert.that(rows).containsInAnyOrder(expectedValues);

    pipeline.run();
  }

  @Test
  public void testReadWithSingleStringParameter() {
    PCollection<TestRow> rows =
        pipeline.apply(
            JdbcIO.<TestRow>read()
                .withDataSourceConfiguration(JdbcIO.DataSourceConfiguration.create(dataSource))
                .withQuery(String.format("select name,id from %s where name = ?", readTableName))
                .withStatementPreparator(
                    preparedStatement -> preparedStatement.setString(1, TestRow.getNameForSeed(1)))
                .withRowMapper(new JdbcTestHelper.CreateTestRowOfNameAndId())
                .withCoder(SerializableCoder.of(TestRow.class)));

    PAssert.thatSingleton(rows.apply("Count All", Count.globally())).isEqualTo(1L);

    Iterable<TestRow> expectedValues = Collections.singletonList(TestRow.fromSeed(1));
    PAssert.that(rows).containsInAnyOrder(expectedValues);

    pipeline.run();
  }

  @Test
  public void testReadRows() {
    SerializableFunction<Void, DataSource> dataSourceProvider = ignored -> dataSource;
    PCollection<Row> rows =
        pipeline.apply(
            JdbcIO.readRows()
                .withDataSourceProviderFn(dataSourceProvider)
                .withQuery(String.format("select name,id from %s where name = ?", readTableName))
                .withStatementPreparator(
                    preparedStatement ->
                        preparedStatement.setString(1, TestRow.getNameForSeed(1))));

    Schema expectedSchema =
        Schema.of(
            Schema.Field.of("NAME", LogicalTypes.variableLengthString(JDBCType.VARCHAR, 500))
                .withNullable(true),
            Schema.Field.of("ID", Schema.FieldType.INT32).withNullable(true));

    assertEquals(expectedSchema, rows.getSchema());

    PCollection<Row> output = rows.apply(Select.fieldNames("NAME", "ID"));
    PAssert.that(output)
        .containsInAnyOrder(
            ImmutableList.of(Row.withSchema(expectedSchema).addValues("Testval1", 1).build()));

    pipeline.run();
  }

  @Test
  public void testReadRowsWithoutStatementPreparator() {
    SerializableFunction<Void, DataSource> dataSourceProvider = ignored -> dataSource;
    String name = TestRow.getNameForSeed(1);
    PCollection<Row> rows =
        pipeline.apply(
            JdbcIO.readRows()
                .withDataSourceProviderFn(dataSourceProvider)
                .withQuery(
                    String.format(
                        "select name,id from %s where name = '%s'", readTableName, name)));

    Schema expectedSchema =
        Schema.of(
            Schema.Field.of("NAME", LogicalTypes.variableLengthString(JDBCType.VARCHAR, 500))
                .withNullable(true),
            Schema.Field.of("ID", Schema.FieldType.INT32).withNullable(true));

    assertEquals(expectedSchema, rows.getSchema());

    PCollection<Row> output = rows.apply(Select.fieldNames("NAME", "ID"));
    PAssert.that(output)
        .containsInAnyOrder(
            ImmutableList.of(Row.withSchema(expectedSchema).addValues(name, 1).build()));

    pipeline.run();
  }

  @Test
  public void testReadWithSchema() {
    SerializableFunction<Void, DataSource> dataSourceProvider = ignored -> dataSource;
    JdbcIO.RowMapper<RowWithSchema> rowMapper =
        rs -> new RowWithSchema(rs.getString("NAME"), rs.getInt("ID"));
    pipeline.getSchemaRegistry().registerJavaBean(RowWithSchema.class);

    PCollection<RowWithSchema> rows =
        pipeline.apply(
            JdbcIO.<RowWithSchema>read()
                .withDataSourceProviderFn(dataSourceProvider)
                .withQuery(String.format("select name,id from %s where name = ?", readTableName))
                .withRowMapper(rowMapper)
                .withCoder(SerializableCoder.of(RowWithSchema.class))
                .withStatementPreparator(
                    preparedStatement ->
                        preparedStatement.setString(1, TestRow.getNameForSeed(1))));

    Schema expectedSchema =
        Schema.of(
            Schema.Field.of("name", Schema.FieldType.STRING),
            Schema.Field.of("id", Schema.FieldType.INT32));

    assertEquals(expectedSchema, rows.getSchema());

    PCollection<Row> output = rows.apply(Select.fieldNames("name", "id"));
    PAssert.that(output)
        .containsInAnyOrder(
            ImmutableList.of(Row.withSchema(expectedSchema).addValues("Testval1", 1).build()));

    pipeline.run();
  }

  @Test
  public void testWrite() throws Exception {
    final long rowsToAdd = 1000L;

    String tableName = DatabaseTestHelper.getTestTableName("UT_WRITE");
    DatabaseTestHelper.createTable(dataSource, tableName);
    try {
      ArrayList<KV<Integer, String>> data = getDataToWrite(rowsToAdd);
      pipeline.apply(Create.of(data)).apply(getJdbcWrite(tableName));

      pipeline.run();

      assertRowCount(tableName, EXPECTED_ROW_COUNT);
    } finally {
      DatabaseTestHelper.deleteTable(dataSource, tableName);
    }
  }

  @Test
  public void testWriteWithResultsAndWaitOn() throws Exception {
    final long rowsToAdd = 1000L;

    String firstTableName = DatabaseTestHelper.getTestTableName("UT_WRITE");
    String secondTableName = DatabaseTestHelper.getTestTableName("UT_WRITE_AFTER_WAIT");
    DatabaseTestHelper.createTable(dataSource, firstTableName);
    DatabaseTestHelper.createTable(dataSource, secondTableName);
    try {
      ArrayList<KV<Integer, String>> data = getDataToWrite(rowsToAdd);

      PCollection<KV<Integer, String>> dataCollection = pipeline.apply(Create.of(data));
      PCollection<Void> rowsWritten =
          dataCollection.apply(getJdbcWrite(firstTableName).withResults());
      dataCollection.apply(Wait.on(rowsWritten)).apply(getJdbcWrite(secondTableName));

      pipeline.run();

      assertRowCount(firstTableName, EXPECTED_ROW_COUNT);
      assertRowCount(secondTableName, EXPECTED_ROW_COUNT);
    } finally {
      DatabaseTestHelper.deleteTable(dataSource, firstTableName);
    }
  }

  private static JdbcIO.Write<KV<Integer, String>> getJdbcWrite(String tableName) {
    return JdbcIO.<KV<Integer, String>>write()
        .withDataSourceConfiguration(
            JdbcIO.DataSourceConfiguration.create(
                "org.apache.derby.jdbc.ClientDriver",
                "jdbc:derby://localhost:" + port + "/target/beam"))
        .withStatement(String.format("insert into %s values(?, ?)", tableName))
        .withBatchSize(10L)
        .withPreparedStatementSetter(
            (element, statement) -> {
              statement.setInt(1, element.getKey());
              statement.setString(2, element.getValue());
            });
  }

  private static ArrayList<KV<Integer, String>> getDataToWrite(long rowsToAdd) {
    ArrayList<KV<Integer, String>> data = new ArrayList<>();
    for (int i = 0; i < rowsToAdd; i++) {
      KV<Integer, String> kv = KV.of(i, "Test");
      data.add(kv);
    }
    return data;
  }

  private static void assertRowCount(String tableName, int expectedRowCount) throws SQLException {
    try (Connection connection = dataSource.getConnection()) {
      try (Statement statement = connection.createStatement()) {
        try (ResultSet resultSet = statement.executeQuery("select count(*) from " + tableName)) {
          resultSet.next();
          int count = resultSet.getInt(1);
          Assert.assertEquals(expectedRowCount, count);
        }
      }
    }
  }

  @Test
  public void testWriteWithBackoff() throws Exception {
    String tableName = DatabaseTestHelper.getTestTableName("UT_WRITE_BACKOFF");
    DatabaseTestHelper.createTable(dataSource, tableName);

    // lock table
    Connection connection = dataSource.getConnection();
    Statement lockStatement = connection.createStatement();
    lockStatement.execute("ALTER TABLE " + tableName + " LOCKSIZE TABLE");
    lockStatement.execute("LOCK TABLE " + tableName + " IN EXCLUSIVE MODE");

    // start a first transaction
    connection.setAutoCommit(false);
    PreparedStatement insertStatement =
        connection.prepareStatement("insert into " + tableName + " values(?, ?)");
    insertStatement.setInt(1, 1);
    insertStatement.setString(2, "TEST");
    insertStatement.execute();

    // try to write to this table
    pipeline
        .apply(Create.of(Collections.singletonList(KV.of(1, "TEST"))))
        .apply(
            JdbcIO.<KV<Integer, String>>write()
                .withDataSourceConfiguration(
                    JdbcIO.DataSourceConfiguration.create(
                        "org.apache.derby.jdbc.ClientDriver",
                        "jdbc:derby://localhost:" + port + "/target/beam"))
                .withStatement(String.format("insert into %s values(?, ?)", tableName))
                .withRetryStrategy(
                    (JdbcIO.RetryStrategy)
                        e -> {
                          return "XJ208"
                              .equals(e.getSQLState()); // we fake a deadlock with a lock here
                        })
                .withPreparedStatementSetter(
                    (element, statement) -> {
                      statement.setInt(1, element.getKey());
                      statement.setString(2, element.getValue());
                    }));

    // starting a thread to perform the commit later, while the pipeline is running into the backoff
    Thread commitThread =
        new Thread(
            () -> {
              try {
                Thread.sleep(10000);
                connection.commit();
              } catch (Exception e) {
                // nothing to do
              }
            });
    commitThread.start();
    pipeline.run();
    commitThread.join();

    // we verify the the backoff has been called thanks to the log message
    expectedLogs.verifyWarn("Deadlock detected, retrying");

    assertRowCount(tableName, 2);
  }

  @After
  public void tearDown() {
    try {
      DatabaseTestHelper.deleteTable(dataSource, BACKOFF_TABLE);
    } catch (Exception e) {
      // nothing to do
    }
  }

  @Test
  public void testWriteWithoutPreparedStatement() throws Exception {
    final int rowsToAdd = 10;

    Schema.Builder schemaBuilder = Schema.builder();
    schemaBuilder.addField(Schema.Field.of("column_boolean", Schema.FieldType.BOOLEAN));
    schemaBuilder.addField(Schema.Field.of("column_string", Schema.FieldType.STRING));
    schemaBuilder.addField(Schema.Field.of("column_int", Schema.FieldType.INT32));
    schemaBuilder.addField(Schema.Field.of("column_long", Schema.FieldType.INT64));
    schemaBuilder.addField(Schema.Field.of("column_float", Schema.FieldType.FLOAT));
    schemaBuilder.addField(Schema.Field.of("column_double", Schema.FieldType.DOUBLE));
    schemaBuilder.addField(Schema.Field.of("column_bigdecimal", Schema.FieldType.DECIMAL));
    schemaBuilder.addField(Schema.Field.of("column_date", LogicalTypes.JDBC_DATE_TYPE));
    schemaBuilder.addField(Schema.Field.of("column_time", LogicalTypes.JDBC_TIME_TYPE));
    schemaBuilder.addField(
        Schema.Field.of("column_timestamptz", LogicalTypes.JDBC_TIMESTAMP_WITH_TIMEZONE_TYPE));
    schemaBuilder.addField(Schema.Field.of("column_timestamp", Schema.FieldType.DATETIME));
    schemaBuilder.addField(Schema.Field.of("column_short", Schema.FieldType.INT16));
    Schema schema = schemaBuilder.build();

    String tableName = DatabaseTestHelper.getTestTableName("UT_WRITE_PS");
    StringBuilder stmt = new StringBuilder("CREATE TABLE ");
    stmt.append(tableName);
    stmt.append(" (");
    stmt.append("column_boolean       BOOLEAN,"); // boolean
    stmt.append("column_string        VARCHAR(254),"); // String
    stmt.append("column_int           INTEGER,"); // int
    stmt.append("column_long          BIGINT,"); // long
    stmt.append("column_float         REAL,"); // float
    stmt.append("column_double        DOUBLE PRECISION,"); // double
    stmt.append("column_bigdecimal    DECIMAL(13,0),"); // BigDecimal
    stmt.append("column_date          DATE,"); // Date
    stmt.append("column_time          TIME,"); // Time
    stmt.append("column_timestamptz   TIMESTAMP,"); // Timestamp
    stmt.append("column_timestamp     TIMESTAMP,"); // Timestamp
    stmt.append("column_short         SMALLINT"); // short
    stmt.append(" )");
    DatabaseTestHelper.createTableWithStatement(dataSource, stmt.toString());
    try {
      ArrayList<Row> data = getRowsToWrite(rowsToAdd, schema);
      pipeline
          .apply(Create.of(data))
          .setRowSchema(schema)
          .apply(
              JdbcIO.<Row>write()
                  .withDataSourceConfiguration(
                      JdbcIO.DataSourceConfiguration.create(
                          "org.apache.derby.jdbc.ClientDriver",
                          "jdbc:derby://localhost:" + port + "/target/beam"))
                  .withBatchSize(10L)
                  .withTable(tableName));
      pipeline.run();
      assertRowCount(tableName, rowsToAdd);
    } finally {
      DatabaseTestHelper.deleteTable(dataSource, tableName);
    }
  }

  @Test
  public void testWriteWithoutPreparedStatementWithReadRows() throws Exception {
    SerializableFunction<Void, DataSource> dataSourceProvider = ignored -> dataSource;
    PCollection<Row> rows =
        pipeline.apply(
            JdbcIO.readRows()
                .withDataSourceProviderFn(dataSourceProvider)
                .withQuery(String.format("select name,id from %s where name = ?", readTableName))
                .withStatementPreparator(
                    preparedStatement ->
                        preparedStatement.setString(1, TestRow.getNameForSeed(1))));

    String writeTableName = DatabaseTestHelper.getTestTableName("UT_WRITE_PS_WITH_READ_ROWS");
    DatabaseTestHelper.createTableForRowWithSchema(dataSource, writeTableName);
    try {
      rows.apply(
          JdbcIO.<Row>write()
              .withDataSourceConfiguration(
                  JdbcIO.DataSourceConfiguration.create(
                      "org.apache.derby.jdbc.ClientDriver",
                      "jdbc:derby://localhost:" + port + "/target/beam"))
              .withBatchSize(10L)
              .withTable(writeTableName));
      pipeline.run();
    } finally {
      DatabaseTestHelper.deleteTable(dataSource, writeTableName);
    }
  }

  @Test
  public void testWriteWithoutPsWithNonNullableTableField() throws Exception {
    final int rowsToAdd = 10;

    Schema.Builder schemaBuilder = Schema.builder();
    schemaBuilder.addField(Schema.Field.of("column_boolean", Schema.FieldType.BOOLEAN));
    schemaBuilder.addField(Schema.Field.of("column_string", Schema.FieldType.STRING));
    Schema schema = schemaBuilder.build();

    String tableName = DatabaseTestHelper.getTestTableName("UT_WRITE");
    StringBuilder stmt = new StringBuilder("CREATE TABLE ");
    stmt.append(tableName);
    stmt.append(" (");
    stmt.append("column_boolean       BOOLEAN,");
    stmt.append("column_int           INTEGER NOT NULL");
    stmt.append(" )");
    DatabaseTestHelper.createTableWithStatement(dataSource, stmt.toString());
    try {
      ArrayList<Row> data = getRowsToWrite(rowsToAdd, schema);
      pipeline
          .apply(Create.of(data))
          .setRowSchema(schema)
          .apply(
              JdbcIO.<Row>write()
                  .withDataSourceConfiguration(
                      JdbcIO.DataSourceConfiguration.create(
                          "org.apache.derby.jdbc.ClientDriver",
                          "jdbc:derby://localhost:" + port + "/target/beam"))
                  .withBatchSize(10L)
                  .withTable(tableName));
      pipeline.run();
    } finally {
      DatabaseTestHelper.deleteTable(dataSource, tableName);
      thrown.expect(RuntimeException.class);
    }
  }

  @Test
  public void testWriteWithoutPreparedStatementAndNonRowType() throws Exception {
    final int rowsToAdd = 10;

    String tableName = DatabaseTestHelper.getTestTableName("UT_WRITE_PS_NON_ROW");
    DatabaseTestHelper.createTableForRowWithSchema(dataSource, tableName);
    try {
      List<RowWithSchema> data = getRowsWithSchemaToWrite(rowsToAdd);

      pipeline
          .apply(Create.of(data))
          .apply(
              JdbcIO.<RowWithSchema>write()
                  .withDataSourceConfiguration(
                      JdbcIO.DataSourceConfiguration.create(
                          "org.apache.derby.jdbc.ClientDriver",
                          "jdbc:derby://localhost:" + port + "/target/beam"))
                  .withBatchSize(10L)
                  .withTable(tableName));
      pipeline.run();
      assertRowCount(tableName, rowsToAdd);
    } finally {
      DatabaseTestHelper.deleteTable(dataSource, tableName);
    }
  }

  @Test
  public void testGetPreparedStatementSetCaller() throws Exception {

    Schema schema =
        Schema.builder()
            .addField("bigint_col", Schema.FieldType.INT64)
            .addField("binary_col", Schema.FieldType.BYTES)
            .addField("bit_col", Schema.FieldType.BOOLEAN)
            .addField("char_col", Schema.FieldType.STRING)
            .addField("decimal_col", Schema.FieldType.DECIMAL)
            .addField("double_col", Schema.FieldType.DOUBLE)
            .addField("float_col", Schema.FieldType.FLOAT)
            .addField("integer_col", Schema.FieldType.INT32)
            .addField("datetime_col", Schema.FieldType.DATETIME)
            .addField("int16_col", Schema.FieldType.INT16)
            .addField("byte_col", Schema.FieldType.BYTE)
            .build();
    Row row =
        Row.withSchema(schema)
            .addValues(
                42L,
                "binary".getBytes(Charset.forName("UTF-8")),
                true,
                "char",
                BigDecimal.valueOf(25L),
                20.5D,
                15.5F,
                10,
                new DateTime(),
                (short) 5,
                Byte.parseByte("1", 2))
            .build();

    PreparedStatement psMocked = mock(PreparedStatement.class);

    JdbcUtil.getPreparedStatementSetCaller(Schema.FieldType.INT64)
        .set(row, psMocked, 0, SchemaUtil.FieldWithIndex.of(schema.getField(0), 0));
    JdbcUtil.getPreparedStatementSetCaller(Schema.FieldType.BYTES)
        .set(row, psMocked, 1, SchemaUtil.FieldWithIndex.of(schema.getField(1), 1));
    JdbcUtil.getPreparedStatementSetCaller(Schema.FieldType.BOOLEAN)
        .set(row, psMocked, 2, SchemaUtil.FieldWithIndex.of(schema.getField(2), 2));
    JdbcUtil.getPreparedStatementSetCaller(Schema.FieldType.STRING)
        .set(row, psMocked, 3, SchemaUtil.FieldWithIndex.of(schema.getField(3), 3));
    JdbcUtil.getPreparedStatementSetCaller(Schema.FieldType.DECIMAL)
        .set(row, psMocked, 4, SchemaUtil.FieldWithIndex.of(schema.getField(4), 4));
    JdbcUtil.getPreparedStatementSetCaller(Schema.FieldType.DOUBLE)
        .set(row, psMocked, 5, SchemaUtil.FieldWithIndex.of(schema.getField(5), 5));
    JdbcUtil.getPreparedStatementSetCaller(Schema.FieldType.FLOAT)
        .set(row, psMocked, 6, SchemaUtil.FieldWithIndex.of(schema.getField(6), 6));
    JdbcUtil.getPreparedStatementSetCaller(Schema.FieldType.INT32)
        .set(row, psMocked, 7, SchemaUtil.FieldWithIndex.of(schema.getField(7), 7));
    JdbcUtil.getPreparedStatementSetCaller(Schema.FieldType.DATETIME)
        .set(row, psMocked, 8, SchemaUtil.FieldWithIndex.of(schema.getField(8), 8));
    JdbcUtil.getPreparedStatementSetCaller(Schema.FieldType.INT16)
        .set(row, psMocked, 9, SchemaUtil.FieldWithIndex.of(schema.getField(9), 9));
    JdbcUtil.getPreparedStatementSetCaller(Schema.FieldType.BYTE)
        .set(row, psMocked, 10, SchemaUtil.FieldWithIndex.of(schema.getField(10), 10));

    verify(psMocked, times(1)).setLong(1, 42L);
    verify(psMocked, times(1)).setBytes(2, "binary".getBytes(Charset.forName("UTF-8")));
    verify(psMocked, times(1)).setBoolean(3, true);
    verify(psMocked, times(1)).setString(4, "char");
    verify(psMocked, times(1)).setBigDecimal(5, BigDecimal.valueOf(25L));
    verify(psMocked, times(1)).setDouble(6, 20.5D);
    verify(psMocked, times(1)).setFloat(7, 15.5F);
    verify(psMocked, times(1)).setInt(8, 10);
    verify(psMocked, times(1))
        .setTimestamp(9, new Timestamp(row.getDateTime("datetime_col").getMillis()));
    verify(psMocked, times(1)).setInt(10, (short) 5);
    verify(psMocked, times(1)).setByte(11, Byte.parseByte("1", 2));
  }

  @Test
  public void testGetPreparedStatementSetCallerForLogicalTypes() throws Exception {

    Schema schema =
        Schema.builder()
            .addField("logical_date_col", LogicalTypes.JDBC_DATE_TYPE)
            .addField("logical_time_col", LogicalTypes.JDBC_TIME_TYPE)
            .addField("logical_time_with_tz_col", LogicalTypes.JDBC_TIMESTAMP_WITH_TIMEZONE_TYPE)
            .build();

    long epochMilli = 1558719710000L;
    DateTime dateTime = new DateTime(epochMilli, ISOChronology.getInstanceUTC());

    Row row =
        Row.withSchema(schema)
            .addValues(
                dateTime.withTimeAtStartOfDay(), dateTime.withDate(new LocalDate(0L)), dateTime)
            .build();

    PreparedStatement psMocked = mock(PreparedStatement.class);

    JdbcUtil.getPreparedStatementSetCaller(LogicalTypes.JDBC_DATE_TYPE)
        .set(row, psMocked, 0, SchemaUtil.FieldWithIndex.of(schema.getField(0), 0));
    JdbcUtil.getPreparedStatementSetCaller(LogicalTypes.JDBC_TIME_TYPE)
        .set(row, psMocked, 1, SchemaUtil.FieldWithIndex.of(schema.getField(1), 1));
    JdbcUtil.getPreparedStatementSetCaller(LogicalTypes.JDBC_TIMESTAMP_WITH_TIMEZONE_TYPE)
        .set(row, psMocked, 2, SchemaUtil.FieldWithIndex.of(schema.getField(2), 2));

    verify(psMocked, times(1)).setDate(1, new Date(row.getDateTime(0).getMillis()));
    verify(psMocked, times(1)).setTime(2, new Time(row.getDateTime(1).getMillis()));

    Calendar cal = Calendar.getInstance(TimeZone.getTimeZone("UTC"));
    cal.setTimeInMillis(epochMilli);

    verify(psMocked, times(1)).setTimestamp(3, new Timestamp(cal.getTime().getTime()), cal);
  }

  @Test
  public void testGetPreparedStatementSetCallerForArray() throws Exception {

    Schema schema =
        Schema.builder()
            .addField("string_array_col", Schema.FieldType.array(Schema.FieldType.STRING))
            .build();

    List<String> stringList = Arrays.asList("string 1", "string 2");

    Row row = Row.withSchema(schema).addValues(stringList).build();

    PreparedStatement psMocked = mock(PreparedStatement.class);
    Connection connectionMocked = mock(Connection.class);
    Array arrayMocked = mock(Array.class);

    when(psMocked.getConnection()).thenReturn(connectionMocked);
    when(connectionMocked.createArrayOf(anyString(), any())).thenReturn(arrayMocked);

    JdbcUtil.getPreparedStatementSetCaller(Schema.FieldType.array(Schema.FieldType.STRING))
        .set(row, psMocked, 0, SchemaUtil.FieldWithIndex.of(schema.getField(0), 0));

    verify(psMocked, times(1)).setArray(1, arrayMocked);
  }

  private static ArrayList<Row> getRowsToWrite(long rowsToAdd, Schema schema) {

    ArrayList<Row> data = new ArrayList<>();
    for (int i = 0; i < rowsToAdd; i++) {
      List<Object> fields = new ArrayList<>();

      Row row =
          schema.getFields().stream()
              .map(field -> dummyFieldValue(field.getType()))
              .collect(Row.toRow(schema));
      data.add(row);
    }
    return data;
  }

  private static ArrayList<RowWithSchema> getRowsWithSchemaToWrite(long rowsToAdd) {

    ArrayList<RowWithSchema> data = new ArrayList<>();
    for (int i = 0; i < rowsToAdd; i++) {
      data.add(new RowWithSchema("Test", i));
    }
    return data;
  }

  private static Object dummyFieldValue(Schema.FieldType fieldType) {
    long epochMilli = 1558719710000L;
    if (fieldType.equals(Schema.FieldType.STRING)) {
      return "string value";
    } else if (fieldType.equals(Schema.FieldType.INT32)) {
      return 100;
    } else if (fieldType.equals(Schema.FieldType.DOUBLE)) {
      return 20.5D;
    } else if (fieldType.equals(Schema.FieldType.BOOLEAN)) {
      return Boolean.TRUE;
    } else if (fieldType.equals(Schema.FieldType.INT16)) {
      return Short.MAX_VALUE;
    } else if (fieldType.equals(Schema.FieldType.INT64)) {
      return Long.MAX_VALUE;
    } else if (fieldType.equals(Schema.FieldType.FLOAT)) {
      return 15.5F;
    } else if (fieldType.equals(Schema.FieldType.DECIMAL)) {
      return BigDecimal.ONE;
    } else if (fieldType.equals(LogicalTypes.JDBC_DATE_TYPE)) {
      return new DateTime(epochMilli, ISOChronology.getInstanceUTC()).withTimeAtStartOfDay();
    } else if (fieldType.equals(LogicalTypes.JDBC_TIME_TYPE)) {
      return new DateTime(epochMilli, ISOChronology.getInstanceUTC()).withDate(new LocalDate(0L));
    } else if (fieldType.equals(LogicalTypes.JDBC_TIMESTAMP_WITH_TIMEZONE_TYPE)) {
      return new DateTime(epochMilli, ISOChronology.getInstanceUTC());
    } else if (fieldType.equals(Schema.FieldType.DATETIME)) {
      return new DateTime(epochMilli, ISOChronology.getInstanceUTC());
    } else {
      return null;
    }
  }

  @Test
  public void testWriteWithEmptyPCollection() {
    pipeline
        .apply(Create.empty(KvCoder.of(VarIntCoder.of(), StringUtf8Coder.of())))
        .apply(
            JdbcIO.<KV<Integer, String>>write()
                .withDataSourceConfiguration(
                    JdbcIO.DataSourceConfiguration.create(
                        "org.apache.derby.jdbc.ClientDriver",
                        "jdbc:derby://localhost:" + port + "/target/beam"))
                .withStatement("insert into BEAM values(?, ?)")
                .withPreparedStatementSetter(
                    (element, statement) -> {
                      statement.setInt(1, element.getKey());
                      statement.setString(2, element.getValue());
                    }));

    pipeline.run();
  }
}
