com.huaweicloud.dws.client.handler.PutActionHandler Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of dws-client Show documentation
There is a newer version: 2.0.0-r0
package com.huaweicloud.dws.client.handler;

import com.huaweicloud.dws.client.DwsConfig;
import com.huaweicloud.dws.client.action.PutAction;
import com.huaweicloud.dws.client.exception.DwsClientException;
import com.huaweicloud.dws.client.exception.DwsClientRecordException;
import com.huaweicloud.dws.client.executor.CopyMergeExecutor;
import com.huaweicloud.dws.client.executor.CopyUpdateExecutor;
import com.huaweicloud.dws.client.executor.CopyUpsertExecutor;
import com.huaweicloud.dws.client.executor.DeleteExecutor;
import com.huaweicloud.dws.client.executor.UpsertExecutor;
import com.huaweicloud.dws.client.model.OperationType;
import com.huaweicloud.dws.client.model.Record;
import com.huaweicloud.dws.client.model.TableName;
import com.huaweicloud.dws.client.model.WriteMode;
import com.huaweicloud.dws.client.util.JdbcUtil;
import com.huaweicloud.dws.client.worker.ConnectionProvider;
import lombok.extern.slf4j.Slf4j;

import java.sql.Connection;
import java.sql.SQLException;
import java.util.List;
import java.util.function.Consumer;
import java.util.stream.Collectors;

/**
 * @ProjectName: dws-connector
 * @ClassName: PutActionHandler
 * @Description:
 * @Date: 2023/1/16 15:14
 * @Version: 1.0
 */
@Slf4j
public class PutActionHandler extends AbstractActionHandler {
    final DwsConfig config;
    final ConnectionProvider connectionProvider;
    final UpsertExecutor upsertExecutor = new UpsertExecutor();
    final DeleteExecutor deleteExecutor = new DeleteExecutor();
    final CopyMergeExecutor copyMergeExecutor = new CopyMergeExecutor();
    final CopyUpsertExecutor copyUpsertExecutor = new CopyUpsertExecutor();
    final CopyUpdateExecutor copyUpdateExecutor = new CopyUpdateExecutor();

    public PutActionHandler(DwsConfig config, ConnectionProvider connectionProvider) {
        this.config = config;
        this.connectionProvider = connectionProvider;
    }

    @Override
    public void handle(PutAction action) throws DwsClientException {
        List records = action.getRecords();
        if (records == null || records.isEmpty()) {
            log.debug("current action records is empty.");
            action.getFuture().complete(null);
            return;
        }
        try {
            doHandle(action, records);
        } catch (DwsClientException e) {
            throw new DwsClientRecordException(e, records);
        }
        Consumer> consumer = config.getSuccessFunction();
        if (consumer != null) {
            consumer.accept(records);
        }
    }

    private void doHandle(PutAction action, List records) throws DwsClientException {
        TableName tableName = records.get(0).getTableSchema().getTableName();
        WriteMode writeMode = config.getTableConfig(tableName).getWriteMode();
        Connection connection = null;
        try {
            connection = connectionProvider.getOrInitConnection();
        } catch (Exception e) {
            log.error("get connection fail", e);
            throw DwsClientException.fromException(e);
        }
        try {
            if (connection.getAutoCommit()) {
                connection.setAutoCommit(false);
            }
            JdbcUtil.executeSql(connection, String.format("SET statement_timeout = %s", config.getTimeOutMs()));
            // 先删除
            deleteExecutor.execute(records.stream().filter(item -> item.getType() == OperationType.DELETE).collect(Collectors.toList()), connection, config);
            // 需要写入的数据
            List writeRecords = records.stream().filter(item -> item.getType() == OperationType.WRITE).collect(Collectors.toList());
            // 强制使用了 upsert 或者自动模式 小于使用copy的量 使用copy执行
            boolean useUpsert = writeMode == WriteMode.AUTO && writeRecords.size() < config.getTableConfig(tableName).getCopyWriteBatchSize();
            if (writeMode == WriteMode.UPSERT || useUpsert) {
                if (config.getTableConfig(tableName).isEnableHstoreUpsertAutocommit()) {
                    JdbcUtil.executeSql(connection, "set enable_hstore_upsert_autocommit=on;");
                }
                // 执行upsert
                try {
                    upsertExecutor.execute(writeRecords, connection, config);
                } finally {
                    if (config.getTableConfig(tableName).isEnableHstoreUpsertAutocommit()) {
                        JdbcUtil.executeSql(connection, "set enable_hstore_upsert_autocommit=off;");
                    }
                }
            } else {
                switch (writeMode) {
                    case COPY_UPSERT:
                    case AUTO:
                        copyUpsertExecutor.execute(writeRecords, connection, config);
                        break;
                    case COPY_MERGE:
                        copyMergeExecutor.execute(writeRecords, connection, config);
                        break;
                    case UPDATE:
                    case COPY_UPDATE:
                    case UPDATE_AUTO:
                        copyUpdateExecutor.execute(writeRecords, connection, config);
                        break;
                    default:
                        log.error("unknown executor write mode {}", writeMode);
                }
            }
            // 提交事务 并完成事件
            connection.commit();
            action.getFuture().complete(null);
        } catch (Throwable e) {
            log.error("hand put action fail.", e);
            try {
                connection.rollback();
            } catch (SQLException sqlE) {
                log.error("rollback fail.", sqlE);
                throw DwsClientException.fromException(sqlE);
            }
            throw DwsClientException.fromException(e);
        }
    }
}