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

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class MockStreamMeta {

  private static final Map, Ingestion.IGSTableType> tableNameMaps = new HashMap<>();
  private static final Map streamSchemaMaps = new HashMap<>();

  static {
    // tableNameMaps
    tableNameMaps.put(new Tuple2<>("mock-db", "mock-tbl-common"), Ingestion.IGSTableType.NORMAL);
    tableNameMaps.put(new Tuple2<>("mock-db", "mock-tbl-cluster"), Ingestion.IGSTableType.CLUSTER);
    tableNameMaps.put(new Tuple2<>("mock-db", "mock-tbl-pk"), Ingestion.IGSTableType.ACID);

    // build stream schema base.
    IngestionV2.StreamSchema streamSchema = IngestionV2.StreamSchema.newBuilder()
        .addDataFields(IngestionV2.DataField.newBuilder()
            .setName("col1")
            .setType(DataType.newBuilder().setCategory(DataTypeCategory.INT32).setNullable(false).setFieldId(1).build())
            .build())
        .addDataFields(IngestionV2.DataField.newBuilder()
            .setName("col2")
            .setType(DataType.newBuilder().setCategory(DataTypeCategory.INT32).setNullable(true).setFieldId(2).build())
            .build())
        .addDataFields(IngestionV2.DataField.newBuilder()
            .setName("col3")
            .setType(DataType.newBuilder().setCategory(DataTypeCategory.STRING).setNullable(true).setFieldId(3).build())
            .build()).build();

    // build DistributionSpec.
    IngestionV2.DistributionSpec distributionSpec = IngestionV2.DistributionSpec.newBuilder()
        .addFieldIds(1)
        .setNumBuckets(1)
        .build();

    // build PrimaryKeySpec.
    IngestionV2.PrimaryKeySpec primaryKeySpec = IngestionV2.PrimaryKeySpec.newBuilder()
        .addFieldIds(1)
        .build();

    // build common table schema & info.
    streamSchemaMaps.put(Ingestion.IGSTableType.NORMAL, streamSchema);

    // build cluster table schema & info.
    IngestionV2.StreamSchema clusterStreamSchema = streamSchema.toBuilder()
        .setDistSpec(distributionSpec).build();
    streamSchemaMaps.put(Ingestion.IGSTableType.CLUSTER, clusterStreamSchema);

    // build pk table schema & info.
    IngestionV2.StreamSchema pkStreamSchema = streamSchema.toBuilder()
        .setDistSpec(distributionSpec)
        .setPrimaryKeySpec(primaryKeySpec).build();
    streamSchemaMaps.put(Ingestion.IGSTableType.ACID, pkStreamSchema);
  }

  public static IngestionV2.StreamSchema getTableMeta(Ingestion.IGSTableType tableType) {
    return streamSchemaMaps.get(tableType);
  }

  public static Ingestion.IGSTableType getTableType(String dbName, String tblName) {
    return tableNameMaps.get(new Tuple2<>(dbName, tblName));
  }

  public static Tuple2 getTableNameByType(Ingestion.IGSTableType tableType) {
    Tuple2 tuple2 = null;
    for (Map.Entry, Ingestion.IGSTableType> entry : tableNameMaps.entrySet()) {
      if (entry.getValue() == tableType) {
        tuple2 = entry.getKey();
        break;
      }
    }
    return tuple2;
  }

  public static List getAllIGSTableTypes() {
    return new ArrayList() {{
      add(Ingestion.IGSTableType.NORMAL);
      add(Ingestion.IGSTableType.CLUSTER);
      add(Ingestion.IGSTableType.ACID);
    }};
  }

  public static List> getAllTables() {
    return new ArrayList<>(tableNameMaps.keySet());
  }
}