code: /*
   * Copyright (c) 2022 Airbyte, Inc., all rights reserved.
   */
  
  package io.airbyte.integrations.source.jdbc.test;
  
  import static org.junit.jupiter.api.Assertions.assertEquals;
  
  import com.fasterxml.jackson.databind.JsonNode;
  import com.google.common.collect.ImmutableMap;
  import com.google.common.collect.Lists;
  import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
  import io.airbyte.commons.json.Jsons;
  import io.airbyte.commons.stream.MoreStreams;
  import io.airbyte.commons.string.Strings;
  import io.airbyte.db.factory.DataSourceFactory;
  import io.airbyte.db.jdbc.DefaultJdbcDatabase;
  import io.airbyte.db.jdbc.JdbcDatabase;
  import io.airbyte.db.jdbc.JdbcUtils;
  import io.airbyte.integrations.source.jdbc.AbstractJdbcSource;
  import io.airbyte.protocol.models.AirbyteCatalog;
  import io.airbyte.protocol.models.AirbyteMessage;
  import io.airbyte.protocol.models.AirbyteMessage.Type;
  import io.airbyte.protocol.models.AirbyteRecordMessage;
  import io.airbyte.protocol.models.CatalogHelpers;
  import io.airbyte.protocol.models.ConfiguredAirbyteCatalog;
  import io.airbyte.protocol.models.ConfiguredAirbyteStream;
  import io.airbyte.protocol.models.DestinationSyncMode;
  import io.airbyte.protocol.models.Field;
  import io.airbyte.protocol.models.JsonSchemaType;
  import io.airbyte.protocol.models.SyncMode;
  import java.math.BigDecimal;
  import java.nio.ByteBuffer;
  import java.util.ArrayList;
  import java.util.BitSet;
  import java.util.Collections;
  import java.util.Iterator;
  import java.util.List;
  import java.util.Optional;
  import org.junit.jupiter.api.Test;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  /**
   * Runs a "large" amount of data through a JdbcSource to ensure that it streams / chunks records.
   */
  // todo (cgardens) - this needs more love and thought. we should be able to test this without having
  // to rewrite so much data. it is enough for now to sanity check that our JdbcSources can actually
  // handle more data than fits in memory.
  @SuppressFBWarnings(
                      value = {"MS_SHOULD_BE_FINAL"},
                      justification = "The static variables are updated in sub classes for convenience, and cannot be final.")
  public abstract class JdbcStressTest {
  
    private static final Logger LOGGER = LoggerFactory.getLogger(JdbcStressTest.class);
  
    // this will get rounded down to the nearest 1000th.
    private static final long TOTAL_RECORDS = 10_000_000L;
    private static final int BATCH_SIZE = 1000;
    public static String TABLE_NAME = "id_and_name";
    public static String COL_ID = "id";
    public static String COL_NAME = "name";
    public static String COL_ID_TYPE = "BIGINT";
    public static String INSERT_STATEMENT = "(%s,'picard-%s')";
  
    private static String streamName;
  
    private BitSet bitSet;
    private JsonNode config;
    private AbstractJdbcSource<?> source;
  
    /**
     * These tests write records without specifying a namespace (schema name). They will be written into
     * whatever the default schema is for the database. When they are discovered they will be namespaced
     * by the schema name (e.g. <default-schema-name>.<table_name>). Thus the source needs to tell the
     * tests what that default schema name is. If the database does not support schemas, then database
     * name should used instead.
     *
     * @return name that will be used to namespace the record.
     */
    public abstract Optional<String> getDefaultSchemaName();
  
    /**
     * A valid configuration to connect to a test database.
     *
     * @return config
     */
    public abstract JsonNode getConfig();
  
    /**
     * Full qualified class name of the JDBC driver for the database.
     *
     * @return driver
     */
    public abstract String getDriverClass();
  
    /**
     * An instance of the source that should be tests.
     *
     * @return source
     */
    public abstract AbstractJdbcSource<?> getSource();
  
    protected String createTableQuery(final String tableName, final String columnClause) {
      return String.format("CREATE TABLE %s(%s)",
          tableName, columnClause);
    }
  
    public void setup() throws Exception {
      LOGGER.info("running for driver:" + getDriverClass());
      bitSet = new BitSet((int) TOTAL_RECORDS);
  
      source = getSource();
      streamName = getDefaultSchemaName().map(val -> val + "." + TABLE_NAME).orElse(TABLE_NAME);
      config = getConfig();
  
      final JsonNode jdbcConfig = source.toDatabaseConfig(config);
      final JdbcDatabase database = new DefaultJdbcDatabase(
          DataSourceFactory.create(
              jdbcConfig.get(JdbcUtils.USERNAME_KEY).asText(),
              jdbcConfig.has(JdbcUtils.PASSWORD_KEY) ? jdbcConfig.get(JdbcUtils.PASSWORD_KEY).asText() : null,
              getDriverClass(),
              jdbcConfig.get(JdbcUtils.JDBC_URL_KEY).asText()));
  
      database.execute(connection -> connection.createStatement().execute(
          createTableQuery("id_and_name", String.format("id %s, name VARCHAR(200)", COL_ID_TYPE))));
      final long batchCount = TOTAL_RECORDS / BATCH_SIZE;
      LOGGER.info("writing {} batches of {}", batchCount, BATCH_SIZE);
      for (int i = 0; i < batchCount; i++) {
        if (i % 1000 == 0)
          LOGGER.info("writing batch: " + i);
        final List<String> insert = new ArrayList<>();
        for (int j = 0; j < BATCH_SIZE; j++) {
          final int recordNumber = (i * BATCH_SIZE) + j;
          insert.add(String.format(INSERT_STATEMENT, recordNumber, recordNumber));
        }
  
        final String sql = prepareInsertStatement(insert);
        database.execute(connection -> connection.createStatement().execute(sql));
      }
  
    }
  
    // todo (cgardens) - restructure these tests so that testFullRefresh() and testIncremental() can be
    // separate tests. current constrained by only wanting to setup the fixture in the database once,
    // but it is not trivial to move them to @BeforeAll because it is static and we are doing
    // inheritance. Not impossible, just needs to be done thoughtfully and for all JdbcSources.
    @Test
    public void stressTest() throws Exception {
      testFullRefresh();
      testIncremental();
    }
  
    private void testFullRefresh() throws Exception {
      runTest(getConfiguredCatalogFullRefresh(), "full_refresh");
    }
  
    private void testIncremental() throws Exception {
      runTest(getConfiguredCatalogIncremental(), "incremental");
    }
  
    private void runTest(final ConfiguredAirbyteCatalog configuredCatalog, final String testName) throws Exception {
      LOGGER.info("running stress test for: " + testName);
      final Iterator<AirbyteMessage> read = source.read(config, configuredCatalog, Jsons.jsonNode(Collections.emptyMap()));
      final long actualCount = MoreStreams.toStream(read)
          .filter(m -> m.getType() == Type.RECORD)
          .peek(m -> {
            if (m.getRecord().getData().get(COL_ID).asLong() % 100000 == 0) {
              LOGGER.info("reading batch: " + m.getRecord().getData().get(COL_ID).asLong() / 1000);
            }
          })
          .peek(m -> assertExpectedMessage(m))
          .count();
      ByteBuffer a;
      final long expectedRoundedRecordsCount = TOTAL_RECORDS - TOTAL_RECORDS % 1000;
      LOGGER.info("expected records count: " + TOTAL_RECORDS);
      LOGGER.info("actual records count: " + actualCount);
      assertEquals(expectedRoundedRecordsCount, actualCount, "testing: " + testName);
      assertEquals(expectedRoundedRecordsCount, bitSet.cardinality(), "testing: " + testName);
    }
  
    // each is roughly 106 bytes.
    private void assertExpectedMessage(final AirbyteMessage actualMessage) {
      final long recordNumber = actualMessage.getRecord().getData().get(COL_ID).asLong();
      bitSet.set((int) recordNumber);
      actualMessage.getRecord().setEmittedAt(null);
  
      final Number expectedRecordNumber =
          getDriverClass().toLowerCase().contains("oracle") ? new BigDecimal(recordNumber)
              : recordNumber;
  
      final AirbyteMessage expectedMessage = new AirbyteMessage().withType(Type.RECORD)
          .withRecord(new AirbyteRecordMessage().withStream(streamName)
              .withData(Jsons.jsonNode(
                  ImmutableMap.of(COL_ID, expectedRecordNumber, COL_NAME, "picard-" + recordNumber))));
      assertEquals(expectedMessage, actualMessage);
    }
  
    private static ConfiguredAirbyteCatalog getConfiguredCatalogFullRefresh() {
      return CatalogHelpers.toDefaultConfiguredCatalog(getCatalog());
    }
  
    private static ConfiguredAirbyteCatalog getConfiguredCatalogIncremental() {
      return new ConfiguredAirbyteCatalog()
          .withStreams(Collections.singletonList(new ConfiguredAirbyteStream().withStream(getCatalog().getStreams().get(0))
              .withCursorField(Collections.singletonList(COL_ID))
              .withSyncMode(SyncMode.INCREMENTAL)
              .withDestinationSyncMode(DestinationSyncMode.APPEND)));
    }
  
    private static AirbyteCatalog getCatalog() {
      return new AirbyteCatalog().withStreams(Lists.newArrayList(CatalogHelpers.createAirbyteStream(
          streamName,
          Field.of(COL_ID, JsonSchemaType.NUMBER),
          Field.of(COL_NAME, JsonSchemaType.STRING))
          .withSupportedSyncModes(Lists.newArrayList(SyncMode.FULL_REFRESH, SyncMode.INCREMENTAL))));
    }
  
    private String prepareInsertStatement(final List<String> inserts) {
      if (getDriverClass().toLowerCase().contains("oracle")) {
        return String.format("INSERT ALL %s SELECT * FROM dual", Strings.join(inserts, " "));
      }
      return String.format("INSERT INTO id_and_name (id, name) VALUES %s", Strings.join(inserts, ", "));
    }
  
  }
  