com.clickzetta.platform.arrow.ArrowRpcResponseRetryCallback 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.arrow;

import com.clickzetta.platform.client.RpcResponseRetryCallback;
import com.clickzetta.platform.client.Session;
import com.clickzetta.platform.client.Table;
import com.clickzetta.platform.client.api.*;
import com.clickzetta.platform.client.message.RequestMessage;
import com.clickzetta.platform.client.message.ResponseMessage;
import com.google.common.collect.Lists;
import cz.proto.ingestion.Ingestion;
import cz.proto.ingestion.v2.IngestionV2;

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

public class ArrowRpcResponseRetryCallback extends RpcResponseRetryCallback<
    RequestMessage,
    ResponseMessage> {

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

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

  @Override
  public Message buildFailureMessage(Table table, RequestMessage request,
                                     ResponseMessage response) {
    return new ArrowFailureMessage(session.getSessionId(), (ArrowTable) 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 Lists.newArrayList(table.getSchemaName());
      }

      @Override
      public List getTableName() {
        return Lists.newArrayList(table.getTableName());
      }

      @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;
      }
    };
  }
}