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

package db.config;

import com.alibaba.druid.pool.DruidDataSource;
import db.AbstractModel;
import db.model.DataSourceModel;
import db.util.ClassScaner;
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.orm.hibernate5.HibernateTemplate;
import org.springframework.orm.hibernate5.HibernateTransactionManager;
import org.springframework.orm.hibernate5.LocalSessionFactoryBean;
import org.springframework.util.ObjectUtils;
import org.springframework.util.ReflectionUtils;

import javax.annotation.PostConstruct;
import javax.annotation.Resource;
import javax.sql.DataSource;
import java.io.IOException;
import java.lang.annotation.Annotation;
import java.lang.reflect.Method;
import java.util.*;

/**
 * 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 String rootKey = "ROOT";

    private static final String tableNote = "tableNote";

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

    private static final String annotation = "db.annotation.";

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

    private static final Map<String, HibernateTemplate> map = new HashMap<>();

    public static final HibernateTemplate get(String key) {
        return map.get(key);
    }

    public static final Map<String, HibernateTemplate> get(){
        return map;
    }

    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());
    }

    private HibernateTemplate initRootDB() {
        SessionFactory rootdb_sessionFactory = initSessionFactory(dataSource(url, username, password), (LocalSessionFactoryBean bean) -> {
            bean.setAnnotatedClasses(rootClass);
        });
        return new HibernateTemplate(rootdb_sessionFactory);
    }

    private HibernateTemplate initDB(DataSourceModel dataSourceModel) {
        SessionFactory sessionFactory = initSessionFactory(dataSource(dataSourceModel), (LocalSessionFactoryBean bean) -> {
            Set<Class> classes;

            Class c;
            try {
                c = Class.forName(annotation + dataSourceModel.getAnnotation());
            } catch (ClassNotFoundException e) {
                return;
            }

            if (Annotation.class.isAssignableFrom(c)) {
                classes = ClassScaner.scan(basePackage, (Class<? extends Annotation>) c);
                bean.setAnnotatedClasses(classes.toArray(new Class[classes.size()]));
            } else {
                log.error(c+"不是注解！！！");
            }
        });

        if (sessionFactory != null) {
            HibernateTransactionManager hibernateTransactionManager = new HibernateTransactionManager();
            hibernateTransactionManager.setSessionFactory(sessionFactory);
            beanFactory.registerSingleton(DBBeanNameManager.getName(dataSourceModel.getDbDesc(),DBBeanNameManager.transactionManager),hibernateTransactionManager);
            return new HibernateTemplate(sessionFactory);
        } else {
            return null;
        }
    }

    public enum DBBeanNameManager {
        hibernateTemplate,
        transactionManager;

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

        public static String getName(DBBeanNameManager type){
            return type.name()+"_root";
        }
    }

    @PostConstruct
    private void init() {

        HibernateTemplate root = initRootDB();
        beanFactory.registerSingleton(DBBeanNameManager.getName(DBBeanNameManager.hibernateTemplate),root);
        map.put(rootKey, root);

        List<DataSourceModel> dataSourceModels = root.loadAll(rootClass);
        for (DataSourceModel dataSourceModel : dataSourceModels) {
            HibernateTemplate hibernateTemplate = initDB(dataSourceModel);
            if (hibernateTemplate != null) {
                beanFactory.registerSingleton(DBBeanNameManager.getName(dataSourceModel.getDbDesc(),DBBeanNameManager.hibernateTemplate),hibernateTemplate);
                map.put(dataSourceModel.getDbDesc(), hibernateTemplate);
            }
        }

    }


    private SessionFactory initSessionFactory(DataSource dataSource, AddConfig config) {
        LocalSessionFactoryBean bean = new LocalSessionFactoryBean();
        Properties properties = new Properties();
        properties.setProperty("hibernate.dialect", dialect);
        properties.setProperty("hibernate.format_sql", "true");
        bean.setHibernateProperties(properties);
        bean.setDataSource(dataSource);
        try {
            config.addConfig(bean);
            bean.afterPropertiesSet();
            checkClass(bean);
            return bean.getObject();
        }  catch (IOException e) {
            log.error(e);
        }
        return null;
    }

    private void checkClass(LocalSessionFactoryBean bean) {
        for (Class c : bean.getMetadataSources().getAnnotatedClasses()) {
            if (AbstractModel.class.isAssignableFrom(c)) {
                try {
                    Method method = ReflectionUtils.findMethod(c, tableNote);
                    Object result = ReflectionUtils.invokeMethod(method, c.newInstance());
                    if (ObjectUtils.isEmpty(result)) {
                        throw new RuntimeException(method + "没有正确重写");
                    }else{
                        log.info("扫描实体类："+c+"，映射表："+result);
                    }
                } catch (IllegalAccessException e) {
                    e.printStackTrace();
                } catch (InstantiationException e) {
                    e.printStackTrace();
                }
            } else {
                log.error(c + "没有继承" + AbstractModel.class);
            }
        }

    }


}