com.scalar.db.api.DistributedStorageSecondaryIndexIntegrationTestBase Maven / Gradle / Ivy

Go to download
package com.scalar.db.api;

import static org.assertj.core.api.Assertions.assertThat;

import com.scalar.db.exception.storage.ExecutionException;
import com.scalar.db.io.DataType;
import com.scalar.db.io.Key;
import com.scalar.db.io.Value;
import com.scalar.db.service.StorageFactory;
import com.scalar.db.util.TestUtils;
import java.io.IOException;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Properties;
import java.util.Random;
import java.util.Set;
import org.junit.jupiter.api.AfterAll;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.TestInstance;

@TestInstance(TestInstance.Lifecycle.PER_CLASS)
public abstract class DistributedStorageSecondaryIndexIntegrationTestBase {

  private static final String TEST_NAME = "storage_secondary_idx";
  private static final String NAMESPACE = "int_test_" + TEST_NAME;
  private static final String PARTITION_KEY = "pkey";
  private static final String INDEX_COL_NAME = "idx_col";
  private static final String COL_NAME = "col";

  private static final int ATTEMPT_COUNT = 50;
  private static final int DATA_NUM = 10;

  private static final Random random = new Random();

  private DistributedStorageAdmin admin;
  private DistributedStorage storage;
  private String namespace;
  private Set secondaryIndexTypes;

  private long seed;

  @BeforeAll
  public void beforeAll() throws Exception {
    initialize(TEST_NAME);
    StorageFactory factory = StorageFactory.create(getProperties(TEST_NAME));
    admin = factory.getAdmin();
    namespace = getNamespace();
    secondaryIndexTypes = getSecondaryIndexTypes();
    createTables();
    storage = factory.getStorage();
    seed = System.currentTimeMillis();
    System.out.println("The seed used in the secondary index integration test is " + seed);
  }

  protected void initialize(String testName) throws Exception {}

  protected abstract Properties getProperties(String testName);

  protected String getNamespace() {
    return NAMESPACE;
  }

  protected Set getSecondaryIndexTypes() {
    return new HashSet<>(Arrays.asList(DataType.values()));
  }

  private void createTables() throws ExecutionException {
    Map options = getCreationOptions();
    admin.createNamespace(namespace, true, options);
    for (DataType secondaryIndexType : secondaryIndexTypes) {
      createTable(secondaryIndexType, options);
    }
  }

  protected Map getCreationOptions() {
    return Collections.emptyMap();
  }

  private void createTable(DataType secondaryIndexType, Map options)
      throws ExecutionException {
    admin.createTable(
        namespace,
        getTableName(secondaryIndexType),
        TableMetadata.newBuilder()
            .addColumn(PARTITION_KEY, DataType.INT)
            .addColumn(INDEX_COL_NAME, secondaryIndexType)
            .addColumn(COL_NAME, DataType.INT)
            .addPartitionKey(PARTITION_KEY)
            .addSecondaryIndex(INDEX_COL_NAME)
            .build(),
        true,
        options);
  }

  @AfterAll
  public void afterAll() throws ExecutionException {
    dropTables();
    admin.close();
    storage.close();
  }

  private void dropTables() throws ExecutionException {
    for (DataType secondaryIndexType : secondaryIndexTypes) {
      admin.dropTable(namespace, getTableName(secondaryIndexType));
    }
    admin.dropNamespace(namespace);
  }

  private void truncateTable(DataType secondaryIndexType) throws ExecutionException {
    admin.truncateTable(namespace, getTableName(secondaryIndexType));
  }

  private String getTableName(DataType secondaryIndexType) {
    return secondaryIndexType.toString();
  }

  @Test
  public void scan_WithRandomSecondaryIndexValue_ShouldReturnProperResult()
      throws ExecutionException, IOException {
    for (DataType secondaryIndexType : secondaryIndexTypes) {
      random.setSeed(seed);

      truncateTable(secondaryIndexType);

      for (int i = 0; i < ATTEMPT_COUNT; i++) {
        // Arrange
        Value secondaryIndexValue = getRandomValue(random, INDEX_COL_NAME, secondaryIndexType);
        prepareRecords(secondaryIndexType, secondaryIndexValue);
        Scan scan =
            new Scan(new Key(secondaryIndexValue))
                .forNamespace(namespace)
                .forTable(getTableName(secondaryIndexType));

        // Act
        List results = scanAll(scan);

        // Assert
        assertResults(results, secondaryIndexValue);
      }
    }
  }

  @Test
  public void scan_WithMaxSecondaryIndexValue_ShouldReturnProperResult()
      throws ExecutionException, IOException {

    for (DataType secondaryIndexType : secondaryIndexTypes) {
      truncateTable(secondaryIndexType);

      // Arrange
      Value secondaryIndexValue = getMaxValue(INDEX_COL_NAME, secondaryIndexType);
      prepareRecords(secondaryIndexType, secondaryIndexValue);
      Scan scan =
          new Scan(new Key(secondaryIndexValue))
              .forNamespace(namespace)
              .forTable(getTableName(secondaryIndexType));

      // Act
      List results = scanAll(scan);

      // Assert
      assertResults(results, secondaryIndexValue);
    }
  }

  @Test
  public void scan_WithMinSecondaryIndexValue_ShouldReturnProperResult()
      throws ExecutionException, IOException {

    for (DataType secondaryIndexType : secondaryIndexTypes) {
      truncateTable(secondaryIndexType);

      // Arrange
      Value secondaryIndexValue = getMinValue(INDEX_COL_NAME, secondaryIndexType);
      prepareRecords(secondaryIndexType, secondaryIndexValue);
      Scan scan =
          new Scan(new Key(secondaryIndexValue))
              .forNamespace(namespace)
              .forTable(getTableName(secondaryIndexType));

      // Act
      List results = scanAll(scan);

      // Assert
      assertResults(results, secondaryIndexValue);
    }
  }

  private void prepareRecords(DataType secondaryIndexType, Value secondaryIndexValue)
      throws ExecutionException {
    for (int i = 0; i < DATA_NUM; i++) {
      Put put =
          new Put(new Key(PARTITION_KEY, i))
              .withValue(secondaryIndexValue)
              .withValue(COL_NAME, 1)
              .forNamespace(namespace)
              .forTable(getTableName(secondaryIndexType));
      storage.put(put);
    }
  }

  private void assertResults(List results, Value secondaryIndexValue) {
    assertThat(results.size()).isEqualTo(DATA_NUM);

    Set partitionKeySet = new HashSet<>();
    for (int i = 0; i < DATA_NUM; i++) {
      partitionKeySet.add(i);
    }

    for (Result result : results) {
      assertThat(result.getValue(PARTITION_KEY).isPresent()).isTrue();
      partitionKeySet.remove(result.getValue(PARTITION_KEY).get().getAsInt());
      assertThat(result.getValue(INDEX_COL_NAME).isPresent()).isTrue();
      assertThat(result.getValue(INDEX_COL_NAME).get()).isEqualTo(secondaryIndexValue);
      assertThat(result.getValue(COL_NAME).isPresent()).isTrue();
      assertThat(result.getValue(COL_NAME).get().getAsInt()).isEqualTo(1);
    }

    assertThat(partitionKeySet).isEmpty();
  }

  private List scanAll(Scan scan) throws ExecutionException, IOException {
    try (Scanner scanner = storage.scan(scan)) {
      return scanner.all();
    }
  }

  protected Value getRandomValue(Random random, String columnName, DataType dataType) {
    return TestUtils.getRandomValue(random, columnName, dataType, true);
  }

  protected Value getMinValue(String columnName, DataType dataType) {
    return TestUtils.getMinValue(columnName, dataType, true);
  }

  protected Value getMaxValue(String columnName, DataType dataType) {
    return TestUtils.getMaxValue(columnName, dataType);
  }
}