com.clickzetta.platform.flusher.FlushTask 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.flusher;

import com.clickzetta.platform.client.RpcRequestCallback;
import com.clickzetta.platform.client.api.ClientContext;
import com.clickzetta.platform.client.Table;
import com.clickzetta.platform.client.api.Listener;
import com.clickzetta.platform.connection.ChannelManager;
import com.clickzetta.platform.operator.OperationsEncoder;
import cz.proto.ingestion.Ingestion;
import org.apache.kudu.RowOperations;
import scala.Tuple2;

import java.util.function.Supplier;

public class FlushTask extends AbstractTask {

  private long batchId;
  private ClientContext context;
  private Buffer buffer;
  private Supplier> channelDataSupplier;
  private RpcRequestCallback rpcRequestCallback;
  private Listener listener;

  public FlushTask(
      long batchId,
      ClientContext context,
      Buffer buffer,
      Supplier> channelDataSupplier,
      RpcRequestCallback rpcResponseCallback,
      Listener listener) {
    this.batchId = batchId;
    this.context = context;
    this.buffer = buffer;
    this.channelDataSupplier = channelDataSupplier;
    this.rpcRequestCallback = rpcResponseCallback;
    this.listener = listener;
  }

  @Override
  public void callInternal() throws Exception {
    if (!buffer.isEmpty()) {
      Table table = buffer.getOperations().get(0).getTable();
      Tuple2 pbData =
          OperationsEncoder.encodeOperations(buffer.getOperations());
      Ingestion.DataMutateRequest.Builder builder = Ingestion.DataMutateRequest.newBuilder()
          .setBatchId(batchId)
          .setInstanceId(context.instanceId())
          .setWorkspace(context.workspace())
          .setWriteTimestamp(System.currentTimeMillis())
          .setSchemaName(table.getSchemaName())
          .setTableName(table.getTableName())
          .setTableType(table.getIgsTableType())
          .setSchema(table.getMeta())
          .setSchemaPb(table.getSchemaPB())
          .setRowOperations(pbData._1)
          .setKeySchemaPb(table.getKeySchemaPB())
          .setKeyRowOperations(pbData._2)
          .setBucketsCount(table.getBucketsNum())
          .addAllKeyNames(table.getSchema().getKeyIndexMap().keySet())
          .addAllSortNames(table.getSchema().getSortIndexMap().keySet())
          .setBatchCount(buffer.getCurrentLines())
          .setIsDispatch(false);
      if (context.authentication()) {
        Ingestion.Account.Builder accountBuilder = Ingestion.Account.newBuilder();
        accountBuilder.setUserName(context.userName());
        if (context.userId() != null) {
          accountBuilder.setUserId(context.userId());
        }
        builder.setToken(context.token()).setAccount(accountBuilder.build());
      }
      Ingestion.DataMutateRequest request = builder.build();
      ChannelManager.ChannelData channelData = channelDataSupplier.get();
      rpcRequestCallback.setTargetHost(channelData.hostPort.toString());
      rpcRequestCallback.onSuccess(request, future);
      try {
        synchronized (channelData.streamObserver) {
          channelData.streamObserver.onNext(request);
        }
      } catch (Throwable t) {
        // roll back rpcRequestCallback onSuccess.
        rpcRequestCallback.onFailure(request, future, t);
      }
    } else {
      future.complete(true);
    }
    listener.onApplyRpcSend();
  }

  @Override
  public long getId() {
    return batchId;
  }

  public Buffer getBuffer() {
    return buffer;
  }
}