001/*
002 *  Copyright (c) 2022-2025, 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.datasource.processor.DataSourceProcessor;
019import com.mybatisflex.core.exception.FlexExceptions;
020import com.mybatisflex.core.util.ClassUtil;
021import org.apache.ibatis.logging.LogFactory;
022
023import javax.sql.DataSource;
024import java.lang.reflect.Method;
025
026/**
027 * @author michael
028 */
029public class DataSourceManager {
030
031    private static DataSourceDecipher decipher;
032
033    public static DataSourceDecipher getDecipher() {
034        return decipher;
035    }
036
037    public static void setDecipher(DataSourceDecipher decipher) {
038        DataSourceManager.decipher = decipher;
039    }
040
041    /**
042     * 动态数据源key取值处理
043     */
044    private static DataSourceProcessor dataSourceProcessor;
045
046    public static DataSourceProcessor getDataSourceProcessor() {
047        return dataSourceProcessor;
048    }
049
050    public static void setDataSourceProcessor(DataSourceProcessor dataSourceProcessor) {
051        DataSourceManager.dataSourceProcessor = dataSourceProcessor;
052    }
053
054    private static DataSourceShardingStrategy dataSourceShardingStrategy;
055
056    public static DataSourceShardingStrategy getDataSourceShardingStrategy() {
057        return dataSourceShardingStrategy;
058    }
059
060    public static void setDataSourceShardingStrategy(DataSourceShardingStrategy dataSourceShardingStrategy) {
061        DataSourceManager.dataSourceShardingStrategy = dataSourceShardingStrategy;
062    }
063
064    public static void decryptDataSource(DataSource dataSource) {
065        if (decipher == null) {
066            return;
067        }
068
069        try {
070            restartDataSource(dataSource);
071        } catch (Exception ignored) {
072            // do nothing here.
073        }
074
075        for (DataSourceProperty property : DataSourceProperty.values()) {
076            Method getterMethod = ClassUtil.getAnyMethod(dataSource.getClass(), property.getGetterMethods());
077            if (getterMethod != null) {
078                String value = invokeMethod(getterMethod, dataSource);
079                if (value != null) {
080                    value = decipher.decrypt(property, value);
081                    Method setter = ClassUtil.getAnyMethod(dataSource.getClass(), property.getSetterMethods());
082                    if (setter != null && value != null) {
083                        invokeMethod(setter, dataSource, value);
084                    }
085                }
086            }
087        }
088    }
089
090    static void restartDataSource(DataSource dataSource) {
091        Method restartMethod = ClassUtil.getFirstMethod(ClassUtil.getUsefulClass(dataSource.getClass())
092            , method -> "restart".equals(method.getName()) && method.getParameterCount() == 0);
093        if (restartMethod != null) {
094            try {
095                restartMethod.invoke(dataSource);
096            } catch (Exception e) {
097                throw FlexExceptions.wrap(e);
098            }
099        }
100    }
101
102
103    static String invokeMethod(Method method, Object object, Object... params) {
104        try {
105            return (String) method.invoke(object, params);
106        } catch (Exception e) {
107            LogFactory.getLog(DataSourceManager.class).error("Can not invoke method: " + method.getName(), e);
108        }
109        return null;
110    }
111
112    static String processDataSourceKey(String dataSourceKey, Object targetOrProxy, Method method, Object[] arguments) {
113        // 如果没有配置 DataSourceProcessor 实例,则不做处理,返回原始值
114        return dataSourceProcessor == null ? dataSourceKey : dataSourceProcessor.process(dataSourceKey, targetOrProxy, method, arguments);
115    }
116
117
118    static String getShardingDsKey(String dataSource, Object mapper, Method method, Object[] args) {
119        return dataSourceShardingStrategy != null ? dataSourceShardingStrategy.doSharding(dataSource, mapper, method, args) : null;
120    }
121}