com.feingto.cloud.config.datasource.dynamic.support.DataSourceContext Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of feingto-auto-configurations Show documentation
The newest version!
package com.feingto.cloud.config.datasource.dynamic.support;

import com.feingto.cloud.kit.AtomicIntegerRoundRobin;
import lombok.extern.slf4j.Slf4j;

import javax.sql.DataSource;
import java.util.LinkedHashMap;
import java.util.Map;
import java.util.Optional;
import java.util.stream.Collectors;

/**
 * 数据源切换上下文
 *
 * @author longfei
 */
@Slf4j
public abstract class DataSourceContext {
    public static final String PRIMARY = "dataSource";
    public static final String SLAVE_PREFIX = "slave-";
    private static final ThreadLocal LOCAL_REQUEST = new ThreadLocal<>();
    private static final Map DATA_SOURCES = new LinkedHashMap<>();
    private static final AtomicIntegerRoundRobin ROUND_ROBIN = new AtomicIntegerRoundRobin();

    public static DataSource getDataSource(String key) {
        return DATA_SOURCES.containsKey(key)
                ? DATA_SOURCES.get(key)
                : DATA_SOURCES.get(DataSourceContext.PRIMARY);
    }

    public static void addDataSource(String key, DataSource dataSource) {
        DATA_SOURCES.put(key, dataSource);
    }

    public static DataSource getDefaultDataSource() {
        return DATA_SOURCES.get(PRIMARY);
    }

    public static Map getDataSources() {
        return DATA_SOURCES;
    }

    public static void removeDataSource(String key) {
        DATA_SOURCES.remove(key);
    }

    /**
     * 切换至只读数据源（前缀：slave-）
     */
    public static void setReadOnlyDataSource() {
        Optional.of(DATA_SOURCES.keySet().stream()
                .filter(datasource -> !PRIMARY.equals(datasource))
                .filter(datasource -> datasource.startsWith(SLAVE_PREFIX))
                .collect(Collectors.toList()))
                .filter(slaves -> slaves.size() > 0)
                // 轮询
                .ifPresent(slaves -> setDbType(slaves.get(ROUND_ROBIN
                        .totalIndex(slaves.size())
                        .getAndIncrement())));
    }

    public static boolean contains(String dataSource) {
        return DATA_SOURCES.containsKey(dataSource);
    }

    public static String getDbType() {
        return LOCAL_REQUEST.get() != null ? LOCAL_REQUEST.get() : PRIMARY;
    }

    public static void setDbType(String dataSource) {
        log.trace("Switch to data source {}", dataSource);
        LOCAL_REQUEST.set(dataSource);
    }

    public static void removeDbType() {
        LOCAL_REQUEST.remove();
    }
}