content: /*
   * Copyright 2020 Confluent Inc.
   *
   * Licensed under the Confluent Community License (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.confluent.io/confluent-community-license
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
   * WARRANTIES OF ANY KIND, either express or implied.  See the License for the
   * specific language governing permissions and limitations under the License.
   */
  
  package io.confluent.connect.jdbc.sink.integration;
  
  import java.sql.Array;
  import java.sql.Connection;
  import java.sql.ResultSet;
  import java.sql.SQLException;
  import java.sql.Statement;
  
  import java.util.Arrays;
  import java.util.Collections;
  import java.util.Map;
  import java.util.UUID;
  import java.util.concurrent.TimeUnit;
  
  import io.confluent.common.utils.IntegrationTest;
  import io.confluent.connect.jdbc.integration.BaseConnectorIT;
  import io.confluent.connect.jdbc.sink.JdbcSinkConfig;
  
  import io.zonky.test.db.postgres.junit.EmbeddedPostgresRules;
  import io.zonky.test.db.postgres.junit.SingleInstancePostgresRule;
  import org.apache.kafka.clients.consumer.ConsumerRecords;
  import org.apache.kafka.connect.data.Schema;
  import org.apache.kafka.connect.data.SchemaBuilder;
  import org.apache.kafka.connect.data.Struct;
  import org.apache.kafka.connect.json.JsonConverter;
  import org.apache.kafka.connect.runtime.errors.ToleranceType;
  import org.junit.After;
  import org.junit.Before;
  import org.junit.Rule;
  import org.junit.Test;
  import org.junit.experimental.categories.Category;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import static io.confluent.connect.jdbc.sink.JdbcSinkConfig.MAX_RETRIES;
  import static org.apache.kafka.connect.runtime.ConnectorConfig.ERRORS_TOLERANCE_CONFIG;
  import static org.apache.kafka.connect.runtime.SinkConnectorConfig.DLQ_TOPIC_NAME_CONFIG;
  import static org.apache.kafka.connect.runtime.SinkConnectorConfig.DLQ_TOPIC_REPLICATION_FACTOR_CONFIG;
  import static org.junit.Assert.assertEquals;
  import static org.junit.Assert.assertNotNull;
  import static org.junit.Assert.assertTrue;
  
  
  /**
   * Integration tests for writing to Postgres with UUID columns.
   */
  @Category(IntegrationTest.class)
  public class PostgresDatatypeIT extends BaseConnectorIT {
  
    private static final Logger LOG = LoggerFactory.getLogger(PostgresDatatypeIT.class);
  
    @Rule
    public SingleInstancePostgresRule pg = EmbeddedPostgresRules.singleInstance();
  
    private String tableName;
    private JsonConverter jsonConverter;
    private Map<String, String> props;
  
    @Before
    public void before() {
      startConnect();
      jsonConverter = jsonConverter();
      props = baseSinkProps();
  
      tableName = "test";
      String jdbcURL = String
          .format("jdbc:postgresql://localhost:%s/postgres", pg.getEmbeddedPostgres().getPort());
      props.put(JdbcSinkConfig.CONNECTION_URL, jdbcURL);
      props.put(JdbcSinkConfig.CONNECTION_USER, "postgres");
      props.put("pk.mode", "none");
      props.put("topics", tableName);
  
      // create topic in Kafka
      connect.kafka().createTopic(tableName, 1);
    }
  
    @After
    public void after() throws SQLException {
      try (Connection c = pg.getEmbeddedPostgres().getPostgresDatabase().getConnection()) {
        try (Statement s = c.createStatement()) {
          s.execute("DROP TABLE IF EXISTS " + tableName);
        }
        LOG.info("Dropped table");
      } finally {
        pg = null;
        stopConnect();
      }
    }
  
    /**
     * Verifies that even when the connector encounters exceptions that would cause a connection
     * with an invalid transaction, the connector sends only the errant record to the error
     * reporter and establishes a valid transaction for subsequent correct records to be sent to
     * the actual database.
     */
    @Test
    public void testPrimaryKeyConstraintsSendsToErrorReporter() throws Exception {
      props.put(ERRORS_TOLERANCE_CONFIG, ToleranceType.ALL.value());
      props.put(DLQ_TOPIC_NAME_CONFIG, DLQ_TOPIC_NAME);
      props.put(DLQ_TOPIC_REPLICATION_FACTOR_CONFIG, "1");
      props.put(MAX_RETRIES, "0");
  
      createTableWithPrimaryKey();
      connect.configureConnector("jdbc-sink-connector", props);
      waitForConnectorToStart("jdbc-sink-connector", 1);
  
      final Schema schema = SchemaBuilder.struct().name("com.example.Person")
          .field("firstname", Schema.STRING_SCHEMA)
          .field("lastname", Schema.STRING_SCHEMA)
          .build();
      final Struct firstStruct = new Struct(schema)
          .put("firstname", "Christina")
          .put("lastname", "Brams");
  
      produceRecord(schema, firstStruct);
      // Send the same record for a PK collision
      produceRecord(schema, firstStruct);
  
      // Now, create and send another normal record
      Struct secondStruct = new Struct(schema)
          .put("firstname", "Brams")
          .put("lastname", "Christina");
  
      produceRecord(schema, secondStruct);
  
      waitForCommittedRecords("jdbc-sink-connector", Collections.singleton(tableName), 3, 1,
          TimeUnit.MINUTES.toMillis(3));
  
      ConsumerRecords<byte[], byte[]> records = connect.kafka().consume(1, CONSUME_MAX_DURATION_MS,
          DLQ_TOPIC_NAME);
  
      assertEquals(1, records.count());
    }
  
    @Test
    public void testRecordSchemaMoreFieldsThanTableSendsToErrorReporter() throws Exception {
      props.put(ERRORS_TOLERANCE_CONFIG, ToleranceType.ALL.value());
      props.put(DLQ_TOPIC_NAME_CONFIG, DLQ_TOPIC_NAME);
      props.put(DLQ_TOPIC_REPLICATION_FACTOR_CONFIG, "1");
  
      createTableWithLessFields();
      connect.configureConnector("jdbc-sink-connector", props);
      waitForConnectorToStart("jdbc-sink-connector", 1);
  
      final Schema schema = SchemaBuilder.struct().name("com.example.Person")
          .field("firstname", Schema.STRING_SCHEMA)
          .field("lastname", Schema.STRING_SCHEMA)
          .field("jsonid", Schema.STRING_SCHEMA)
          .field("userid", Schema.STRING_SCHEMA)
          .build();
      final Struct struct = new Struct(schema)
          .put("firstname", "Christina")
          .put("lastname", "Brams")
          .put("jsonid", "5")
          .put("userid", UUID.randomUUID().toString());
  
      produceRecord(schema, struct);
  
      waitForCommittedRecords("jdbc-sink-connector", Collections.singleton(tableName), 1, 1,
          TimeUnit.MINUTES.toMillis(2));
  
      ConsumerRecords<byte[], byte[]> records = connect.kafka().consume(1, CONSUME_MAX_DURATION_MS,
          DLQ_TOPIC_NAME);
  
      assertEquals(1, records.count());
    }
  
    @Test
    public void testWriteToTableWithUuidColumn() throws Exception {
      createTableWithUuidColumns();
      connect.configureConnector("jdbc-sink-connector", props);
      waitForConnectorToStart("jdbc-sink-connector", 1);
  
      final Schema schema = SchemaBuilder.struct().name("com.example.Person")
                                         .field("firstname", Schema.STRING_SCHEMA)
                                         .field("lastname", Schema.STRING_SCHEMA)
                                         .field("jsonid", Schema.STRING_SCHEMA)
                                         .field("userid", Schema.STRING_SCHEMA)
                                         .build();
      final Struct struct = new Struct(schema)
          .put("firstname", "Christina")
          .put("lastname", "Brams")
          .put("jsonid", "5")
          .put("userid", UUID.randomUUID().toString());
  
      produceRecord(schema, struct);
  
      waitForCommittedRecords("jdbc-sink-connector", Collections.singleton(tableName), 1, 1,
          TimeUnit.MINUTES.toMillis(2));
  
      try (Connection c = pg.getEmbeddedPostgres().getPostgresDatabase().getConnection()) {
        try (Statement s = c.createStatement()) {
          try (ResultSet rs = s.executeQuery("SELECT * FROM " + tableName)) {
            assertTrue(rs.next());
            assertEquals(struct.getString("firstname"), rs.getString("firstname"));
            assertEquals(struct.getString("lastname"), rs.getString("lastname"));
            assertEquals(struct.getString("jsonid"), rs.getString("jsonid"));
            assertEquals(struct.getString("userid"), rs.getString("userid"));
          }
        }
      }
    }
  
    @Test
    public void testWriteToTableWithIntArrayColumn() throws SQLException, InterruptedException {
      createTableWithIntArrayColumns();
      connect.configureConnector("jdbc-sink-connector", props);
      waitForConnectorToStart("jdbc-sink-connector", 1);
  
      final Schema schema = SchemaBuilder.struct().name("com.example.Person")
          .field("firstname", Schema.STRING_SCHEMA)
          .field("lastname", Schema.STRING_SCHEMA)
          .field("friends", SchemaBuilder.array(Schema.INT32_SCHEMA).build())
          .field("friendnames", SchemaBuilder.array(Schema.STRING_SCHEMA).build())
          .build();
  
      final Struct struct = new Struct(schema)
          .put("firstname", "Christina")
          .put("lastname", "Brams")
          .put("friends", Arrays.asList(10, 6221))
          .put("friendnames", Arrays.asList("Lucas", "Tom"));
      produceRecord(schema, struct);
  
      waitForCommittedRecords("jdbc-sink-connector", Collections.singleton(tableName), 1, 1,
          TimeUnit.MINUTES.toMillis(2));
      try (Connection c = pg.getEmbeddedPostgres().getPostgresDatabase().getConnection()) {
        try (Statement s = c.createStatement()) {
          try (ResultSet rs = s.executeQuery("SELECT * FROM " + tableName)) {
            assertTrue(rs.next());
            assertEquals(struct.getString("firstname"), rs.getString("firstname"));
            assertEquals(struct.getString("lastname"), rs.getString("lastname"));
            assertJDBCArray(rs, "friends", struct);
            assertJDBCArray(rs, "friendnames", struct);
          }
        }
      }
    }
  
    @Test
    public void testWriteToTableWithIntArrayColumnMissingFields() throws SQLException, InterruptedException {
      createTableWithIntArrayColumnsMissing();
      props.put(JdbcSinkConfig.AUTO_EVOLVE, "true");
      connect.configureConnector("jdbc-sink-connector", props);
      waitForConnectorToStart("jdbc-sink-connector", 1);
  
      final Schema schema = SchemaBuilder.struct().name("com.example.Person")
          .field("firstname", Schema.STRING_SCHEMA)
          .field("lastname", Schema.STRING_SCHEMA)
          .field("friends", SchemaBuilder.array(Schema.INT32_SCHEMA).optional().build())
          .field("friendnames", SchemaBuilder.array(Schema.STRING_SCHEMA).optional().build())
          .build();
  
      final Struct struct = new Struct(schema)
          .put("firstname", "Christina")
          .put("lastname", "Brams")
          .put("friends", Arrays.asList(10, 6221))
          .put("friendnames", Arrays.asList("Lucas", "Tom"));
      produceRecord(schema, struct);
  
      waitForCommittedRecords("jdbc-sink-connector", Collections.singleton(tableName), 1, 1,
          TimeUnit.MINUTES.toMillis(2));
      try (Connection c = pg.getEmbeddedPostgres().getPostgresDatabase().getConnection()) {
        try (Statement s = c.createStatement()) {
          try (ResultSet rs = s.executeQuery("SELECT * FROM " + tableName)) {
            assertTrue(rs.next());
            assertEquals(struct.getString("firstname"), rs.getString("firstname"));
            assertEquals(struct.getString("lastname"), rs.getString("lastname"));
            assertJDBCArray(rs, "friends", struct);
            assertJDBCArray(rs, "friendnames", struct);
          }
        }
      }
    }
  
    private void assertJDBCArray(ResultSet rs, String fieldName, Struct struct) throws SQLException {
      Array array = rs.getArray(fieldName);
      assertNotNull(array);
      assertEquals(struct.getArray(fieldName), Arrays.asList((Object[])array.getArray()));
    }
  
    private void createTableWithIntArrayColumnsMissing() throws SQLException {
      LOG.info("Creating table {} with UUID column", tableName);
      try (Connection c = pg.getEmbeddedPostgres().getPostgresDatabase().getConnection()) {
        c.setAutoCommit(false);
        try (Statement s = c.createStatement()) {
          String sql = String.format(
              "CREATE TABLE %s(firstName TEXT, lastName TEXT, jsonid json)",
              tableName
          );
          LOG.info("Executing statement: {}", sql);
          s.execute(sql);
          c.commit();
        }
      }
      LOG.info("Created table {} with UUID column", tableName);
    }
  
    private void createTableWithIntArrayColumns() throws SQLException {
      LOG.info("Creating table {} with UUID column", tableName);
      try (Connection c = pg.getEmbeddedPostgres().getPostgresDatabase().getConnection()) {
        c.setAutoCommit(false);
        try (Statement s = c.createStatement()) {
          String sql = String.format(
              "CREATE TABLE %s(firstName TEXT, lastName TEXT, jsonid json, friends int[], friendnames text[])",
              tableName
          );
          LOG.info("Executing statement: {}", sql);
          s.execute(sql);
          c.commit();
        }
      }
      LOG.info("Created table {} with UUID column", tableName);
    }
  
    private void createTable(String columnsSql) throws SQLException {
      try (Connection c = pg.getEmbeddedPostgres().getPostgresDatabase().getConnection()) {
        c.setAutoCommit(false);
        try (Statement s = c.createStatement()) {
          String sql = String.format(
              columnsSql,
              tableName
          );
          LOG.info("Executing statement: {}", sql);
          s.execute(sql);
          c.commit();
        }
      }
    }
  
    private void createTableWithUuidColumns() throws SQLException {
      LOG.info("Creating table {} with UUID column", tableName);
      createTable("CREATE TABLE %s(firstName TEXT, lastName TEXT, jsonid json, userid UUID)");
      LOG.info("Created table {} with UUID column", tableName);
    }
  
    private void createTableWithLessFields() throws SQLException {
      LOG.info("Creating table {} with less fields", tableName);
      createTable("CREATE TABLE %s(firstName TEXT, jsonid json, userid UUID)");
      LOG.info("Created table {} with less fields", tableName);
    }
  
    private void createTableWithPrimaryKey() throws SQLException {
      LOG.info("Creating table {} with a primary key", tableName);
      createTable("CREATE TABLE %s(firstName TEXT PRIMARY KEY, lastName TEXT)");
      LOG.info("Created table {} with a primary key", tableName);
    }
  
    private void produceRecord(Schema schema, Struct struct) {
      String kafkaValue = new String(jsonConverter.fromConnectData(tableName, schema, struct));
      connect.kafka().produce(tableName, null, kafkaValue);
    }
  }
  