001/*
002 *  Copyright (c) 2022-2024, Mybatis-Flex (fuhai999@gmail.com).
003 *  <p>
004 *  Licensed under the Apache License, Version 2.0 (the "License");
005 *  you may not use this file except in compliance with the License.
006 *  You may obtain a copy of the License at
007 *  <p>
008 *  http://www.apache.org/licenses/LICENSE-2.0
009 *  <p>
010 *  Unless required by applicable law or agreed to in writing, software
011 *  distributed under the License is distributed on an "AS IS" BASIS,
012 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
013 *  See the License for the specific language governing permissions and
014 *  limitations under the License.
015 */
016package com.mybatisflex.core.datasource;
017
018import com.mybatisflex.core.exception.FlexAssert;
019
020import java.lang.reflect.Method;
021import java.util.ArrayDeque;
022import java.util.Deque;
023import java.util.function.Supplier;
024
025/**
026 * @author michael
027 */
028public class DataSourceKey {
029
030    private static ThreadLocal<Deque<String>> lookup = ThreadLocal.withInitial(ArrayDeque::new);
031
032    private DataSourceKey() {
033    }
034
035    public static void use(String dataSourceKey) {
036        Deque<String> deque = lookup.get();
037        if (deque == null) {
038            deque = new ArrayDeque<>(1);
039            lookup.set(deque);
040        }
041        deque.push(dataSourceKey);
042    }
043
044    public static String get() {
045        Deque<String> deque = lookup.get();
046        return deque != null ? deque.peek() : null;
047    }
048
049    public static void clear() {
050        Deque<String> deque = lookup.get();
051        if (deque != null) {
052            deque.pop();
053            if (deque.isEmpty()) {
054                lookup.remove();
055            }
056        }
057    }
058
059    public static void forceClear() {
060        lookup.remove();
061    }
062
063    public static void use(String dataSourceKey, Runnable runnable) {
064        try {
065            use(dataSourceKey);
066            runnable.run();
067        } finally {
068            clear();
069        }
070    }
071
072    public static <T> T use(String dataSourceKey, Supplier<T> supplier) {
073        try {
074            use(dataSourceKey);
075            return supplier.get();
076        } finally {
077            clear();
078        }
079    }
080
081    public static void setThreadLocal(ThreadLocal<Deque<String>> threadLocal) {
082        FlexAssert.notNull(threadLocal, "threadLocal");
083        if (threadLocal.get() == null) {
084            threadLocal.set(lookup.get());
085        }
086        lookup = threadLocal;
087    }
088
089    public static String processDataSourceKey(String dataSourceKey, Object targetOrProxy, Method method, Object[] arguments) {
090        String dsKey = DataSourceManager.processDataSourceKey(dataSourceKey, targetOrProxy, method, arguments);
091        return dsKey != null ? dsKey : dataSourceKey;
092    }
093
094
095    public static String getShardingDsKey(String dataSource, Object mapper, Method method, Object[] args) {
096        String shardingDsKey = DataSourceManager.getShardingDsKey(dataSource, mapper, method, args);
097        return shardingDsKey != null ? shardingDsKey : dataSource;
098    }
099
100}