com.clickzetta.platform.client.api.multi.MultiRpcResponseRetryCallback Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of clickzetta-java Show documentation
The java SDK for clickzetta's Lakehouse
There is a newer version: 2.0.0
package com.clickzetta.platform.client.api.multi;

import com.clickzetta.platform.client.RpcResponseRetryCallback;
import com.clickzetta.platform.client.Session;
import com.clickzetta.platform.client.Table;
import com.clickzetta.platform.client.api.Message;
import com.clickzetta.platform.client.api.Options;
import com.clickzetta.platform.client.api.Row;
import com.clickzetta.platform.client.message.RequestMessage;
import com.clickzetta.platform.client.message.ResponseMessage;
import cz.proto.ingestion.Ingestion;
import cz.proto.ingestion.v2.IngestionV2;

import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.stream.Collectors;

public class MultiRpcResponseRetryCallback extends RpcResponseRetryCallback<
    RequestMessage,
    ResponseMessage> {

  public final Set EMPTY_SET = new HashSet<>(0);

  public MultiRpcResponseRetryCallback(Session session, Options options) {
    super(session, options, options.getRequestFailedRetryStatus().getV2CodeList()
        .stream().map(IngestionV2.Code::getNumber).collect(Collectors.toList()));
  }

  public final Set getRetryStatus() {
    if (retryEnable) {
      return this.retryStatus;
    } else {
      return EMPTY_SET;
    }
  }

  @Override
  public boolean retryCondition(Set retryStatus, ResponseMessage multiResponse) {
    return multiResponse.getOriginal().getMutateResponsesList().stream().allMatch(
        response -> {
          if (response.getStatus().getCode() == IngestionV2.Code.SUCCESS) {
            return true;
          } else {
            return retryStatus.contains(response.getStatus().getCode().getNumber());
          }
        });
  }

  @Override
  public Message buildSuccessMessage(RequestMessage request,
                                     ResponseMessage response) {
    return new MultiSuccessMessage(session.getSessionId(), request.getOriginal(), response.getOriginal());
  }

  @Override
  public Message buildFailureMessage(Table table, RequestMessage request,
                                     ResponseMessage response) {
    return new MultiFailureMessage(session.getSessionId(), (MultiTable) table, request.getOriginal(), response.getOriginal());
  }

  @Override
  public Message buildOnErrorMessage(Table table, RequestMessage request) {
    return new Message() {
      @Override
      public String getSessionId() {
        return session.getSessionId();
      }

      @Override
      public long getBatchId() {
        return request.getRequestId();
      }

      @Override
      public long getTimestamp() {
        return request.getTimestamp();
      }

      @Override
      public List getSchemaName() {
        return request.getOriginal().getMutateRequestsList().stream()
            .map(mutateRequest -> mutateRequest.getTableIdent().getSchemaName()).collect(Collectors.toList());
      }

      @Override
      public List getTableName() {
        return request.getOriginal().getMutateRequestsList().stream()
            .map(mutateRequest -> mutateRequest.getTableIdent().getTableName()).collect(Collectors.toList());
      }

      @Override
      public long getTotalRowsCounts() {
        return request.getBatchCount();
      }

      @Override
      public List getRowStates() {
        return null;
      }

      @Override
      public int getErrorRowsCounts() {
        return 0;
      }

      @Override
      public List getErrorRows() {
        return null;
      }
    };
  }
}