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        lookup.get().push(dataSourceKey);
037    }
038
039    public static String get() {
040        return lookup.get().peek();
041    }
042
043    public static void clear() {
044        Deque<String> deque = lookup.get();
045        deque.pop();
046        if (deque.isEmpty()) {
047            lookup.remove();
048        }
049    }
050
051    public static void forceClear() {
052        lookup.remove();
053    }
054
055    public static void use(String dataSourceKey, Runnable runnable) {
056        try {
057            use(dataSourceKey);
058            runnable.run();
059        } finally {
060            clear();
061        }
062    }
063
064    public static <T> T use(String dataSourceKey, Supplier<T> supplier) {
065        try {
066            use(dataSourceKey);
067            return supplier.get();
068        } finally {
069            clear();
070        }
071    }
072
073    public static void setThreadLocal(ThreadLocal<Deque<String>> threadLocal) {
074        FlexAssert.notNull(threadLocal, "threadLocal");
075        if (threadLocal.get() == null) {
076            threadLocal.set(lookup.get());
077        }
078        lookup = threadLocal;
079    }
080
081    public static String getShardingDsKey(String dataSource, Object mapper, Method method, Object[] args) {
082        String shardingDsKey = DataSourceManager.getShardingDsKey(dataSource, mapper, method, args);
083        return shardingDsKey != null ? shardingDsKey : dataSource;
084    }
085
086    // === For Removal ===
087
088    @Deprecated
089    public static String getByManual() {
090        throw new UnsupportedOperationException("使用 DataSource.get() 代替。");
091    }
092    @Deprecated
093    public static String getByAnnotation() {
094        throw new UnsupportedOperationException("使用 DataSource.get() 代替。");
095    }
096    @Deprecated
097    public static void useWithAnnotation(String dataSourceKey) {
098        throw new UnsupportedOperationException("使用 DataSource.use(String) 代替。");
099    }
100    @Deprecated
101    public static void setAnnotationKeyThreadLocal(ThreadLocal<String> annotationKeyThreadLocal) {
102        throw new UnsupportedOperationException("使用 DataSource.setThreadLocal(ThreadLocal<Deque<String>>) 代替。");
103    }
104    @Deprecated
105    public static void setManualKeyThreadLocal(ThreadLocal<String> manualKeyThreadLocal) {
106        throw new UnsupportedOperationException("使用 DataSource.setThreadLocal(ThreadLocal<Deque<String>>) 代替。");
107    }
108
109}