All Downloads are FREE. Search and download functionalities are using the official Maven repository.

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

Go to download

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.checkNotNull;

import com.google.common.primitives.UnsignedBytes;
import com.scalar.db.api.Consistency;
import com.scalar.db.api.Get;
import com.scalar.db.api.Scan;
import com.scalar.db.api.Scan.Ordering;
import com.scalar.db.api.Scan.Ordering.Order;
import com.scalar.db.api.ScanAll;
import com.scalar.db.api.Scanner;
import com.scalar.db.api.Selection;
import com.scalar.db.api.TableMetadata;
import com.scalar.db.common.EmptyScanner;
import com.scalar.db.common.TableMetadataManager;
import com.scalar.db.exception.storage.ExecutionException;
import com.scalar.db.io.Column;
import com.scalar.db.io.Key;
import com.scalar.db.storage.dynamo.bytes.BytesUtils;
import com.scalar.db.storage.dynamo.bytes.KeyBytesEncoder;
import com.scalar.db.util.ScalarDbUtils;
import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import javax.annotation.Nonnull;
import javax.annotation.concurrent.ThreadSafe;
import software.amazon.awssdk.core.SdkBytes;
import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
import software.amazon.awssdk.services.dynamodb.model.DynamoDbException;
import software.amazon.awssdk.services.dynamodb.model.DynamoDbRequest;
import software.amazon.awssdk.services.dynamodb.model.GetItemRequest;
import software.amazon.awssdk.services.dynamodb.model.QueryRequest;
import software.amazon.awssdk.services.dynamodb.model.ScanRequest;

/**
 * A handler class for select statement
 *
 * @author Yuji Ito, Toshihiro Suzuki
 */
@ThreadSafe
public class SelectStatementHandler {
  private final DynamoDbClient client;
  private final TableMetadataManager metadataManager;
  private final String namespacePrefix;

  /**
   * Constructs a {@code SelectStatementHandler} with the specified {@link DynamoDbClient} and a new
   * {@link TableMetadataManager}
   *
   * @param client {@code DynamoDbClient}
   * @param metadataManager {@code TableMetadataManager}
   * @param namespacePrefix a namespace prefix
   */
  @SuppressFBWarnings("EI_EXPOSE_REP2")
  public SelectStatementHandler(
      DynamoDbClient client,
      TableMetadataManager metadataManager,
      Optional namespacePrefix) {
    this.client = checkNotNull(client);
    this.metadataManager = checkNotNull(metadataManager);
    this.namespacePrefix = namespacePrefix.orElse("");
  }

  @Nonnull
  public Scanner handle(Selection selection) throws ExecutionException {
    TableMetadata tableMetadata = metadataManager.getTableMetadata(selection);
    if (selection instanceof Get) {
      selection = copyAndAppendNamespacePrefix((Get) selection);
    } else {
      selection = copyAndAppendNamespacePrefix((Scan) selection);
    }

    try {
      if (!(selection instanceof ScanAll)
          && ScalarDbUtils.isSecondaryIndexSpecified(selection, tableMetadata)) {
        return executeScanWithIndex(selection, tableMetadata);
      }

      if (selection instanceof Get) {
        return executeGet((Get) selection, tableMetadata);
      } else if (selection instanceof ScanAll) {
        return executeFullScan((ScanAll) selection, tableMetadata);
      } else {
        return executeScan((Scan) selection, tableMetadata);
      }

    } catch (DynamoDbException e) {
      throw new ExecutionException(e.getMessage(), e);
    }
  }

  private Scanner executeGet(Get get, TableMetadata tableMetadata) {
    DynamoOperation dynamoOperation = new DynamoOperation(get, tableMetadata);

    GetItemRequest.Builder builder =
        GetItemRequest.builder()
            .tableName(dynamoOperation.getTableName())
            .key(dynamoOperation.getKeyMap());

    if (!get.getProjections().isEmpty()) {
      Map expressionAttributeNames = new HashMap<>();
      projectionExpression(builder, get, expressionAttributeNames);
      builder.expressionAttributeNames(expressionAttributeNames);
    }

    if (get.getConsistency() != Consistency.EVENTUAL) {
      builder.consistentRead(true);
    }

    return new GetItemScanner(
        client, builder.build(), new ResultInterpreter(get.getProjections(), tableMetadata));
  }

  private Scanner executeScanWithIndex(Selection selection, TableMetadata tableMetadata) {
    DynamoOperation dynamoOperation = new DynamoOperation(selection, tableMetadata);
    Column keyColumn = selection.getPartitionKey().getColumns().get(0);
    String column = keyColumn.getName();
    String indexTable = dynamoOperation.getGlobalIndexName(column);
    QueryRequest.Builder builder =
        QueryRequest.builder().tableName(dynamoOperation.getTableName()).indexName(indexTable);

    String expressionColumnName = DynamoOperation.COLUMN_NAME_ALIAS + "0";
    String condition = expressionColumnName + " = " + DynamoOperation.VALUE_ALIAS + "0";
    ValueBinder binder = new ValueBinder(DynamoOperation.VALUE_ALIAS);
    keyColumn.accept(binder);
    Map bindMap = binder.build();
    builder.keyConditionExpression(condition).expressionAttributeValues(bindMap);

    Map expressionAttributeNames = new HashMap<>();
    expressionAttributeNames.put(expressionColumnName, column);

    if (!selection.getProjections().isEmpty()) {
      projectionExpression(builder, selection, expressionAttributeNames);
    }

    builder.expressionAttributeNames(expressionAttributeNames);

    if (selection instanceof Scan) {
      Scan scan = (Scan) selection;
      if (scan.getLimit() > 0) {
        builder.limit(scan.getLimit());
      }
    }
    com.scalar.db.storage.dynamo.request.QueryRequest request =
        new com.scalar.db.storage.dynamo.request.QueryRequest(client, builder.build());
    return new QueryScanner(
        request, new ResultInterpreter(selection.getProjections(), tableMetadata));
  }

  private Scanner executeScan(Scan scan, TableMetadata tableMetadata) {
    DynamoOperation dynamoOperation = new DynamoOperation(scan, tableMetadata);
    QueryRequest.Builder builder = QueryRequest.builder().tableName(dynamoOperation.getTableName());

    if (!setConditions(builder, scan, tableMetadata)) {
      // if setConditions() fails, return an empty scanner
      return new EmptyScanner();
    }

    if (!scan.getOrderings().isEmpty()) {
      Ordering ordering = scan.getOrderings().get(0);
      if (ordering.getOrder() != tableMetadata.getClusteringOrder(ordering.getColumnName())) {
        // reverse scan
        builder.scanIndexForward(false);
      }
    }

    if (scan.getLimit() > 0) {
      builder.limit(scan.getLimit());
    }

    if (!scan.getProjections().isEmpty()) {
      Map expressionAttributeNames = new HashMap<>();
      projectionExpression(builder, scan, expressionAttributeNames);
      builder.expressionAttributeNames(expressionAttributeNames);
    }

    if (scan.getConsistency() != Consistency.EVENTUAL) {
      builder.consistentRead(true);
    }
    com.scalar.db.storage.dynamo.request.QueryRequest queryRequest =
        new com.scalar.db.storage.dynamo.request.QueryRequest(client, builder.build());
    return new QueryScanner(
        queryRequest, new ResultInterpreter(scan.getProjections(), tableMetadata));
  }

  private Scanner executeFullScan(ScanAll scan, TableMetadata tableMetadata) {
    DynamoOperation dynamoOperation = new DynamoOperation(scan, tableMetadata);
    ScanRequest.Builder builder = ScanRequest.builder().tableName(dynamoOperation.getTableName());

    if (scan.getLimit() > 0) {
      builder.limit(scan.getLimit());
    }

    if (!scan.getProjections().isEmpty()) {
      Map expressionAttributeNames = new HashMap<>();
      projectionExpression(builder, scan, expressionAttributeNames);
      builder.expressionAttributeNames(expressionAttributeNames);
    }

    if (scan.getConsistency() != Consistency.EVENTUAL) {
      builder.consistentRead(true);
    }
    com.scalar.db.storage.dynamo.request.ScanRequest requestWrapper =
        new com.scalar.db.storage.dynamo.request.ScanRequest(client, builder.build());
    return new QueryScanner(
        requestWrapper, new ResultInterpreter(scan.getProjections(), tableMetadata));
  }

  private void projectionExpression(
      DynamoDbRequest.Builder builder,
      Selection selection,
      Map expressionAttributeNames) {
    assert builder instanceof GetItemRequest.Builder
        || builder instanceof QueryRequest.Builder
        || builder instanceof ScanRequest.Builder;

    List projections = new ArrayList<>(selection.getProjections().size());
    for (String projection : selection.getProjections()) {
      String alias = DynamoOperation.COLUMN_NAME_ALIAS + expressionAttributeNames.size();
      projections.add(alias);
      expressionAttributeNames.put(alias, projection);
    }
    String projectionExpression = String.join(",", projections);

    if (builder instanceof GetItemRequest.Builder) {
      ((GetItemRequest.Builder) builder).projectionExpression(projectionExpression);
    } else if (builder instanceof QueryRequest.Builder) {
      ((QueryRequest.Builder) builder).projectionExpression(projectionExpression);
    } else {
      ((ScanRequest.Builder) builder).projectionExpression(projectionExpression);
    }
  }

  private boolean setConditions(
      QueryRequest.Builder builder, Scan scan, TableMetadata tableMetadata) {
    List conditions = new ArrayList<>();
    Map bindMap = new HashMap<>();

    setConditionForPartitionKey(scan, tableMetadata, conditions, bindMap);

    // If the scan is for DESC clustering order, use the end clustering key as a start key and the
    // start clustering key as an end key
    boolean scanForDescClusteringOrder = isScanForDescClusteringOrder(scan, tableMetadata);
    Optional startKey =
        scanForDescClusteringOrder ? scan.getEndClusteringKey() : scan.getStartClusteringKey();
    boolean startInclusive =
        scanForDescClusteringOrder ? scan.getEndInclusive() : scan.getStartInclusive();
    Optional endKey =
        scanForDescClusteringOrder ? scan.getStartClusteringKey() : scan.getEndClusteringKey();
    boolean endInclusive =
        scanForDescClusteringOrder ? scan.getStartInclusive() : scan.getEndInclusive();

    if (startKey.isPresent() && endKey.isPresent()) {
      if (!setBetweenCondition(
          startKey.get(),
          startInclusive,
          endKey.get(),
          endInclusive,
          tableMetadata,
          conditions,
          bindMap)) {
        return false;
      }
    } else {
      if (startKey.isPresent()) {
        if (startKey.get().size() == 1) {
          if (!setStartCondition(
              startKey.get(), startInclusive, tableMetadata, conditions, bindMap)) {
            return false;
          }
        } else {
          // if a start key with multiple values specified and no end key specified, use between
          // condition and use a key based on the start key without the last value as an end key
          if (!setBetweenCondition(
              startKey.get(),
              startInclusive,
              getKeyWithoutLastValue(startKey.get()),
              true,
              tableMetadata,
              conditions,
              bindMap)) {
            return false;
          }
        }
      }

      if (endKey.isPresent()) {
        if (endKey.get().size() == 1) {
          setEndCondition(endKey.get(), endInclusive, tableMetadata, conditions, bindMap);
        } else {
          // if an end key with multiple values specified and no start key specified, use between
          // condition and use a key based on the end key without the last value as a start key
          if (!setBetweenCondition(
              getKeyWithoutLastValue(endKey.get()),
              true,
              endKey.get(),
              endInclusive,
              tableMetadata,
              conditions,
              bindMap)) {
            return false;
          }
        }
      }
    }

    builder
        .keyConditionExpression(String.join(" AND ", conditions))
        .expressionAttributeValues(bindMap);
    return true;
  }

  private Key getKeyWithoutLastValue(Key originalKey) {
    Key.Builder keyBuilder = Key.newBuilder();
    for (int i = 0; i < originalKey.get().size() - 1; i++) {
      keyBuilder.add(originalKey.get().get(i));
    }
    return keyBuilder.build();
  }

  private void setConditionForPartitionKey(
      Scan scan,
      TableMetadata tableMetadata,
      List conditions,
      Map bindMap) {
    conditions.add(DynamoOperation.PARTITION_KEY + " = " + DynamoOperation.PARTITION_KEY_ALIAS);

    DynamoOperation dynamoOperation = new DynamoOperation(scan, tableMetadata);
    ByteBuffer concatenatedPartitionKey = dynamoOperation.getConcatenatedPartitionKey();
    bindMap.put(
        DynamoOperation.PARTITION_KEY_ALIAS,
        AttributeValue.builder().b(SdkBytes.fromByteBuffer(concatenatedPartitionKey)).build());
  }

  private boolean setBetweenCondition(
      Key startKey,
      boolean startInclusive,
      Key endKey,
      boolean endInclusive,
      TableMetadata tableMetadata,
      List conditions,
      Map bindMap) {

    ByteBuffer startKeyBytes = getKeyBytes(startKey, tableMetadata);
    if (!startInclusive) {
      // if exclusive scan for the start key, we use the closest next bytes of the start key bytes
      Optional closestNextBytes = BytesUtils.getClosestNextBytes(startKeyBytes);
      if (!closestNextBytes.isPresent()) {
        // if we can't find the closest next bytes of the start key bytes, return false. That means
        // we should return empty results in this case
        return false;
      }
      startKeyBytes = closestNextBytes.get();
    }

    ByteBuffer endKeyBytes = getKeyBytes(endKey, tableMetadata);
    boolean fullClusteringKeySpecified =
        endKey.size() == tableMetadata.getClusteringKeyNames().size();
    if (fullClusteringKeySpecified) {
      if (!endInclusive) {
        // if full end key specified, and it's an exclusive scan for the end key, we use the closest
        // previous bytes of the end key bytes for the between condition
        Optional closestPreviousBytes = BytesUtils.getClosestPreviousBytes(endKeyBytes);
        if (!closestPreviousBytes.isPresent()) {
          // if we can't find the closest previous bytes of the end key bytes, return false. That
          // means we should return empty results in this case
          return false;
        }
        endKeyBytes = closestPreviousBytes.get();
      }
    } else {
      if (endInclusive) {
        // if partial end key specified, and it's an inclusive scan for the end key, we use the
        // closest next bytes of the end key bytes for the between condition
        Optional closestNextBytes = BytesUtils.getClosestNextBytes(endKeyBytes);
        if (!closestNextBytes.isPresent()) {
          // if we can't find the closest next bytes of the end key bytes, set start condition with
          // the start key
          return setStartCondition(startKey, startInclusive, tableMetadata, conditions, bindMap);
        }
        endKeyBytes = closestNextBytes.get();
      }
    }

    byte[] start = BytesUtils.toBytes(startKeyBytes);
    byte[] end = BytesUtils.toBytes(endKeyBytes);
    if (UnsignedBytes.lexicographicalComparator().compare(start, end) > 0) {
      // if the start key bytes are greater than the end key bytes, return false. That means we
      // should return empty results in this case. This situation could happen when full clustering
      // keys specified and scanning exclusively
      return false;
    }

    conditions.add(
        DynamoOperation.CLUSTERING_KEY
            + " BETWEEN "
            + DynamoOperation.START_CLUSTERING_KEY_ALIAS
            + " AND "
            + DynamoOperation.END_CLUSTERING_KEY_ALIAS);

    bindMap.put(
        DynamoOperation.START_CLUSTERING_KEY_ALIAS,
        AttributeValue.builder().b(SdkBytes.fromByteArray(start)).build());
    bindMap.put(
        DynamoOperation.END_CLUSTERING_KEY_ALIAS,
        AttributeValue.builder().b(SdkBytes.fromByteArray(end)).build());

    return true;
  }

  private boolean setStartCondition(
      Key startKey,
      boolean startInclusive,
      TableMetadata tableMetadata,
      List conditions,
      Map bindMap) {
    ByteBuffer startKeyBytes = getKeyBytes(startKey, tableMetadata);

    boolean fullClusteringKeySpecified =
        startKey.size() == tableMetadata.getClusteringKeyNames().size();
    if (fullClusteringKeySpecified) {
      conditions.add(
          DynamoOperation.CLUSTERING_KEY
              + (startInclusive ? " >= " : " > ")
              + DynamoOperation.START_CLUSTERING_KEY_ALIAS);
      bindMap.put(
          DynamoOperation.START_CLUSTERING_KEY_ALIAS,
          AttributeValue.builder().b(SdkBytes.fromByteBuffer(startKeyBytes)).build());
    } else {
      if (startInclusive) {
        conditions.add(
            DynamoOperation.CLUSTERING_KEY + " >= " + DynamoOperation.START_CLUSTERING_KEY_ALIAS);
        bindMap.put(
            DynamoOperation.START_CLUSTERING_KEY_ALIAS,
            AttributeValue.builder().b(SdkBytes.fromByteBuffer(startKeyBytes)).build());
      } else {
        // if partial start key specified, and it's an exclusive scan for the start key, we use
        // the closest next bytes of the start key bytes for the grater than or equal condition
        Optional closestNextBytes = BytesUtils.getClosestNextBytes(startKeyBytes);
        if (closestNextBytes.isPresent()) {
          conditions.add(
              DynamoOperation.CLUSTERING_KEY + " >= " + DynamoOperation.START_CLUSTERING_KEY_ALIAS);
          bindMap.put(
              DynamoOperation.START_CLUSTERING_KEY_ALIAS,
              AttributeValue.builder().b(SdkBytes.fromByteBuffer(closestNextBytes.get())).build());
        } else {
          // if we can't find the closest next bytes of the start key bytes, return false. That
          // means we should return empty results in this case
          return false;
        }
      }
    }
    return true;
  }

  private void setEndCondition(
      Key endKey,
      boolean endInclusive,
      TableMetadata tableMetadata,
      List conditions,
      Map bindMap) {
    ByteBuffer endKeyBytes = getKeyBytes(endKey, tableMetadata);

    boolean fullClusteringKeySpecified =
        endKey.size() == tableMetadata.getClusteringKeyNames().size();
    if (fullClusteringKeySpecified) {
      conditions.add(
          DynamoOperation.CLUSTERING_KEY
              + (endInclusive ? " <= " : " < ")
              + DynamoOperation.END_CLUSTERING_KEY_ALIAS);
      bindMap.put(
          DynamoOperation.END_CLUSTERING_KEY_ALIAS,
          AttributeValue.builder().b(SdkBytes.fromByteBuffer(endKeyBytes)).build());
    } else {
      if (endInclusive) {
        // if partial end key specified, and it's an inclusive scan for the end key, we use the
        // closest next bytes of the end key bytes for the less than condition
        BytesUtils.getClosestNextBytes(endKeyBytes)
            .ifPresent(
                k -> {
                  conditions.add(
                      DynamoOperation.CLUSTERING_KEY
                          + " < "
                          + DynamoOperation.END_CLUSTERING_KEY_ALIAS);
                  bindMap.put(
                      DynamoOperation.END_CLUSTERING_KEY_ALIAS,
                      AttributeValue.builder().b(SdkBytes.fromByteBuffer(k)).build());
                });
      } else {
        conditions.add(
            DynamoOperation.CLUSTERING_KEY + " < " + DynamoOperation.END_CLUSTERING_KEY_ALIAS);
        bindMap.put(
            DynamoOperation.END_CLUSTERING_KEY_ALIAS,
            AttributeValue.builder().b(SdkBytes.fromByteBuffer(endKeyBytes)).build());
      }
    }
  }

  private ByteBuffer getKeyBytes(Key key, TableMetadata tableMetadata) {
    return new KeyBytesEncoder().encode(key, tableMetadata.getClusteringOrders());
  }

  private boolean isScanForDescClusteringOrder(Scan scan, TableMetadata tableMetadata) {
    if (scan.getStartClusteringKey().isPresent()) {
      Key startClusteringKey = scan.getStartClusteringKey().get();
      String lastValueName = startClusteringKey.get().get(startClusteringKey.size() - 1).getName();
      return tableMetadata.getClusteringOrder(lastValueName) == Order.DESC;
    }
    if (scan.getEndClusteringKey().isPresent()) {
      Key endClusteringKey = scan.getEndClusteringKey().get();
      String lastValueName = endClusteringKey.get().get(endClusteringKey.size() - 1).getName();
      return tableMetadata.getClusteringOrder(lastValueName) == Order.DESC;
    }
    return false;
  }

  private Get copyAndAppendNamespacePrefix(Get get) {
    assert get.forNamespace().isPresent();
    return Get.newBuilder(get).namespace(namespacePrefix + get.forNamespace().get()).build();
  }

  private Scan copyAndAppendNamespacePrefix(Scan scan) {
    assert scan.forNamespace().isPresent();
    return Scan.newBuilder(scan).namespace(namespacePrefix + scan.forNamespace().get()).build();
  }
}




© 2015 - 2024 Weber Informatics LLC | Privacy Policy