com.clickzetta.platform.bulkload.BulkLoadStreamImpl 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.bulkload;

import com.clickzetta.platform.client.CZClient;
import com.clickzetta.platform.client.Table;
import com.clickzetta.platform.client.api.BulkLoadCommitOptions;
import com.clickzetta.platform.client.api.BulkLoadOperation;
import com.clickzetta.platform.client.api.BulkLoadState;
import com.clickzetta.platform.client.api.BulkLoadStream;
import com.clickzetta.platform.client.api.BulkLoadWriter;
import com.clickzetta.platform.common.Schema;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.util.List;
import java.util.Optional;

public class BulkLoadStreamImpl implements BulkLoadStream {

  private static final Logger LOG = LoggerFactory.getLogger(BulkLoadStreamImpl.class);
  private final CZClient client;
  private BulkLoadMetaData metaData;
  private volatile String location = null;
  public BulkLoadStreamImpl(CZClient client, BulkLoadMetaData metaData) {
    this.client = client;
    this.metaData = metaData;
    LOG.info("BulkLoadStream is created: {}", metaData.getStreamId());
  }

  @Override
  public String getStreamId() {
    return metaData.getStreamId();
  }

  @Override
  public BulkLoadOperation getOperation() {
    return metaData.getOperation();
  }

  @Override
  public BulkLoadState getStreamState() throws IOException {
    // Issue RPC call to get real-time stream state.
    metaData = client.getBulkLoadStreamMetaData(
      metaData.getSchemaName(), metaData.getTableName(), metaData.getStreamId());
    return metaData.getState();
  }

  @Override
  public String getSqlErrorMsg() {
    return metaData.getSqlErrorMsg();
  }

  @Override
  public Schema getSchema() {
    return getTable().getSchema();
  }

  @Override
  public Table getTable() {
    return metaData.getArrowTable();
  }

  @Override
  public Optional getPartitionSpecs() {
    return metaData.getPartitionSpecs();
  }

  @Override
  public List getRecordKeys() {
    return metaData.getRecordKeys();
  }

  @Override
  public BulkLoadWriter openWriter(int partitionId) throws IOException {
    LOG.info("Open writer for BulkLoadStream {} with partitionId {}", metaData.getStreamId(), partitionId);

    BulkLoadConfig config = client.openBulkLoadStreamWriterV2(
      metaData.getInstanceId(),
      metaData.getWorkspace(),
      metaData.getSchemaName(),
      metaData.getTableName(),
      metaData.getStreamId(),
      partitionId);

    if (location == null) {
      synchronized (this) {
        if (location == null) {
          location = config.getStagingConfig().getPath();
        }
      }
    }

    return new BulkLoadWriterImpl(client, metaData, config, partitionId);
  }

  @Override
  public void commit(BulkLoadCommitOptions options) throws IOException {
    LOG.info("Committing BulkLoadStream {}, location: {}", metaData.getStreamId(), location != null ? location : "");

    client.commitBulkLoadStreamV2(
      metaData.getInstanceId(),
      metaData.getWorkspace(),
      metaData.getSchemaName(),
      metaData.getTableName(),
      metaData.getStreamId(),
      options.getWorkspace(),
      options.getVc(),
      BulkLoadCommitMode.COMMIT_STREAM);
  }

  @Override
  public void cancel() throws IOException {
    client.commitBulkLoadStreamV2(
      metaData.getInstanceId(),
      metaData.getWorkspace(),
      metaData.getSchemaName(),
      metaData.getTableName(),
      metaData.getStreamId(),
      "",
      "",
      BulkLoadCommitMode.ABORT_STREAM);
  }

  @Override
  public void close() throws IOException {
    // No-op
  }
}