com.scalar.db.storage.cosmos.BatchHandler 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.cosmos;

import com.azure.cosmos.CosmosClient;
import com.azure.cosmos.CosmosException;
import com.scalar.db.api.Mutation;
import com.scalar.db.api.TableMetadata;
import com.scalar.db.common.TableMetadataManager;
import com.scalar.db.exception.storage.ExecutionException;
import com.scalar.db.exception.storage.NoMutationException;
import com.scalar.db.exception.storage.RetriableExecutionException;
import java.util.ArrayList;
import java.util.List;
import javax.annotation.concurrent.ThreadSafe;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * A handler for a batch
 *
 * @author Yuji Ito
 */
@ThreadSafe
public class BatchHandler {
  private static final Logger logger = LoggerFactory.getLogger(BatchHandler.class);
  private static final String MUTATION_STORED_PROCEDURE = "mutate.js";
  private final CosmosClient client;
  private final TableMetadataManager metadataManager;

  /**
   * Constructs a {@code BatchHandler} with the specified {@link CosmosClient} and {@link
   * TableMetadataManager}
   *
   * @param client {@code CosmosClient} to create a statement with
   * @param metadataManager {@code TableMetadataManager}
   */
  public BatchHandler(CosmosClient client, TableMetadataManager metadataManager) {
    this.client = client;
    this.metadataManager = metadataManager;
  }

  /**
   * Execute the specified list of {@link Mutation}s in batch. All the {@link Mutation}s in the list
   * must be for the same partition.
   *
   * @param mutations a list of {@code Mutation}s to execute
   * @throws NoMutationException if at least one of conditional {@code Mutation}s fails because it
   *     didn't meet the condition
   */
  public void handle(List mutations) throws ExecutionException {
    TableMetadata tableMetadata = metadataManager.getTableMetadata(mutations.get(0));
    try {
      executeStoredProcedure(mutations, tableMetadata);
    } catch (CosmosException e) {
      throwException(e);
    }
  }

  private void executeStoredProcedure(
      List mutations, TableMetadata tableMetadata) throws CosmosException {
    List types = new ArrayList<>();
    List records = new ArrayList<>();
    List queries = new ArrayList<>();

    mutations.forEach(
        mutation -> {
          CosmosMutation cosmosMutation = new CosmosMutation(mutation, tableMetadata);
          types.add(cosmosMutation.getMutationType().ordinal());
          records.add(cosmosMutation.makeRecord());
          queries.add(cosmosMutation.makeConditionalQuery());
        });
    List