com.huaweicloud.dws.client.DwsClient Maven / Gradle / Ivy

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

import com.huaweicloud.dws.client.action.ScanAction;
import com.huaweicloud.dws.client.action.SqlAction;
import com.huaweicloud.dws.client.collector.ActionCollector;
import com.huaweicloud.dws.client.collector.ITableCollector;
import com.huaweicloud.dws.client.exception.DwsClientException;
import com.huaweicloud.dws.client.exception.ExceptionCode;
import com.huaweicloud.dws.client.exception.InvalidException;
import com.huaweicloud.dws.client.function.SqlExceptionFunction;
import com.huaweicloud.dws.client.model.Record;
import com.huaweicloud.dws.client.model.TableName;
import com.huaweicloud.dws.client.model.TableSchema;
import com.huaweicloud.dws.client.model.TypeDefinition;
import com.huaweicloud.dws.client.op.Delete;
import com.huaweicloud.dws.client.op.Get;
import com.huaweicloud.dws.client.op.Operate;
import com.huaweicloud.dws.client.op.Scan;
import com.huaweicloud.dws.client.op.Write;
import com.huaweicloud.dws.client.util.AssertUtil;
import com.huaweicloud.dws.client.util.CacheUtil;
import com.huaweicloud.dws.client.util.JdbcUtil;
import com.huaweicloud.dws.client.util.LogUtil;
import com.huaweicloud.dws.client.worker.ExecutionPool;

import lombok.extern.slf4j.Slf4j;

import java.io.Closeable;
import java.io.IOException;
import java.sql.Connection;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.stream.Collectors;

/**
 * @ProjectName: dws-connector
 * @ClassName: DwsClient
 * @Description: 提供通用、便捷、高性能的入库能力
 * @Date: 2022/12/22 11:15
 * @Version: 1.0
 */
@Slf4j
public class DwsClient implements Closeable {

    /**
     * 事件收集器
     */
    private final ActionCollector collector;

    /**
     * 资源池
     */
    private final ExecutionPool pool;

    /**
     * 缓存工具
     */
    private final CacheUtil cacheUtil;

    private final DwsConfig config;

    /**
     * 初始化client
     *
     * @param config 配置参数
     * @throws InvalidException 配置无效时抛出参数无效异常
     */
    public DwsClient(DwsConfig config, String name) {
        AssertUtil.nonNull(config, new InvalidException("config is null"));
        this.pool = new ExecutionPool(name, config, this);
        collector = this.pool.getCollector();
        this.config = config;
        this.cacheUtil = new CacheUtil(config, this);
    }

    public DwsClient(DwsConfig config) {
        this(config, "dwsClient");
    }

    public DwsConfig getConfig() {
        return config;
    }

    /**
     * 原生sql执行
     */
    public  T sql(SqlExceptionFunction func) throws DwsClientException {
        SqlAction action = new SqlAction<>(func, config);
        while (!pool.submit(action)) {
            LogUtil.withLogSwitch(config, () -> log.warn("try submit."));
        }
        return action.getResult();
    }

    /**
     * 获取表结构定义
     */
    public TableSchema getTableSchema(TableName tableName) throws DwsClientException {
        if (tableName == null) {
            throw new InvalidException("tableName is null");
        }
        return cacheUtil.getTableSchema(tableName);
    }

    /**
     * 获取数据类型定义
     */
    public TypeDefinition getTypeDefinition(TableName typeName) throws DwsClientException {
        if (typeName == null) {
            throw new InvalidException("typeName is null");
        }
        return cacheUtil.getTypeDefinition(typeName);
    }

    /**
     * 根据表名开启一个写入数据API（如果非默认schema需要带表名全称 schema.name）
     *
     * @deprecated 写入方式增加了update模式，为避免歧义统一使用{@link DwsClient#write(java.lang.String)}
     */
    @Deprecated
    public Operate upsert(String tableName) throws DwsClientException {
        TableName name = TableName.valueOf(tableName);
        return write(name);
    }

    /**
     * @deprecated {@link DwsClient#write(TableName)}
     */
    @Deprecated
    public Operate upsert(TableName tableName) throws DwsClientException {
        TableSchema tableSchema = getTableSchema(tableName);
        return write(tableSchema);
    }

    /**
     * @deprecated {@link DwsClient#write(TableSchema)}
     */
    @Deprecated
    public Operate upsert(TableSchema schema) {
        return write(schema);
    }

    public Operate write(String tableName) throws DwsClientException {
        TableName name = TableName.valueOf(tableName);
        return write(name);
    }

    public Operate write(String tableName, Map schema) throws DwsClientException {
        TableName name = TableName.valueOf(tableName);
        return write(name, schema);
    }

    public Operate write(TableName tableName) throws DwsClientException {
        TableSchema tableSchema = getTableSchema(tableName);
        return write(tableSchema);
    }

    public Operate write(TableName tableName, Map schema) throws DwsClientException {
        TableSchema tableSchema = getTableSchema(tableName);
        tableSchema = checkSchema(tableSchema, schema);
        return write(tableSchema);
    }

    /**
     * 检查表结构是否包含用户指定字段
     */
    private TableSchema checkSchema(TableSchema tableSchema, Map schema) throws DwsClientException {
        if (schema == null) {
            return tableSchema;
        }
        Map columnMap = tableSchema.getColumnIndexMap();
        List addColumnList =
            schema.keySet().stream().filter(key -> !columnMap.containsKey(key)).collect(Collectors.toList());
        if (addColumnList.isEmpty()) {
            return tableSchema;
        }
        cacheUtil.removeCache(tableSchema.getTableName());
        log.info("dws table {}, not find columns {}", tableSchema.getTableName().getFullName(), addColumnList);
        for (String column : addColumnList) {
            String sql = JdbcUtil.getAddColumn(schema, column, tableSchema.getTableName().getFullName());
            try {
                sql((SqlExceptionFunction) connection -> {
                    connection.createStatement().execute(sql);
                    return null;
                });
            } catch (DwsClientException e) {
                // 如果报当前列已经存在的错误，直接忽略，并发时可能出现
                if (e.getCode() != ExceptionCode.ALREADY_EXISTS) {
                    throw e;
                }
            }
        }
        return getTableSchema(tableSchema.getTableName());
    }

    public Operate write(TableSchema schema) {
        return new Write(schema, this.collector.getTableActionCollector(schema));
    }

    /**
     * 获取当前表得收集器
     */
    public ITableCollector collector(TableSchema schema) {
        return collector.getTableActionCollector(schema);
    }

    /**
     * 根据表名开启一个删除数据API（如果非默认schema需要带表名全称 schema.name）
     */
    public Operate delete(String tableName) throws DwsClientException {
        TableName name = TableName.valueOf(tableName);
        return delete(name);
    }

    public Operate delete(TableName tableName) throws DwsClientException {
        TableSchema tableSchema = getTableSchema(tableName);
        return delete(tableSchema);
    }

    public Operate delete(TableSchema schema) {
        return new Delete(schema, this.collector(schema));
    }

    /**
     * 根据表名开启一个查询数据API（如果非默认schema需要带表名全称 schema.name）
     */
    public CompletableFuture get(TableName tableName, String keyColumn, Object keyValue)
        throws DwsClientException {
        TableSchema tableSchema = getTableSchema(tableName);
        return get(tableSchema, keyColumn, keyValue);
    }

    public CompletableFuture get(TableSchema schema, String keyColumn, Object keyValue)
        throws DwsClientException {
        Get get = Get.newBuilder(schema, this.collector).setPrimaryKey(keyColumn, keyValue).build();
        commit(schema, get);
        return get.getFuture();
    }

    public CompletableFuture get(TableSchema schema, String[] keyColumns, Object[] keyValues)
        throws DwsClientException {
        Get get = Get.newBuilder(schema, this.collector)
            .setPrimaryKeys(Arrays.asList(keyColumns), Arrays.asList(keyValues))
            .build();
        commit(schema, get);
        return get.getFuture();
    }

    public CompletableFuture get(Record record) throws DwsClientException {
        Get get = new Get(record, this.collector, true);
        commit(record.getTableSchema(), get);
        return get.getFuture();
    }

    /**
     * 根据表名开启一个scan查询数据API（如果非默认schema需要带表名全称 schema.name）
     */
    public CompletableFuture> scan(TableSchema schema, String[] keyColumns, Object[] keyValues)
        throws DwsClientException {
        return scan(schema, keyColumns, keyValues, schema.getColumnNames().toArray(new String[0]));
    }

    public CompletableFuture> scan(TableSchema schema, String[] keyColumns, Object[] keyValues,
        String[] selectedColumns) throws DwsClientException {
        Scan.Builder builder =
            Scan.newBuilder(schema).setLookUpKeys(Arrays.asList(keyColumns), Arrays.asList(keyValues));
        if (selectedColumns != null) {
            builder.withSelectedColumns(selectedColumns);
        }
        ScanAction action = new ScanAction(builder.build(), config);
        while (!pool.submit(action)) {
            LogUtil.withLogSwitch(config, () -> log.warn("try submit."));
        }
        return action.getFuture();
    }

    public CompletableFuture> scan(Scan scan) throws DwsClientException {
        ScanAction action = new ScanAction(scan, config);
        while (!pool.submit(action)) {
            LogUtil.withLogSwitch(config, () -> log.warn("try submit."));
        }
        return action.getFuture();
    }

    public void commit(TableSchema schema, Get get) throws DwsClientException {
        Get.check(get);
        if (get.isFullColumn()) {
            for (int i = 0; i < schema.getColumns().size(); ++i) {
                if (!get.getRecord().isSet(i)) {
                    get.getRecord().setValue(i, null);
                }
            }
        }
        try {
            get.commit();
        } catch (DwsClientException e) {
            get.getFuture().completeExceptionally(e);
        }
    }

    /**
     * 将缓存数据强制刷入数据库
     */
    public void flush() throws DwsClientException {
        pool.flush();
        pool.tryException();
    }

    /**
     * 关闭client
     * 1、关闭新数据写入能力
     * 2、刷缓存数据入库
     * 3、释放client开辟资源
     */
    @Override
    public void close() throws IOException {
        Exception exception = null;
        // 先将本地缓存刷库
        try {
            flush();
        } catch (Exception e) {
            log.error("dws client close error", e);
            exception = e;
        }
        // 关闭资源
        pool.close();
        AssertUtil.isNull(exception, new IOException(exception));
    }

    /**
     * 将对应表的schema信息从缓存中删掉
     *
     * @param tableName 表名
     */
    public void removeTableSchema(TableName tableName) {
        cacheUtil.removeCache(tableName);
    }
}