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

Go to download
package com.huaweicloud.dws.client.handler;

import com.huaweicloud.dws.client.DwsConfig;
import com.huaweicloud.dws.client.TableConfig;
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.ColumnKey;
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.util.LocalUtil;
import com.huaweicloud.dws.client.worker.ConnectionProvider;
import lombok.extern.slf4j.Slf4j;

import java.sql.Connection;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
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;
    final CopyUpsertExecutor copyUpsertExecutor;
    final CopyUpdateExecutor copyUpdateExecutor;

    public PutActionHandler(DwsConfig config, ConnectionProvider connectionProvider) {
        this.config = config;
        this.connectionProvider = connectionProvider;
        copyMergeExecutor = new CopyMergeExecutor(connectionProvider);
        copyUpsertExecutor = new CopyUpsertExecutor(connectionProvider);
        copyUpdateExecutor = new CopyUpdateExecutor(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();
        TableConfig tableConfig = config.getTableConfig(tableName);
        Connection connection = null;
        try {
            connection = connectionProvider.getOrInitConnection();
        } catch (Exception e) {
            log.error("get connection fail", e);
            throw DwsClientException.fromException(e);
        }
        try {
            LocalUtil.setTableConfig(config.getTableConfig(tableName));
            if (!tableConfig.isAutoCommit() && connection.getAutoCommit()) {
                connection.setAutoCommit(false);
            }
            // 先删除
            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());
            Map> cache = new HashMap<>();
            for (Record record : writeRecords) {
                cache.computeIfAbsent(new ColumnKey(record.getColumnBit(), record.getIgnoreUpdate()),
                        columnKey -> new ArrayList<>(writeRecords.size())).add(record);
            }
            for (Map.Entry> entry : cache.entrySet()) {
                writeRecord(tableName, writeMode, connection, entry.getValue());
            }
            // 提交事务 并完成事件
            if (!tableConfig.isAutoCommit()) {
                connection.commit();
            }
            action.getFuture().complete(null);
            log.info("action successful.");
        } catch (Throwable e) {
            log.error("hand put action fail.", e);
            // 异常后清掉临时表缓存，避免事务回滚丢失临时表
            connectionProvider.cleanTempTable(records.get(0).getTableSchema());
            if (!tableConfig.isAutoCommit()) {
                try {
                    connection.rollback();
                } catch (SQLException sqlE) {
                    log.error("rollback fail.", sqlE);
                    throw DwsClientException.fromException(sqlE);
                }
            }
            throw DwsClientException.fromException(e);
        } finally {
            LocalUtil.removeTableConfig();
        }
    }

    private void writeRecord(TableName tableName, WriteMode writeMode, Connection connection, List writeRecords) throws SQLException, DwsClientException {
        if (writeRecords.isEmpty()) {
            log.info("no data write, table = {}", tableName.getFullName());
            return;
        }
        // 强制使用了 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);
            }
        }
    }
}