webcrawler/db/src/main/java/db/config/HibernateConfig.java

package db.config;

import com.alibaba.druid.pool.DruidDataSource;
import db.annotation.Model;
import db.annotation.TableInfo;
import db.model.DataSourceModel;
import org.apache.commons.lang3.ClassUtils;
import org.apache.commons.lang3.StringUtils;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.hibernate.SessionFactory;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.beans.factory.support.DefaultListableBeanFactory;
import org.springframework.core.type.filter.AnnotationTypeFilter;
import org.springframework.orm.hibernate5.HibernateTemplate;
import org.springframework.orm.hibernate5.HibernateTransactionManager;
import org.springframework.orm.hibernate5.LocalSessionFactoryBean;

import javax.annotation.PostConstruct;
import javax.annotation.Resource;
import javax.sql.DataSource;
import java.io.IOException;
import java.lang.annotation.Annotation;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Properties;

/**
 * Created by reborn on 2017/7/28.
 */
public class HibernateConfig {

    private static Logger log = LogManager.getLogger();

    @Value("${hibernate.dialect}")
    private String dialect;
    @Value("${hibernate.connection.driver_class}")
    private String driver;

    @Value("${db_url}")
    private String url;
    @Value("${db_username}")
    private String username;
    @Value("${db_password}")
    private String password;

    @Resource
    private DefaultListableBeanFactory beanFactory;

    private static final Class<DataSourceModel> rootClass = DataSourceModel.class;

    private static final String basePackage = "db.model";

    private static HibernateTemplate rootHibernateTemplate;

    private static final Map<Class<? extends Annotation>, HibernateTemplate> map = new HashMap<>();

    public final static HibernateTemplate getRootHibernateTemplate() {
        return rootHibernateTemplate;
    }

    public static final HibernateTemplate get(Class<? extends Annotation> key) {
        if(map.containsKey(key)){
            return map.get(key);
        }else{
            throw new RuntimeException("找不到key="+key+"的HibernateTemplate对象");
        }
    }

    public static final Map<Class<? extends Annotation>, HibernateTemplate> get(){
        return map;
    }

    private static final Map<Class,TableInfo> tableNotes=new HashMap<>();

    public static TableInfo getTableNote(Class c){
        return tableNotes.get(c);
    }

    /**
     * 注册数据源
     * @param url
     * @param username
     * @param password
     * @return
     */
    private DruidDataSource dataSource(String url, String username, String password) {
        DruidDataSource dataSource = new DruidDataSource();
        dataSource.setUrl(url);
        dataSource.setUsername(username);
        dataSource.setPassword(password);
        dataSource.setDriverClassName(driver);
        log.info("数据源初始化\nurl=" + url + "\nusername=" + username + "\npassword=" + password);
        return dataSource;
    }

    private DruidDataSource dataSource(DataSourceModel dataSourceModel) {
        String url = "jdbc:mysql://" + dataSourceModel.getHostname() + ":" + dataSourceModel.getDbPort() + "/" + dataSourceModel.getDbName() + "?serverTimezone=UTC&useSSL=false";
        return dataSource(url, dataSourceModel.getUsername(), dataSourceModel.getDbPassword());
    }

    /**
     * 初始化根数据源
     * @return
     */
    private HibernateTemplate initRootDB() {
        SessionFactory rootdb_sessionFactory = initSessionFactory(dataSource(url, username, password), (LocalSessionFactoryBean bean) -> {
            bean.setAnnotatedClasses(rootClass);
        });
        return new HibernateTemplate(rootdb_sessionFactory);
    }

    /**
     * 初始化数据源
     * @param dataSourceModel
     * @return
     */
    private HibernateTemplate initDB(DataSourceModel dataSourceModel) {
        DruidDataSource dataSource=dataSource(dataSourceModel);
        SessionFactory sessionFactory = initSessionFactory(dataSource, (LocalSessionFactoryBean bean) -> {
            Class c;
            try {
                //按照规则拼接实体类注解标记全类名
                c = Class.forName( dataSourceModel.getAnnotation());
                if(!Annotation.class.isAssignableFrom(c)){
                    log.error(c+"不是注解标记！！！");
                    return;
                }
            } catch (ClassNotFoundException e) {
                log.error(e);
                return;
            }

            //按照注解标记从指定包扫描实体类，并且注册到session工厂
            bean.setPackagesToScan(basePackage);
            bean.setEntityTypeFilters(new AnnotationTypeFilter(c));
        },dataSourceModel.getDbName());

        if (sessionFactory != null) {
            HibernateTransactionManager hibernateTransactionManager = new HibernateTransactionManager();
            hibernateTransactionManager.setSessionFactory(sessionFactory);
            hibernateTransactionManager.setDataSource(dataSource);
            hibernateTransactionManager.afterPropertiesSet();
            //添加事务管理器
            beanFactory.registerSingleton(DBBeanNameManager.getName(dataSourceModel.getDbId(),DBBeanNameManager.transactionManager),hibernateTransactionManager);
            return new HibernateTemplate(sessionFactory);
        } else {
            return null;
        }
    }

    /**
     * bean命名规则
     */
    public enum DBBeanNameManager {
        hibernateTemplate("H"),
        transactionManager("T");

        private String prefix;

        DBBeanNameManager(String prefix) {
            this.prefix = prefix;
        }

        public static String getName(String key, DBBeanNameManager type){
            return type.prefix+key;
        }

        public static String getName(DBBeanNameManager type){
            return type.prefix+"root";
        }
    }

    /**
     * 初始化数据源，session工厂
     */
    @PostConstruct
    private void init() {

        //连接根数据源
        rootHibernateTemplate = initRootDB();
        beanFactory.registerSingleton(DBBeanNameManager.getName(DBBeanNameManager.hibernateTemplate),rootHibernateTemplate);
        //从根数据源读取其他数据源信息，并初始化
        List<DataSourceModel> dataSourceModels = rootHibernateTemplate.loadAll(rootClass);
        for (DataSourceModel dataSourceModel : dataSourceModels) {
            HibernateTemplate hibernateTemplate = initDB(dataSourceModel);
            if (hibernateTemplate != null) {

                beanFactory.registerSingleton(DBBeanNameManager.getName(dataSourceModel.getDbId(),DBBeanNameManager.hibernateTemplate),hibernateTemplate);

                Class c;
                try {
                    c = ClassUtils.getClass(dataSourceModel.getDbId());
                } catch (ClassNotFoundException e) {
                    throw new RuntimeException("类"+dataSourceModel.getDbId()+"不存在");
                }
                if(c.isAnnotationPresent(Model.class)){
                    Model model= (Model) c.getAnnotation(Model.class);
                    map.put(model.value(), hibernateTemplate);
                }else{
                    throw new RuntimeException(dataSourceModel.getDbId()+"缺少"+Model.class+"注解");
                }
            }
        }

    }

    /**
     * 初始化session工厂
     * @param dataSource
     * @param config
     * @return
     */
    private SessionFactory initSessionFactory(DataSource dataSource, SessionFactoryConfig config,String defaultSchema) {
        LocalSessionFactoryBean bean = new LocalSessionFactoryBean();
        Properties properties = new Properties();
        properties.setProperty("hibernate.dialect", dialect);
        properties.setProperty("hibernate.format_sql", "true");
        if(StringUtils.isNotEmpty(defaultSchema)) {
            properties.setProperty("hibernate.default_schema", defaultSchema);
        }
        bean.setHibernateProperties(properties);
        bean.setDataSource(dataSource);
        try {
            config.call(bean);
            bean.afterPropertiesSet();
            return bean.getObject();
        }  catch (IOException e) {
            log.error(e);
        }
        return null;
    }

    private SessionFactory initSessionFactory(DataSource dataSource, SessionFactoryConfig config) {
        return initSessionFactory(dataSource, config,null);
    }

}