Source MySQL\MsSql\Postgres: added RDS base performance tests

airbyte-integrations/bases/standard-source-test/build.gradle

@@ -22,6 +22,7 @@ dependencies {
     runtimeOnly 'org.junit.jupiter:junit-jupiter-engine:5.4.2'
     implementation 'org.junit.platform:junit-platform-launcher:1.7.0'
     implementation 'org.junit.jupiter:junit-jupiter-api:5.4.2'
+    implementation 'org.junit.jupiter:junit-jupiter-params:5.8.1'
 }
 
 def getFullPath(String className) {

airbyte-integrations/bases/standard-source-test/src/main/java/io/airbyte/integrations/standardtest/source/AbstractSourceConnectorTest.java

@@ -9,6 +9,7 @@
 
 import com.fasterxml.jackson.databind.JsonNode;
 import io.airbyte.config.JobGetSpecConfig;
+import io.airbyte.config.ResourceRequirements;
 import io.airbyte.config.StandardCheckConnectionInput;
 import io.airbyte.config.StandardCheckConnectionOutput;
 import io.airbyte.config.StandardDiscoverCatalogInput;
@@ -24,6 +25,7 @@
 import io.airbyte.workers.DefaultDiscoverCatalogWorker;
 import io.airbyte.workers.DefaultGetSpecWorker;
 import io.airbyte.workers.WorkerException;
+import io.airbyte.workers.WorkerUtils;
 import io.airbyte.workers.process.AirbyteIntegrationLauncher;
 import io.airbyte.workers.process.DockerProcessFactory;
 import io.airbyte.workers.process.ProcessFactory;
@@ -33,18 +35,22 @@
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 /**
  * This abstract class contains helpful functionality and boilerplate for testing a source
  * connector.
  */
 public abstract class AbstractSourceConnectorTest {
 
+  protected static final Logger LOGGER = LoggerFactory.getLogger(AbstractSourceConnectorTest.class);
   private TestDestinationEnv environment;
   private Path jobRoot;
   protected Path localRoot;
@@ -53,6 +59,11 @@ public abstract class AbstractSourceConnectorTest {
   private static final String JOB_ID = String.valueOf(0L);
   private static final int JOB_ATTEMPT = 0;
 
+  private static final String CPU_REQUEST_FIELD_NAME = "cpuRequest";
+  private static final String CPU_LIMIT_FIELD_NAME = "cpuLimit";
+  private static final String MEMORY_REQUEST_FIELD_NAME = "memoryRequest";
+  private static final String MEMORY_LIMIT_FIELD_NAME = "memoryLimit";
+
   /**
    * Name of the docker image that the tests will run against.
    *
@@ -169,7 +180,9 @@ protected Map<String, Integer> runReadVerifyNumberOfReceivedMsgs(final Configure
         .withState(state == null ? null : new State().withState(state))
         .withCatalog(catalog);
 
-    final AirbyteSource source = new DefaultAirbyteSource(new AirbyteIntegrationLauncher(JOB_ID, JOB_ATTEMPT, getImageName(), processFactory));
+    final Map<String, String> mapOfResourceRequirementsParams = prepareResourceRequestMapBySystemProperties();
+    final AirbyteSource source =
+        prepareAirbyteSource(!mapOfResourceRequirementsParams.isEmpty() ? prepareResourceRequirements(mapOfResourceRequirementsParams) : null);
     source.start(sourceConfig, jobRoot);
 
     while (!source.isFinished()) {
@@ -186,4 +199,36 @@ protected Map<String, Integer> runReadVerifyNumberOfReceivedMsgs(final Configure
     return mapOfExpectedRecordsCount;
   }
 
+  protected ResourceRequirements prepareResourceRequirements(Map<String, String> mapOfResourceRequirementsParams) {
+    return new ResourceRequirements().withCpuRequest(mapOfResourceRequirementsParams.get(CPU_REQUEST_FIELD_NAME))
+        .withCpuLimit(mapOfResourceRequirementsParams.get(CPU_LIMIT_FIELD_NAME))
+        .withMemoryRequest(mapOfResourceRequirementsParams.get(MEMORY_REQUEST_FIELD_NAME))
+        .withMemoryLimit(mapOfResourceRequirementsParams.get(MEMORY_LIMIT_FIELD_NAME));
+  }
+
+  private AirbyteSource prepareAirbyteSource(ResourceRequirements resourceRequirements) {
+    var integrationLauncher = resourceRequirements == null ? new AirbyteIntegrationLauncher(JOB_ID, JOB_ATTEMPT, getImageName(), processFactory)
+        : new AirbyteIntegrationLauncher(JOB_ID, JOB_ATTEMPT, getImageName(), processFactory, resourceRequirements);
+    return new DefaultAirbyteSource(integrationLauncher);
+  }
+
+  private static Map<String, String> prepareResourceRequestMapBySystemProperties() {
+    var cpuLimit = System.getProperty(CPU_LIMIT_FIELD_NAME);
+    var memoryLimit = System.getProperty(MEMORY_LIMIT_FIELD_NAME);
+    if (cpuLimit.isBlank() || cpuLimit.isEmpty()) {
+      cpuLimit = WorkerUtils.DEFAULT_RESOURCE_REQUIREMENTS.getCpuLimit();
+    }
+    if (memoryLimit.isBlank() || memoryLimit.isEmpty()) {
+      memoryLimit = WorkerUtils.DEFAULT_RESOURCE_REQUIREMENTS.getMemoryLimit();
+    }
+    LOGGER.error("cpu limit -->> {}", cpuLimit);
+    LOGGER.error("memory limit -->> {}", memoryLimit);
+    Map<String, String> result = new HashMap<>();
+    result.put(CPU_REQUEST_FIELD_NAME, WorkerUtils.DEFAULT_RESOURCE_REQUIREMENTS.getCpuRequest());
+    result.put(CPU_LIMIT_FIELD_NAME, cpuLimit);
+    result.put(MEMORY_REQUEST_FIELD_NAME, WorkerUtils.DEFAULT_RESOURCE_REQUIREMENTS.getMemoryRequest());
+    result.put(MEMORY_LIMIT_FIELD_NAME, memoryLimit);
+    return result;
+  }
+
 }

airbyte-integrations/bases/standard-source-test/src/main/java/io/airbyte/integrations/standardtest/source/performancetest/AbstractSourceBasePerformanceTest.java

@@ -0,0 +1,44 @@
+/*
+ * Copyright (c) 2021 Airbyte, Inc., all rights reserved.
+ */
+
+package io.airbyte.integrations.standardtest.source.performancetest;
+
+import io.airbyte.integrations.standardtest.source.AbstractSourceConnectorTest;
+import io.airbyte.integrations.standardtest.source.TestDestinationEnv;
+
+/**
+ * This abstract class contains common methods for both steams - Fill Db scripts and Performance
+ * tests.
+ */
+public abstract class AbstractSourceBasePerformanceTest extends AbstractSourceConnectorTest {
+
+  private static final String TEST_COLUMN_NAME = "test_column";
+  private static final String TEST_STREAM_NAME_TEMPLATE = "test_%S";
+
+  /**
+   * The column name will be used for a test column in the test tables. Override it if default name is
+   * not valid for your source.
+   *
+   * @return Test column name
+   */
+  protected String getTestColumnName() {
+    return TEST_COLUMN_NAME;
+  }
+
+  /**
+   * The stream name template will be used for a test tables. Override it if default name is not valid
+   * for your source.
+   *
+   * @return Test steam name template
+   */
+  protected String getTestStreamNameTemplate() {
+    return TEST_STREAM_NAME_TEMPLATE;
+  }
+
+  @Override
+  protected void setupEnvironment(final TestDestinationEnv environment) throws Exception {
+    // DO NOTHING. Mandatory to override. DB will be setup as part of each test
+  }
+
+}

airbyte-integrations/bases/standard-source-test/src/main/java/io/airbyte/integrations/standardtest/source/performancetest/AbstractSourceFillDbWithTestData.java

@@ -0,0 +1,138 @@
+/*
+ * Copyright (c) 2021 Airbyte, Inc., all rights reserved.
+ */
+
+package io.airbyte.integrations.standardtest.source.performancetest;
+
+import io.airbyte.db.Database;
+import java.util.StringJoiner;
+import java.util.stream.Stream;
+import org.junit.jupiter.api.Disabled;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.MethodSource;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * This abstract class contains common methods for Fill Db scripts.
+ */
+public abstract class AbstractSourceFillDbWithTestData extends AbstractSourceBasePerformanceTest {
+
+  private static final String CREATE_DB_TABLE_TEMPLATE = "CREATE TABLE %s.%s(id INTEGER PRIMARY KEY, %s)";
+  private static final String INSERT_INTO_DB_TABLE_QUERY_TEMPLATE = "INSERT INTO %s.%s (%s) VALUES %s";
+  private static final String TEST_DB_FIELD_TYPE = "varchar(10)";
+
+  protected static final Logger c = LoggerFactory.getLogger(AbstractSourceFillDbWithTestData.class);
+  private static final String TEST_VALUE_TEMPLATE_POSTGRES = "\'Value id_placeholder\'";
+  protected static Stream testArgs;
+
+  /**
+   * Setup the test database. All tables and data described in the registered tests will be put there.
+   *
+   * @return configured test database
+   * @throws Exception - might throw any exception during initialization.
+   */
+  protected abstract Database setupDatabase(String dbName) throws Exception;
+
+  /**
+   * The test added test data to a new DB. 1. Set DB creds in static variables above 2. Set desired
+   * number for streams, coolumns and records 3. Run the test
+   */
+  @Disabled
+  @ParameterizedTest
+  @MethodSource("provideParameters")
+  public void addTestData(String dbName,
+                          String schemaName,
+                          int numberOfDummyRecords,
+                          int numberOfBatches,
+                          int numberOfColumns,
+                          int numberOfStreams)
+      throws Exception {
+
+    final Database database = setupDatabase(dbName);
+
+    database.query(ctx -> {
+      for (int currentSteamNumber = 0; currentSteamNumber < numberOfStreams; currentSteamNumber++) {
+
+        String currentTableName = String.format(getTestStreamNameTemplate(), currentSteamNumber);
+
+        ctx.fetch(prepareCreateTableQuery(schemaName, numberOfColumns, currentTableName));
+        for (int i = 0; i < numberOfBatches; i++) {
+          String insertQueryTemplate = prepareInsertQueryTemplate(schemaName, i,
+              numberOfColumns,
+              numberOfDummyRecords);
+          ctx.fetch(String.format(insertQueryTemplate, currentTableName));
+        }
+
+        c.info("Finished processing for stream " + currentSteamNumber);
+      }
+      return null;
+    });
+
+    database.close();
+
+  }
+
+  /**
+   * This is a data provider for fill DB script,, Each argument's group would be ran as a separate
+   * test. Set the "testArgs" in test class of your DB in @BeforeTest method.
+   *
+   * 1st arg - a name of DB that will be used in jdbc connection string. 2nd arg - a schemaName that
+   * will be ised as a NameSpace in Configured Airbyte Catalog. 3rd arg - a number of expected records
+   * retrieved in each stream. 4th arg - a number of columns in each stream\table that will be use for
+   * Airbyte Cataloq configuration 5th arg - a number of streams to read in configured airbyte
+   * Catalog. Each stream\table in DB should be names like "test_0", "test_1",..., test_n.
+   *
+   * Stream.of( Arguments.of("your_db_name", "your_schema_name", 100, 2, 240, 1000) );
+   */
+  private static Stream<Arguments> provideParameters() {
+    return testArgs;
+  }
+
+  protected String prepareCreateTableQuery(final String dbSchemaName,
+                                           final int numberOfColumns,
+                                           final String currentTableName) {
+
+    StringJoiner sj = new StringJoiner(",");
+    for (int i = 0; i < numberOfColumns; i++) {
+      sj.add(String.format(" %s%s %s", getTestColumnName(), i, TEST_DB_FIELD_TYPE));
+    }
+
+    return String.format(CREATE_DB_TABLE_TEMPLATE, dbSchemaName, currentTableName, sj.toString());
+  }
+
+  protected String prepareInsertQueryTemplate(final String dbSchemaName,
+                                              final int batchNumber,
+                                              final int numberOfColumns,
+                                              final int recordsNumber) {
+
+    StringJoiner fieldsNames = new StringJoiner(",");
+    fieldsNames.add("id");
+
+    StringJoiner baseInsertQuery = new StringJoiner(",");
+    baseInsertQuery.add("id_placeholder");
+
+    for (int i = 0; i < numberOfColumns; i++) {
+      fieldsNames.add(getTestColumnName() + i);
+      baseInsertQuery.add(TEST_VALUE_TEMPLATE_POSTGRES);
+    }
+
+    StringJoiner insertGroupValuesJoiner = new StringJoiner(",");
+
+    int batchMessages = batchNumber * 100;
+
+    for (int currentRecordNumber = batchMessages;
+        currentRecordNumber < recordsNumber + batchMessages;
+        currentRecordNumber++) {
+      insertGroupValuesJoiner
+          .add("(" + baseInsertQuery.toString()
+              .replaceAll("id_placeholder", String.valueOf(currentRecordNumber)) + ")");
+    }
+
+    return String
+        .format(INSERT_INTO_DB_TABLE_QUERY_TEMPLATE, dbSchemaName, "%s", fieldsNames.toString(),
+            insertGroupValuesJoiner.toString());
+  }
+
+}