com.scalar.db.storage.dynamo.Dynamo Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of scalardb Show documentation
A universal transaction manager that achieves database-agnostic transactions and distributed transactions that span multiple databases
There is a newer version: 3.14.0
Show newest version
package com.scalar.db.storage.dynamo;

import static com.google.common.base.Preconditions.checkArgument;

import com.google.common.annotations.VisibleForTesting;
import com.google.inject.Inject;
import com.scalar.db.api.Delete;
import com.scalar.db.api.DistributedStorage;
import com.scalar.db.api.Get;
import com.scalar.db.api.Mutation;
import com.scalar.db.api.Put;
import com.scalar.db.api.Result;
import com.scalar.db.api.Scan;
import com.scalar.db.api.Scanner;
import com.scalar.db.common.AbstractDistributedStorage;
import com.scalar.db.common.TableMetadataManager;
import com.scalar.db.common.checker.OperationChecker;
import com.scalar.db.config.DatabaseConfig;
import com.scalar.db.exception.storage.ExecutionException;
import com.scalar.db.util.ScalarDbUtils;
import java.io.IOException;
import java.net.URI;
import java.util.List;
import java.util.Optional;
import javax.annotation.Nonnull;
import javax.annotation.concurrent.ThreadSafe;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
import software.amazon.awssdk.regions.Region;
import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
import software.amazon.awssdk.services.dynamodb.DynamoDbClientBuilder;

/**
 * A storage implementation with DynamoDB for {@link DistributedStorage}
 *
 * @author Yuji Ito
 */
@ThreadSafe
public class Dynamo extends AbstractDistributedStorage {
  private static final Logger logger = LoggerFactory.getLogger(Dynamo.class);

  private final DynamoDbClient client;
  private final SelectStatementHandler selectStatementHandler;
  private final PutStatementHandler putStatementHandler;
  private final DeleteStatementHandler deleteStatementHandler;
  private final BatchHandler batchHandler;
  private final OperationChecker operationChecker;

  @Inject
  public Dynamo(DatabaseConfig databaseConfig) {
    super(databaseConfig);
    DynamoConfig config = new DynamoConfig(databaseConfig);

    DynamoDbClientBuilder builder = DynamoDbClient.builder();
    config.getEndpointOverride().ifPresent(e -> builder.endpointOverride(URI.create(e)));
    client =
        builder
            .credentialsProvider(
                StaticCredentialsProvider.create(
                    AwsBasicCredentials.create(
                        config.getAccessKeyId(), config.getSecretAccessKey())))
            .region(Region.of(config.getRegion()))
            .build();

    TableMetadataManager metadataManager =
        new TableMetadataManager(
            new DynamoAdmin(client, config), databaseConfig.getMetadataCacheExpirationTimeSecs());
    operationChecker = new DynamoOperationChecker(metadataManager);

    selectStatementHandler =
        new SelectStatementHandler(client, metadataManager, config.getNamespacePrefix());
    putStatementHandler =
        new PutStatementHandler(client, metadataManager, config.getNamespacePrefix());
    deleteStatementHandler =
        new DeleteStatementHandler(client, metadataManager, config.getNamespacePrefix());
    batchHandler = new BatchHandler(client, metadataManager, config.getNamespacePrefix());

    logger.info("DynamoDB object is created properly.");
  }

  @VisibleForTesting
  Dynamo(
      DatabaseConfig databaseConfig,
      DynamoDbClient client,
      SelectStatementHandler select,
      PutStatementHandler put,
      DeleteStatementHandler delete,
      BatchHandler batch,
      OperationChecker operationChecker) {
    super(databaseConfig);
    this.client = client;
    this.selectStatementHandler = select;
    this.putStatementHandler = put;
    this.deleteStatementHandler = delete;
    this.batchHandler = batch;
    this.operationChecker = operationChecker;
  }

  @Override
  @Nonnull
  public Optional get(Get get) throws ExecutionException {
    get = copyAndSetTargetToIfNot(get);
    operationChecker.check(get);

    Scanner scanner = null;
    try {
      scanner = selectStatementHandler.handle(get);
      Optional ret = scanner.one();
      if (scanner.one().isPresent()) {
        throw new IllegalArgumentException("please use scan() for non-exact match selection");
      }
      return ret;
    } finally {
      if (scanner != null) {
        try {
          scanner.close();
        } catch (IOException e) {
          logger.warn("failed to close the scanner", e);
        }
      }
    }
  }

  @Override
  public Scanner scan(Scan scan) throws ExecutionException {
    scan = copyAndSetTargetToIfNot(scan);
    operationChecker.check(scan);

    if (ScalarDbUtils.isRelational(scan)) {
      throw new UnsupportedOperationException(
          "scanning all records with orderings or conditions is not supported in DynamoDB");
    }

    return selectStatementHandler.handle(scan);
  }

  @Override
  public void put(Put put) throws ExecutionException {
    put = copyAndSetTargetToIfNot(put);
    operationChecker.check(put);

    putStatementHandler.handle(put);
  }

  @Override
  public void put(List puts) throws ExecutionException {
    mutate(puts);
  }

  @Override
  public void delete(Delete delete) throws ExecutionException {
    delete = copyAndSetTargetToIfNot(delete);
    operationChecker.check(delete);

    deleteStatementHandler.handle(delete);
  }

  @Override
  public void delete(List deletes) throws ExecutionException {
    mutate(deletes);
  }

  @Override
  public void mutate(List mutations) throws ExecutionException {
    checkArgument(mutations.size() != 0);
    if (mutations.size() == 1) {
      Mutation mutation = mutations.get(0);
      if (mutation instanceof Put) {
        put((Put) mutation);
      } else if (mutation instanceof Delete) {
        delete((Delete) mutation);
      }
      return;
    }

    mutations = copyAndSetTargetToIfNot(mutations);
    operationChecker.check(mutations);
    batchHandler.handle(mutations);
  }

  @Override
  public void close() {
    client.close();
  }
}