com.clickzetta.platform.test.BaseIgsController 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.test;

import com.clickzetta.platform.client.proxy.RequestConstructor;
import com.clickzetta.platform.client.proxy.ResponseConstructor;
import com.google.common.annotations.VisibleForTesting;
import com.google.protobuf.AbstractMessage;
import cz.proto.ingestion.IGSControllerServiceGrpc;
import cz.proto.ingestion.Ingestion;
import io.grpc.stub.StreamObserver;

import java.io.IOException;
import java.util.List;
import java.util.function.Supplier;

@VisibleForTesting
public class BaseIgsController extends IGSControllerServiceGrpc.IGSControllerServiceImplBase {

  private Supplier> getWorkersSupplier;
  private BaseIgsControllerV2 controllerV2;

  public BaseIgsController(Supplier> getWorkersSupplier) {
    this.getWorkersSupplier = getWorkersSupplier;
  }

  public void setControllerV2(BaseIgsControllerV2 controllerV2) {
    this.controllerV2 = controllerV2;
  }

  @Override
  public void getTableMeta(Ingestion.GetTableMetaRequest request,
                           StreamObserver responseObserver) {
    try {
      responseObserver.onNext(buildGatewayRpcCallResponse(
          Ingestion.MethodEnum.GET_TABLE_META, RequestConstructor.toReqString(request)));
    } catch (IOException e) {
      throw new RuntimeException(e);
    }
    responseObserver.onCompleted();
  }

  @Override
  public void createTablet(Ingestion.ControllerCreateTabletRequest request,
                           StreamObserver responseObserver) {
    try {
      responseObserver.onNext(buildGatewayRpcCallResponse(
          Ingestion.MethodEnum.CREATE_TABLET, RequestConstructor.toReqString(request)));
    } catch (IOException e) {
      throw new RuntimeException(e);
    }
    responseObserver.onCompleted();
  }

  @Override
  public void getMutateWorkers(Ingestion.GetMutateWorkersRequest request,
                               StreamObserver responseObserver) {
    try {
      responseObserver.onNext(buildGatewayRpcCallResponse(
          Ingestion.MethodEnum.GET_MUTATE_WORKER, RequestConstructor.toReqString(request)));
    } catch (IOException e) {
      throw new RuntimeException(e);
    }
    responseObserver.onCompleted();
  }

  @Override
  public void commitTablet(Ingestion.CommitTabletRequest request,
                           StreamObserver responseObserver) {
    try {
      responseObserver.onNext(buildGatewayRpcCallResponse(
          Ingestion.MethodEnum.COMMIT_TABLET, RequestConstructor.toReqString(request)));
    } catch (IOException e) {
      throw new RuntimeException(e);
    }
    responseObserver.onCompleted();
  }

  @Override
  public void dropTablet(Ingestion.DropTabletRequest request,
                         StreamObserver responseObserver) {
    try {
      responseObserver.onNext(buildGatewayRpcCallResponse(
          Ingestion.MethodEnum.DROP_TABLET, RequestConstructor.toReqString(request)));
    } catch (IOException e) {
      throw new RuntimeException(e);
    }
    responseObserver.onCompleted();
  }

  @Override
  public void checkTableExists(Ingestion.CheckTableExistsRequest request,
                               StreamObserver responseObserver) {
    try {
      responseObserver.onNext(buildGatewayRpcCallResponse(
          Ingestion.MethodEnum.CHECK_TABLE_EXISTS, RequestConstructor.toReqString(request)));
    } catch (IOException e) {
      throw new RuntimeException(e);
    }
    responseObserver.onCompleted();
  }

  @Override
  public void gatewayRpcCall(Ingestion.GatewayRequest request, StreamObserver responseObserver) {
    Ingestion.MethodEnum method = Ingestion.MethodEnum.forNumber(request.getMethodEnumValue());
    String message = request.getMessage();
    AbstractMessage response = buildGatewayRpcCallResponse(method, message);
    try {
      responseObserver.onNext(Ingestion.GatewayResponse.newBuilder()
          .setStatus(Ingestion.ResponseStatus.newBuilder().setCode(Ingestion.Code.SUCCESS).build())
          .setMessage(ResponseConstructor.toRespString(response)).build());
    } catch (IOException e) {
      throw new RuntimeException(e);
    }
    responseObserver.onCompleted();
  }

  public  T buildGatewayRpcCallResponse(Ingestion.MethodEnum method, String message) {
    try {
      switch (method) {
        case GET_TABLE_META: {
          Ingestion.GetTableMetaRequest request = RequestConstructor.getRequest(Ingestion.MethodEnum.GET_TABLE_META, message);
          Ingestion.IGSTableType igsTableType = MockMeta.getTableType(request.getSchemaName(), request.getTableName());
          return (T) Ingestion.GetTableMetaResponse.newBuilder()
              .setTableMeta(MockMeta.getTableMeta(igsTableType))
              .setTableType(igsTableType)
              .setInstanceId(request.getInstanceId()).build();
        }
        case CREATE_TABLET: {
          return (T) Ingestion.CreateTabletResponse.newBuilder()
              .setStatus(Ingestion.ResponseStatus.newBuilder()
                  .setCode(Ingestion.Code.SUCCESS).build()).build();
        }
        case GET_MUTATE_WORKER: {
          return (T) Ingestion.GetMutateWorkersResponse.newBuilder().addAllTuple(getWorkersSupplier.get()).build();
        }
        case COMMIT_TABLET: {
          return (T) Ingestion.CommitTabletResponse.newBuilder()
              .setStatus(Ingestion.ResponseStatus.newBuilder()
                  .setCode(Ingestion.Code.SUCCESS).build()).build();
        }
        case DROP_TABLET: {
          return (T) Ingestion.DropTabletResponse.newBuilder()
              .setStatus(Ingestion.ResponseStatus.newBuilder()
                  .setCode(Ingestion.Code.SUCCESS).build()).build();
        }
        case CHECK_TABLE_EXISTS: {
          return (T) Ingestion.CheckTableExistsResponse.newBuilder()
              .setStatus(Ingestion.ResponseStatus.newBuilder()
                  .setCode(Ingestion.Code.SUCCESS).build()).build();
        }
        case CREATE_OR_GET_STREAM_V2:
        case CLOSE_STREAM_V2:
        case GET_ROUTE_WORKER_V2:
        case CREATE_BULK_LOAD_STREAM_V2:
        case GET_BULK_LOAD_STREAM_V2:
        case COMMIT_BULK_LOAD_STREAM_V2:
        case OPEN_BULK_LOAD_STREAM_WRITER_V2:
        case FINISH_BULK_LOAD_STREAM_WRITER_V2:
        case GET_BULK_LOAD_STREAM_STS_TOKEN_V2:
        case COMMIT_V2:
        case ASYNC_COMMIT_V2:
        case CHECK_COMMIT_RESULT_V2: {
          return controllerV2.buildGatewayRpcCallResponse(method, message);
        }
        case GATEWAY_RPC_CALL:
        default:
          throw new UnsupportedOperationException("not support method " + method);
      }
    } catch (IOException ioe) {
      throw new RuntimeException(ioe);
    }
  }
}