com.clickzetta.platform.client.CZSession 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;

import com.clickzetta.platform.client.api.*;
import com.clickzetta.platform.client.message.RequestMessage;
import com.clickzetta.platform.client.message.ResponseMessage;
import com.clickzetta.platform.client.message.kudu.KuduRequestMessage;
import com.clickzetta.platform.client.message.kudu.KuduResponseMessage;
import com.clickzetta.platform.client.proxy.RpcProxy;
import com.clickzetta.platform.common.CZException;
import com.clickzetta.platform.common.NotifyScheduledExecutorService;
import com.clickzetta.platform.connection.ChannelManager;
import com.clickzetta.platform.flusher.AbstractTask;
import com.clickzetta.platform.flusher.Buffer;
import com.clickzetta.platform.flusher.FlushTask;
import com.clickzetta.platform.flusher.StreamTask;
import com.clickzetta.platform.util.Util;
import cz.proto.ingestion.IGSWorkerServiceGrpc;
import cz.proto.ingestion.Ingestion;
import io.grpc.CallCredentials;
import io.grpc.CallOptions;
import io.grpc.ManagedChannel;
import io.grpc.Metadata;
import io.grpc.stub.StreamObserver;

import java.io.IOException;
import java.util.concurrent.Executor;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Supplier;

import static com.clickzetta.platform.connection.ChannelManager.RPC_CLIENT_STREAM_ID;

public final class CZSession
    extends AbstractSession {

  private CZSession(RpcProxy client, CZStream stream, Options options) {
    super(client, stream, options,
        new RpcCallbackFactory<
            RequestMessage,
            ResponseMessage>() {
          @Override
          public RpcCallback<
              RequestMessage,
              ResponseMessage>
          buildRpcResponseCallback(Session session, Options options) {
            return new CZRpcResponseRetryCallback(session, options);
          }

          @Override
          public RpcResponseHandler<
              RequestMessage,
              ResponseMessage> buildResponseHandler(Session session) {
            return new RpcResponseHandler,
                ResponseMessage>(session) {
              @Override
              public void handleCallback(
                  RpcCallback,
                      ResponseMessage> responseCallback,
                  RequestMessage req,
                  ResponseMessage resp) {
                if (resp.getStatusCode() == Ingestion.Code.SUCCESS.getNumber()) {
                  responseCallback.onSuccess(req, resp);
                } else {
                  // TODO get error row handler.
                  responseCallback.onFailure(req, resp, new CZException(resp.getOriginal().getStatus().toString()));
                }
              }
            };
          }
        });
  }

  public static CZSession build(CZClient czClient, CZStream stream, Options options) {
    return new CZSession(czClient, stream, options);
  }

  @Override
  public StreamObserver buildStreamObserver(
      ManagedChannel channel,
      CallOptions callOptions,
      StreamObserver streamObserver) {
    IGSWorkerServiceGrpc.IGSWorkerServiceStub stub = IGSWorkerServiceGrpc.newStub(channel);
    if (callOptions.getDeadline() != null) {
      stub = stub.withDeadline(callOptions.getDeadline());
    }
    if (callOptions.getCompressor() != null) {
      stub = stub.withCompression(callOptions.getCompressor());
    }
    if (callOptions.getMaxInboundMessageSize() != null) {
      stub = stub.withMaxInboundMessageSize(callOptions.getMaxInboundMessageSize());
    }
    if (callOptions.getMaxOutboundMessageSize() != null) {
      stub = stub.withMaxOutboundMessageSize(callOptions.getMaxOutboundMessageSize());
    }
    stub = stub.withCallCredentials(new CallCredentials() {
      @Override
      public void applyRequestMetadata(RequestInfo requestInfo, Executor executor, MetadataApplier metadataApplier) {
        Metadata metadata = new Metadata();
        metadata.put(Metadata.Key.of(RPC_CLIENT_STREAM_ID, Metadata.ASCII_STRING_MARSHALLER), Util.generatorStreamId(sessionId));
        metadataApplier.apply(metadata);
        LOG.info("session {} apply request metadata with streamId {}", sessionId, metadata);
      }

      @Override
      public void thisUsesUnstableApi() {
      }
    });
    return stub.mutate(streamObserver);
  }

  @Override
  public RequestMessage buildRequestMessage(Ingestion.DataMutateRequest request) {
    return new KuduRequestMessage(request);
  }

  @Override
  public ResponseMessage buildResponseMessage(Ingestion.DataMutateResponse response) {
    return new KuduResponseMessage(response);
  }

  @Override
  public Ingestion.DataMutateResponse buildRetryResponseMessage(long requestId, Ingestion.DataMutateRequest request) {
    throw new UnsupportedOperationException();
  }

  @Override
  public AbstractTask buildStreamTask(Session session, NotifyScheduledExecutorService retryThreadPool,
                                      AtomicInteger retryRequestCnt, Table table, ClientContext context,
                                      long internalMs, RetryMode retryMode,
                                      Ingestion.DataMutateRequest request,
                                      Ingestion.DataMutateResponse response,
                                      Supplier> channelDataSupplier,
                                      RpcRequestCallback rpcRequestCallback,
                                      Listener listener) {
    return new StreamTask(session, retryThreadPool, retryRequestCnt, table, context,
        internalMs, retryMode, request, response, channelDataSupplier, rpcRequestCallback, listener);
  }

  @Override
  public AbstractTask buildFlushTask(long requestId, ClientContext clientContext, Buffer buffer,
                                     Supplier> channelDataSupplier,
                                     RpcRequestCallback requestCallback, Listener listener) {
    return new FlushTask(requestId, clientContext, buffer, channelDataSupplier, requestCallback, listener);
  }

  @Override
  public void closeHook(RpcProxy client) throws IOException {
    ClientContext context = client.getClientContext();
    client.checkTableExists(context.instanceId(), context.workspace(), stream.getTable().getSchemaName(),
        stream.getTable().getTableName(), stream.getTable().getTableId());
  }
}