com.clickzetta.platform.test.MockBulkLoadStream 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 cz.proto.FileFormatTypeOuterClass;
import cz.proto.ingestion.Ingestion;
import cz.proto.ingestion.v2.IngestionV2;
import scala.Tuple2;

import java.util.HashMap;
import java.util.Map;
import java.util.UUID;

public class MockBulkLoadStream {

  public enum StreamStatus {
    CREATE_ABLE,
    CREATE_FAILED,
    JOIN_ABLE,
    NON_JOIN_ABLE,
    GOOD_TO_COMMIT,
    ALREADY_COMMITTED
  }

  public static String MOCK_SCHEMA_NAME = "mock-db";
  public static String MOCK_CREATE_FAILED_TABLE_NAME = "mock-tbl-create-failed";
  public static String MOCK_NON_JOIN_ABLE_TABLE_NAME = "mock-tbl-non-join-able";
  public static String MOCK_COMMITTED_TABLE_NAME = "mock-tbl-committed";

  private static final Map, StreamStatus> tableNameMaps = new HashMap<>();

  static {
    for (Tuple2 tableNames : MockMeta.getAllTables()) {
      tableNameMaps.put(tableNames, StreamStatus.CREATE_ABLE);
    }
    tableNameMaps.put(new Tuple2<>(MOCK_SCHEMA_NAME, MOCK_CREATE_FAILED_TABLE_NAME), StreamStatus.CREATE_FAILED);
    tableNameMaps.put(new Tuple2<>(MOCK_SCHEMA_NAME, MOCK_NON_JOIN_ABLE_TABLE_NAME), StreamStatus.NON_JOIN_ABLE);
    tableNameMaps.put(new Tuple2<>(MOCK_SCHEMA_NAME, MOCK_COMMITTED_TABLE_NAME), StreamStatus.ALREADY_COMMITTED);
  }

  public static StreamStatus getStreamStatus(IngestionV2.TableIdentifier identifier) {
    return tableNameMaps.get(new Tuple2<>(identifier.getSchemaName(), identifier.getTableName()));
  }

  private static IngestionV2.BulkLoadStreamState getStreamStateV2(StreamStatus status) {
    switch (status) {
      case CREATE_ABLE:
      case JOIN_ABLE:
      case GOOD_TO_COMMIT:
        return IngestionV2.BulkLoadStreamState.BL_CREATED;
      case CREATE_FAILED:
        return IngestionV2.BulkLoadStreamState.UNRECOGNIZED;
      case NON_JOIN_ABLE:
        return IngestionV2.BulkLoadStreamState.BL_COMMIT_SUBMITTED;
      case ALREADY_COMMITTED:
        return IngestionV2.BulkLoadStreamState.BL_COMMIT_SUCCESS;
      default:
        throw new RuntimeException("unknown stream status: " + status);
    }
  }

  public static String getStreamId(StreamStatus status) {
    switch (status) {
      case CREATE_ABLE:
      case JOIN_ABLE:
      case GOOD_TO_COMMIT:
        return "mock-bulk-load-stream-id";
      case CREATE_FAILED:
        return "mock-create-failed-stream-id";
      case NON_JOIN_ABLE:
        return "mock-non-join-able-stream-id";
      case ALREADY_COMMITTED:
        return "mock-already-committed-stream-id";
      default:
        throw new RuntimeException("unknown stream status: " + status);
    }
  }

  public static IngestionV2.BulkLoadStreamWriterConfig getWriteConfig(IngestionV2.TableIdentifier identifier,
                                                                      boolean useCosStagingPath) {
    if (useCosStagingPath) {
      return IngestionV2.BulkLoadStreamWriterConfig.newBuilder()
        .setStagingPath(IngestionV2.StagingPathInfo.newBuilder()
          .setCosPath(IngestionV2.CosStagingPathInfo.newBuilder()
            .setPath("/tmp/" + UUID.randomUUID() + "_staging_path")
            .setStsAkId("mock_ak_id")
            .setStsAkSecret("mock_ak_secret")
            .setStsToken("mock_sts_token")
            .setCosRegion("mock_oss_endpoint")
            .setCosExpireTime(Long.MAX_VALUE)
            .build())
          .build())
        .setFileFormat(FileFormatTypeOuterClass.FileFormatType.PARQUET)
        .setMaxNumRowsPerFile(10)
        .setMaxSizeInBytesPerFile(256 * 1024 * 1024)
        .build();
    } else {
      return IngestionV2.BulkLoadStreamWriterConfig.newBuilder()
        .setStagingPath(IngestionV2.StagingPathInfo.newBuilder()
          .setOssPath(IngestionV2.OssStagingPathInfo.newBuilder()
            .setPath("/tmp/" + UUID.randomUUID() + "_staging_path")
            .setStsAkId("mock_ak_id")
            .setStsAkSecret("mock_ak_secret")
            .setStsToken("mock_sts_token")
            .setOssEndpoint("mock_oss_endpoint")
            .setOssExpireTime(Long.MAX_VALUE)
            .build())
          .build())
        .setFileFormat(FileFormatTypeOuterClass.FileFormatType.PARQUET)
        .setMaxNumRowsPerFile(10)
        .setMaxSizeInBytesPerFile(256 * 1024 * 1024)
        .build();
    }
  }

  public static IngestionV2.BulkLoadStreamInfo getStreamInfo(IngestionV2.TableIdentifier identifier,
                                                             IngestionV2.BulkLoadStreamOperation mode) {
    Ingestion.IGSTableType tableType = MockMeta.getTableType(identifier.getSchemaName(), identifier.getTableName());
    StreamStatus status = getStreamStatus(identifier);
    return IngestionV2.BulkLoadStreamInfo.newBuilder()
      .setIdentifier(identifier)
      .setStreamId(getStreamId(status))
      .setStreamSchema(MockStreamMeta.getTableMeta(tableType))
      .setOperation(mode)
      .setStreamState(getStreamStateV2(status))
      .build();
  }

}