001/*
002 *  Copyright (c) 2022-2023, 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.FlexExceptions;
019import com.mybatisflex.core.util.ClassUtil;
020import org.apache.ibatis.logging.LogFactory;
021
022import javax.sql.DataSource;
023import java.lang.reflect.Method;
024
025/**
026 * @author michael
027 */
028public class DataSourceManager {
029
030    private static DataSourceDecipher decipher;
031
032    public static DataSourceDecipher getDecipher() {
033        return decipher;
034    }
035
036    public static void setDecipher(DataSourceDecipher decipher) {
037        DataSourceManager.decipher = decipher;
038    }
039
040    private static DataSourceShardingStrategy dataSourceShardingStrategy;
041
042    public static DataSourceShardingStrategy getDataSourceShardingStrategy() {
043        return dataSourceShardingStrategy;
044    }
045
046    public static void setDataSourceShardingStrategy(DataSourceShardingStrategy dataSourceShardingStrategy) {
047        DataSourceManager.dataSourceShardingStrategy = dataSourceShardingStrategy;
048    }
049
050    public static void decryptDataSource(DataSource dataSource) {
051        if (decipher == null) {
052            return;
053        }
054
055        try {
056            restartDataSource(dataSource);
057        } catch (Exception ignored) {
058            // do nothing here.
059        }
060
061        for (DataSourceProperty property : DataSourceProperty.values()) {
062            Method getterMethod = ClassUtil.getAnyMethod(dataSource.getClass(), property.getGetterMethods());
063            if (getterMethod != null) {
064                String value = invokeMethod(getterMethod, dataSource);
065                if (value != null) {
066                    value = decipher.decrypt(property, value);
067                    Method setter = ClassUtil.getAnyMethod(dataSource.getClass(), property.getSetterMethods());
068                    if (setter != null && value != null) {
069                        invokeMethod(setter, dataSource, value);
070                    }
071                }
072            }
073        }
074    }
075
076    static void restartDataSource(DataSource dataSource) {
077        Method restartMethod = ClassUtil.getFirstMethod(ClassUtil.getUsefulClass(dataSource.getClass())
078            , method -> "restart".equals(method.getName()) && method.getParameterCount() == 0);
079        if (restartMethod != null) {
080            try {
081                restartMethod.invoke(dataSource);
082            } catch (Exception e) {
083                throw FlexExceptions.wrap(e);
084            }
085        }
086    }
087
088
089    static String invokeMethod(Method method, Object object, Object... params) {
090        try {
091            return (String) method.invoke(object, params);
092        } catch (Exception e) {
093            LogFactory.getLog(DataSourceManager.class).error("Can not invoke method: " + method.getName(), e);
094        }
095        return null;
096    }
097
098
099    static String getByShardingStrategy(String dataSource, Object mapper, Method method, Object[] args) {
100        return dataSourceShardingStrategy != null ? dataSourceShardingStrategy.doSharding(dataSource, mapper, method, args) : null;
101    }
102}