org.postgresql.adba.submissions.TransactionSubmission Maven / Gradle / Ivy

Go to download
package org.postgresql.adba.submissions;

import jdk.incubator.sql2.TransactionOutcome;
import org.postgresql.adba.PgSubmission;
import org.postgresql.adba.communication.packets.CommandComplete;
import org.postgresql.adba.communication.packets.DataRow;
import org.postgresql.adba.operations.helpers.ParameterHolder;

import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionStage;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.Consumer;
import java.util.function.Supplier;
import java.util.stream.Collector;

public class TransactionSubmission implements PgSubmission {
  private final Supplier cancel;
  private CompletableFuture publicStage;
  private String sql;
  private final AtomicBoolean sendConsumed = new AtomicBoolean(false);
  private Consumer errorHandler;
  private ParameterHolder holder = new ParameterHolder();

  /**
   * Creates the submission object of the transation operations.
   *
   * @param cancel cancel method
   * @param errorHandler error handler method
   * @param sql the query
   */
  public TransactionSubmission(Supplier cancel, Consumer errorHandler, String sql) {
    this.cancel = cancel;
    this.errorHandler = errorHandler;
    this.sql = sql;
  }

  @Override
  public String getSql() {
    return sql;
  }

  @Override
  public AtomicBoolean getSendConsumed() {
    return sendConsumed;
  }

  @Override
  public ParameterHolder getHolder() {
    return holder;
  }

  @Override
  public Types getCompletionType() {
    return Types.TRANSACTION;
  }

  @Override
  public void setCollector(Collector collector) {

  }

  @Override
  public Object finish(Object finishObject) {
    CommandComplete.Types type = (CommandComplete.Types)finishObject;
    if (type == CommandComplete.Types.ROLLBACK) {
      ((CompletableFuture) getCompletionStage())
          .complete(TransactionOutcome.ROLLBACK);
    } else if (type == CommandComplete.Types.COMMIT) {
      ((CompletableFuture) getCompletionStage())
          .complete(TransactionOutcome.COMMIT);
    } else {
      ((CompletableFuture) getCompletionStage())
          .complete(TransactionOutcome.UNKNOWN);
    }
    return null;
  }

  @Override
  public void addRow(DataRow row) {

  }

  @Override
  public List getParamTypes() throws ExecutionException, InterruptedException {
    return null;
  }

  @Override
  public int numberOfQueryRepetitions() throws ExecutionException, InterruptedException {
    return 1;
  }

  @Override
  public Consumer getErrorHandler() {
    return errorHandler;
  }

  @Override
  public CompletionStage cancel() {
    return null;
  }

  @Override
  public CompletionStage getCompletionStage() {
    if (publicStage == null) {
      publicStage = new CompletableFuture<>();
    }

    return publicStage;
  }
}