com.clickzetta.client.BulkloadStream 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.client;

import com.clickzetta.platform.client.Table;
import com.clickzetta.platform.client.api.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.sql.SQLException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

public class BulkloadStream implements RowStream {
    private static Logger logger = LoggerFactory.getLogger(BulkloadStream.class);
    private ClickZettaClient client;
    private BulkLoadStream stream;
    private Map writerMap = new ConcurrentHashMap<>();
    private boolean isClosed = false;

    BulkloadStream(ClickZettaClient client, String schema, String table, BulkLoadOptions options)
            throws IOException {
        this.client = client;
        this.stream = client.getIgsClient().createBulkLoadStream(schema, table, options);
    }

    BulkloadStream(ClickZettaClient client, String streamId) throws IOException {
        this.client = client;
        this.stream = client.getIgsClient().getBulkLoadStream(null, null, streamId);
    }

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

    @Override
    public Row createRow() throws SQLException {
        return createRow(0);
    }

    /**
     * Not Suggested in BulkloadStream.
     * @param operator
     * @return
     * @throws SQLException
     */
    @Override
    public Row createRow(Stream.Operator operator) throws SQLException {
        return createRow();
    }

    /**
     * Create a Row with partitionId for BulkloadStream.
     * @param partitionId partitionId is used to identify which partition the row belongs to.
     * @return
     * @throws SQLException
     */
    @Override
    public Row createRow(int partitionId) throws SQLException {
        writerMap.computeIfAbsent(partitionId, k -> {
            try {
                return stream.openWriter(partitionId);
            } catch (IOException e) {
                e.printStackTrace();
            }
            return null;
        });
        return writerMap.get(partitionId).createRow();
    }

    /**
     * Not Avaliable in BulkloadStream.
     * @param row
     * @throws IOException
     * @throws SQLException
     */
    @Override
    public void apply(Row row) throws IOException, SQLException {
        apply(row, 0);
    }

    /**
     * Write a Row into Files which belong to the partitionId.
     * @param row
     * @param partitionId partitionId is used to identify which partition the row belongs to.
     * @throws SQLException
     */
    @Override
    public void apply(Row row, int partitionId) throws SQLException {
        writerMap.computeIfPresent(partitionId, (k, v) -> {
            try {
                v.write(row);
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
            return v;
        });
    }

    @Override
    public void close() throws IOException, SQLException {
        close(3600);
    }

    /**
     * Flush data to server for BulkloadStream with timeout.
     * @throws IOException
     * @throws SQLException
     */

    public void close(int timeout) throws IOException, SQLException {
        if (isClosed) {
            logger.warn("Bulkload stream already closed!");
            return;
        }
        if (timeout <= 0) {
            throw new SQLException("Bulkload stream close timeout should be positive");
        }
        isClosed = true;
        // close writer
        if (writerMap != null) {
            writerMap.forEach((k, v) -> {
                try {
                    v.close();
                } catch (IOException e) {
                    throw new RuntimeException(e);
                }
            });
            writerMap.clear();
        }

        // bulkload commit
        BulkLoadCommitOptions options = BulkLoadCommitOptions.newBuilder()
                .withWorkspace(client.getWorkspace())
                .withVc(client.getVCluster())
                .build();
        stream.commit(options);

        // wait for bulkload commit finish
        int waitInterval = 1000;
        int maxWaitCount = timeout * 1000 / waitInterval;
        while ((stream.getStreamState() == BulkLoadState.SEALED ||
                stream.getStreamState() == BulkLoadState.COMMIT_SUBMITTED) && maxWaitCount-- > 0) {
            try {
                Thread.sleep(waitInterval);
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
                logger.warn("Thread interrupted while sleeping", e);
            }
        }

        if (maxWaitCount <= 0) {
            throw new SQLException(stream.getStreamId() + " commit not finished after waiting " + timeout + " seconds");
        } else {
            if (stream.getStreamState() != BulkLoadState.COMMIT_SUCCESS) {
                throw new SQLException(stream.getStreamId() + " commit failed, error message: " + stream.getSqlErrorMsg());
            }
            logger.info("Bulkload {} commit finished successfully, state: {}",stream.getStreamId(), stream.getStreamState());
        }
    }

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

    @Override
    public StreamState getState() {
        try {
            BulkLoadState state = stream.getStreamState();
            if (state == BulkLoadState.COMMIT_SUCCESS) {
                return StreamState.SUCCESS;
            } else if (state == BulkLoadState.COMMIT_FAILED) {
                return StreamState.FAILED;
            } else {
                return StreamState.RUNNING;
            }
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

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