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