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

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

import org.postgresql.adba.PgSubmission;
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 GroupSubmission implements PgSubmission {
  private final Supplier cancel;
  private CompletableFuture publicStage;
  private CompletionStage membersTail;
  private Consumer errorHandler;
  private final AtomicBoolean sendConsumed = new AtomicBoolean(false);
  private Collector collector;
  private Object collectorHolder;

  public GroupSubmission(Supplier cancel, Consumer errorHandler) {
    this.cancel = cancel;
    this.errorHandler = errorHandler;
  }

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

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

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

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

  @Override
  public void setCollector(Collector collector) {
    this.collector = collector;

    collectorHolder = collector.supplier().get();
  }

  @Override
  public Object finish(Object finishObject) {
    try {
      Object o = null;
      if (collector != null) {
        o = collector.finisher().apply(collectorHolder);
      }
      ((CompletableFuture)getCompletionStage()).complete((T)o);
    } catch (Throwable t) {
      ((CompletableFuture)getCompletionStage()).completeExceptionally(t);
    }
    return null;
  }

  @Override
  public void addRow(DataRow row) {

  }

  /**
   * A group submission accumulates the results of it's members.
   * @param result result from a member operation
   */
  public void addGroupResult(Object result) {
    try {
      collector.accumulator().accept(collectorHolder, result);
    } catch (Throwable e) {
      publicStage.completeExceptionally(e);
    }
  }

  @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 new CompletableFuture().completeAsync(cancel);
  }

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

    return publicStage;
  }

  /**
   * the group submission have a monad of operations, and this adds another on top of it.
   * @param completionStage operation to add
   */
  public void stackFuture(CompletableFuture completionStage) {
    if (membersTail == null) {
      membersTail = getCompletionStage();
    }

    membersTail.exceptionally(e -> {
      completionStage.completeExceptionally(e);
      return null;
    });
    membersTail = completionStage;
  }
}