com.huaweicloud.dws.client.util.CacheUtil Maven / Gradle / Ivy

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

import com.github.benmanes.caffeine.cache.Caffeine;
import com.github.benmanes.caffeine.cache.LoadingCache;
import com.github.benmanes.caffeine.cache.RemovalListener;
import com.huaweicloud.dws.client.DwsClient;
import com.huaweicloud.dws.client.DwsConfig;
import com.huaweicloud.dws.client.exception.DwsClientException;
import com.huaweicloud.dws.client.model.TableName;
import com.huaweicloud.dws.client.model.TableSchema;
import com.huaweicloud.dws.client.model.TypeDefinition;

import lombok.extern.slf4j.Slf4j;

import java.sql.SQLException;
import java.util.concurrent.TimeUnit;

/**
 * @ProjectName: dws-connector
 * @ClassName: CacheUtil
 * @Description:
 * @Date: 2023/1/28 10:17
 * @Version: 1.0
 */
@Slf4j
public class CacheUtil {
    /**
     * 使用缓存组件对表结构做缓存，避免大量重复查询提升性能
     */
    private final LoadingCache TABLE_CACHE;

    private final LoadingCache typeCache;

    private final DwsClient client;

    private static final int MAX_CACHE_SIZE = 6000;

    public CacheUtil(DwsConfig config, DwsClient client) {
        TABLE_CACHE = initCache(config);
        typeCache = initTypeCache(config);
        this.client = client;
    }

    private LoadingCache initCache(DwsConfig config) {
        return Caffeine.newBuilder()
            .maximumSize(MAX_CACHE_SIZE)
            .expireAfterWrite(config.getMetadataCacheSeconds(), TimeUnit.SECONDS)
            .removalListener((RemovalListener) (key, value, cause) -> log.debug("table schema cache expire. key {}", key))
            .recordStats()
            .build(tableName -> {
                try {
                    log.debug("load table schema cache tableName = {}", tableName);
                    return getTableSchema(tableName, true);
                } catch (Exception e) {
                    log.error("load table schema error", e);
                    throw DwsClientException.fromException(e);
                }
            });

    }

    private LoadingCache initTypeCache(DwsConfig config) {
        return Caffeine.newBuilder()
            .maximumSize(MAX_CACHE_SIZE)
            .expireAfterWrite(config.getMetadataCacheSeconds(), TimeUnit.SECONDS)
            .removalListener((RemovalListener) (key, value, cause) -> log.debug("type schema cache expire. key {}", key))
            .recordStats()
            .build(typeName -> {
                try {
                    log.debug("load type schema cache scheme = {}", typeName);
                    return getTypeDefinition(typeName, true);
                } catch (Exception e) {
                    log.error("load type schema error", e);
                    throw DwsClientException.fromException(e);
                }
            });

    }

    public TableSchema getTableSchema(TableName tableName) throws DwsClientException {
        return getTableSchema(tableName, false);
    }

    public TableSchema getTableSchema(TableName tableName, boolean noCache) throws DwsClientException {
        if (!noCache) {
            try {
                return TABLE_CACHE.get(tableName);
            } catch (Exception e) {
                throw DwsClientException.fromException(e);
            }
        }
        return client.sql(connection -> {
            try {
                return JdbcUtil.getTableSchema(connection, tableName,
                    client.getConfig().getTableConfig(tableName).getUniqueKeys(),
                    client.getConfig().getTableConfig(tableName).isCaseSensitive());
            } catch (DwsClientException e) {
                if (e.getOriginal() instanceof SQLException) {
                    throw (SQLException) e.getOriginal();
                }
                throw new SQLException(e);
            }
        });
    }

    public TypeDefinition getTypeDefinition(TableName typeName) throws DwsClientException {
        return getTypeDefinition(typeName, false);
    }

    public TypeDefinition getTypeDefinition(TableName typeName, boolean noCache) throws DwsClientException {
        if (!noCache) {
            try {
                return typeCache.get(typeName);
            } catch (Exception e) {
                throw DwsClientException.fromException(e);
            }
        }
        return client.sql(connection -> {
            try {
                return JdbcUtil.getTypeDefinition(connection, typeName);
            } catch (DwsClientException e) {
                if (e.getOriginal() instanceof SQLException) {
                    throw (SQLException) e.getOriginal();
                }
                throw new SQLException(e);
            }
        });
    }

    public void removeCache(TableName tableName) {
        TABLE_CACHE.invalidate(tableName);
    }
}