com.clickzetta.platform.test.BaseIgsControllerV2 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.google.common.annotations.VisibleForTesting;
import com.google.protobuf.AbstractMessage;
import cz.proto.ingestion.Ingestion;
import cz.proto.ingestion.v2.IngestionControllerServiceGrpc;
import cz.proto.ingestion.v2.IngestionV2;
import io.grpc.stub.StreamObserver;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Supplier;

@VisibleForTesting
public class BaseIgsControllerV2 extends IngestionControllerServiceGrpc.IngestionControllerServiceImplBase {

  private Supplier> getWorkersSupplier;

  // redirect status & max retry times.
  private volatile Map needRedirectMap;

  public BaseIgsControllerV2(Supplier> getWorkersSupplier) {
    this.getWorkersSupplier = getWorkersSupplier;
    this.needRedirectMap = new ConcurrentHashMap<>();
  }

  public void setNeedRedirectMap(IngestionV2.Code code, int times) {
    this.needRedirectMap.put(code, times);
  }

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

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

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

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

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

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

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

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

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

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

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

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

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

  public  T buildGatewayRpcCallResponse(Ingestion.MethodEnum method, String message) {
    try {
      switch (method) {
        case CREATE_OR_GET_STREAM_V2: {
          IngestionV2.CreateOrGetStreamRequest request = RequestConstructor.getRequest(Ingestion.MethodEnum.CREATE_OR_GET_STREAM_V2, message);
          IngestionV2.TableIdentifier tableIdentifier = request.getTableIdent();
          String schemaName = tableIdentifier.getSchemaName();
          String tableName = tableIdentifier.getTableName();
          IngestionV2.StreamSchema streamSchema = null;

          // get base stream schema or struct stream schema.
          if (schemaName.contains("json")) {
            Ingestion.IGSTableType igsTableType = MockJsonStructStreamMeta.getTableType(schemaName, tableName);
            streamSchema = MockJsonStructStreamMeta.getTableMeta(igsTableType);
          } else if (!schemaName.contains("struct")) {
            Ingestion.IGSTableType igsTableType = MockStreamMeta.getTableType(schemaName, tableName);
            streamSchema = MockStreamMeta.getTableMeta(igsTableType);
          } else {
            Ingestion.IGSTableType igsTableType = MockStructStreamMeta.getTableType(schemaName, tableName);
            streamSchema = MockStructStreamMeta.getTableMeta(igsTableType);
          }
          if (needRedirectMap.containsKey(IngestionV2.Code.NEED_REDIRECT)) {
            int times = needRedirectMap.get(IngestionV2.Code.NEED_REDIRECT);
            if (times > 0) {
              needRedirectMap.put(IngestionV2.Code.NEED_REDIRECT, times - 1);
              return (T) IngestionV2.CreateOrGetStreamResponse.newBuilder()
                  .setStatus(IngestionV2.ResponseStatus.newBuilder()
                      .setCode(IngestionV2.Code.NEED_REDIRECT).build())
                  .build();
            }
          }
          return (T) IngestionV2.CreateOrGetStreamResponse.newBuilder()
              .setStatus(IngestionV2.ResponseStatus.newBuilder()
                  .setCode(IngestionV2.Code.SUCCESS).build())
              .setDataSchema(streamSchema)
              .setAlreadyExists(true)
              .build();
        }
        case CLOSE_STREAM_V2: {
          IngestionV2.CloseStreamRequest request = RequestConstructor.getRequest(Ingestion.MethodEnum.CLOSE_STREAM_V2, message);
          IngestionV2.TableIdentifier tableIdentifier = request.getTableIdent();
          return (T) IngestionV2.CloseStreamResponse.newBuilder()
              .setStatus(IngestionV2.ResponseStatus.newBuilder()
                  .setCode(IngestionV2.Code.SUCCESS).build())
              .build();
        }
        case GET_ROUTE_WORKER_V2: {
          List list = new ArrayList<>();
          for (Ingestion.HostPortTuple hostPortTuple : getWorkersSupplier.get()) {
            list.add(IngestionV2.HostPortTuple.newBuilder()
                .setHost(hostPortTuple.getHost())
                .setPort(hostPortTuple.getPort()).build());
          }
          return (T) IngestionV2.GetRouteWorkersResponse.newBuilder().addAllTuple(list).build();
        }
        case COMMIT_V2: {
          return (T) IngestionV2.CommitResponse.newBuilder()
              .setStatus(IngestionV2.ResponseStatus.newBuilder()
                  .setCode(IngestionV2.Code.SUCCESS).build())
              .build();
        }
        case ASYNC_COMMIT_V2:
          return (T) IngestionV2.CommitResponse.newBuilder()
              .setStatus(IngestionV2.ResponseStatus.newBuilder()
                  .setCode(IngestionV2.Code.SUCCESS).build())
              .setCommitId(1001)
              .build();
        case CHECK_COMMIT_RESULT_V2:
          return (T) IngestionV2.CheckCommitResultResponse.newBuilder()
              .setStatus(IngestionV2.ResponseStatus.newBuilder()
                  .setCode(IngestionV2.Code.SUCCESS).build())
              .setFinished(true)
              .build();
        case SCHEMA_CHANGE:
          return (T) IngestionV2.SchemaChangeResponse.newBuilder()
              .setStatus(IngestionV2.ResponseStatus.newBuilder()
                  .setCode(IngestionV2.Code.SUCCESS).build())
              .build();

        case CREATE_BULK_LOAD_STREAM_V2: {
          IngestionV2.CreateBulkLoadStreamRequest request =
              RequestConstructor.getRequest(Ingestion.MethodEnum.CREATE_BULK_LOAD_STREAM_V2, message);
          IngestionV2.TableIdentifier identifier = request.getIdentifier();
          MockBulkLoadStream.StreamStatus status = MockBulkLoadStream.getStreamStatus(identifier);
          if (status == MockBulkLoadStream.StreamStatus.CREATE_ABLE) {
            IngestionV2.BulkLoadStreamInfo info = MockBulkLoadStream.getStreamInfo(identifier, request.getOperation());
            return (T) IngestionV2.CreateBulkLoadStreamResponse.newBuilder()
                .setStatus(IngestionV2.ResponseStatus.newBuilder()
                    .setCode(IngestionV2.Code.SUCCESS).build())
                .setInfo(info)
                .build();
          } else {
            return (T) IngestionV2.CreateBulkLoadStreamResponse.newBuilder()
                .setStatus(IngestionV2.ResponseStatus.newBuilder()
                    .setCode(IngestionV2.Code.FAILED).build())
                .build();
          }
        }

        case GET_BULK_LOAD_STREAM_V2: {
          IngestionV2.GetBulkLoadStreamRequest request =
              RequestConstructor.getRequest(Ingestion.MethodEnum.GET_BULK_LOAD_STREAM_V2, message);
          IngestionV2.TableIdentifier identifier = request.getIdentifier();
          MockBulkLoadStream.StreamStatus status = MockBulkLoadStream.getStreamStatus(identifier);
          if (status == MockBulkLoadStream.StreamStatus.JOIN_ABLE
              || status == MockBulkLoadStream.StreamStatus.CREATE_ABLE) {
            IngestionV2.BulkLoadStreamInfo info =
                MockBulkLoadStream.getStreamInfo(identifier, IngestionV2.BulkLoadStreamOperation.BL_APPEND);
            return (T) IngestionV2.GetBulkLoadStreamResponse.newBuilder()
                .setStatus(IngestionV2.ResponseStatus.newBuilder()
                    .setCode(IngestionV2.Code.SUCCESS).build())
                .setInfo(info)
                .build();
          } else {
            return (T) IngestionV2.GetBulkLoadStreamResponse.newBuilder()
                .setStatus(IngestionV2.ResponseStatus.newBuilder()
                    .setCode(IngestionV2.Code.FAILED).build())
                .build();
          }
        }

        case COMMIT_BULK_LOAD_STREAM_V2: {
          return (T) IngestionV2.CommitBulkLoadStreamResponse.newBuilder()
              .setStatus(IngestionV2.ResponseStatus.newBuilder()
                  .setCode(IngestionV2.Code.SUCCESS).build()).build();
        }

        case OPEN_BULK_LOAD_STREAM_WRITER_V2: {
          IngestionV2.OpenBulkLoadStreamWriterRequest request =
              RequestConstructor.getRequest(Ingestion.MethodEnum.OPEN_BULK_LOAD_STREAM_WRITER_V2, message);
          IngestionV2.TableIdentifier identifier = request.getIdentifier();
          MockBulkLoadStream.StreamStatus status = MockBulkLoadStream.getStreamStatus(identifier);
          if (status == MockBulkLoadStream.StreamStatus.JOIN_ABLE
              || status == MockBulkLoadStream.StreamStatus.CREATE_ABLE) {
            IngestionV2.BulkLoadStreamWriterConfig config = MockBulkLoadStream.getWriteConfig(identifier, true);
            return (T) IngestionV2.OpenBulkLoadStreamWriterResponse.newBuilder()
                .setStatus(IngestionV2.ResponseStatus.newBuilder()
                    .setCode(IngestionV2.Code.SUCCESS).build())
                .setConfig(config)
                .build();
          } else {
            return (T) IngestionV2.OpenBulkLoadStreamWriterResponse.newBuilder()
                .setStatus(IngestionV2.ResponseStatus.newBuilder()
                    .setCode(IngestionV2.Code.FAILED).build())
                .build();
          }
        }

        case FINISH_BULK_LOAD_STREAM_WRITER_V2: {
          return (T) IngestionV2.FinishBulkLoadStreamWriterResponse.newBuilder()
              .setStatus(IngestionV2.ResponseStatus.newBuilder()
                  .setCode(IngestionV2.Code.SUCCESS)
                  .build())
              .build();
        }

        case GET_BULK_LOAD_STREAM_STS_TOKEN_V2: {
          IngestionV2.GetBulkLoadStreamStsTokenRequest request =
              RequestConstructor.getRequest(Ingestion.MethodEnum.GET_BULK_LOAD_STREAM_STS_TOKEN_V2, message);
          IngestionV2.TableIdentifier identifier = request.getIdentifier();
          MockBulkLoadStream.StreamStatus status = MockBulkLoadStream.getStreamStatus(identifier);
          if (status == MockBulkLoadStream.StreamStatus.JOIN_ABLE
              || status == MockBulkLoadStream.StreamStatus.CREATE_ABLE) {
            IngestionV2.BulkLoadStreamWriterConfig config = MockBulkLoadStream.getWriteConfig(identifier, false);
            return (T) IngestionV2.GetBulkLoadStreamStsTokenResponse.newBuilder()
                .setStatus(IngestionV2.ResponseStatus.newBuilder()
                    .setCode(IngestionV2.Code.SUCCESS).build())
                .setStagingPath(config.getStagingPath())
                .build();
          } else {
            return (T) IngestionV2.GetBulkLoadStreamStsTokenResponse.newBuilder()
                .setStatus(IngestionV2.ResponseStatus.newBuilder()
                    .setCode(IngestionV2.Code.FAILED).build())
                .build();
          }
        }

        case GATEWAY_RPC_CALL:
        default:
          throw new UnsupportedOperationException("not support method " + method);
      }
    } catch (IOException ioe) {
      throw new RuntimeException(ioe);
    }
  }
}