Mybatis分頁攔截器

版權聲明：本文為部落客原創文章，未經部落客允許不得轉載。 https://blog.csdn.net/u010741376/article/details/53517869

import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.plugin.*;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.SQLException;
import java.util.Map;
import java.util.Properties;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.locks.ReentrantLock;

/**
 * Mybatis - 通用分頁攔截器
 *
 */
@SuppressWarnings("rawtypes")
@Intercepts(@Signature(type = Executor.class, method = "query", args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class}))
public class PageHelper implements Interceptor {
    //sql工具類
    private SqlUtil sqlUtil;
    //屬性參數資訊
    private Properties properties;
    //配置對象方式
    private SqlUtilConfig sqlUtilConfig;
    //自動擷取dialect,如果沒有setProperties或setSqlUtilConfig，也可以正常進行
    private boolean autoDialect = true;
    //運作時自動擷取dialect
    private boolean autoRuntimeDialect;
    //多資料源時，擷取jdbcurl後是否關閉資料源
    private boolean closeConn = true;
    //緩存
    private Map<String, SqlUtil> urlSqlUtilMap = new ConcurrentHashMap<String, SqlUtil>();
    private ReentrantLock lock = new ReentrantLock();

    /**
     * 擷取任意查詢方法的count總數
     *
     * @param select
     * @return
     */
    public static long count(ISelect select) {
        Page<?> page = startPage(1, -1, true);
        select.doSelect();
        return page.getTotal();
    }

    /**
     * 開始分頁
     *
     * @param pageNum  頁碼
     * @param pageSize 每頁顯示數量
     */
    public static <E> Page<E> startPage(int pageNum, int pageSize) {
        return startPage(pageNum, pageSize, true);
    }

    /**
     * 開始分頁
     *
     * @param pageNum  頁碼
     * @param pageSize 每頁顯示數量
     * @param count    是否進行count查詢
     */
    public static <E> Page<E> startPage(int pageNum, int pageSize, boolean count) {
        return startPage(pageNum, pageSize, count, null);
    }

    /**
     * 開始分頁
     *
     * @param pageNum  頁碼
     * @param pageSize 每頁顯示數量
     * @param orderBy  排序
     */
    public static <E> Page<E> startPage(int pageNum, int pageSize, String orderBy) {
        Page<E> page = startPage(pageNum, pageSize);
        page.setOrderBy(orderBy);
        return page;
    }

    /**
     * 開始分頁
     *
     * @param offset 頁碼
     * @param limit  每頁顯示數量
     */
    public static <E> Page<E> offsetPage(int offset, int limit) {
        return offsetPage(offset, limit, true);
    }

    /**
     * 開始分頁
     *
     * @param offset 頁碼
     * @param limit  每頁顯示數量
     * @param count  是否進行count查詢
     */
    public static <E> Page<E> offsetPage(int offset, int limit, boolean count) {
        Page<E> page = new Page<E>(new int[]{offset, limit}, count);
        //當已經執行過orderBy的時候
        Page<E> oldPage = SqlUtil.getLocalPage();
        if (oldPage != null && oldPage.isOrderByOnly()) {
            page.setOrderBy(oldPage.getOrderBy());
        }
        SqlUtil.setLocalPage(page);
        return page;
    }

    /**
     * 開始分頁
     *
     * @param offset  頁碼
     * @param limit   每頁顯示數量
     * @param orderBy 排序
     */
    public static <E> Page<E> offsetPage(int offset, int limit, String orderBy) {
        Page<E> page = offsetPage(offset, limit);
        page.setOrderBy(orderBy);
        return page;
    }

    /**
     * 開始分頁
     *
     * @param pageNum    頁碼
     * @param pageSize   每頁顯示數量
     * @param count      是否進行count查詢
     * @param reasonable 分頁合理化,null時用預設配置
     */
    public static <E> Page<E> startPage(int pageNum, int pageSize, boolean count, Boolean reasonable) {
        return startPage(pageNum, pageSize, count, reasonable, null);
    }

    /**
     * 開始分頁
     *
     * @param pageNum      頁碼
     * @param pageSize     每頁顯示數量
     * @param count        是否進行count查詢
     * @param reasonable   分頁合理化,null時用預設配置
     * @param pageSizeZero true且pageSize=0時傳回全部結果，false時分頁,null時用預設配置
     */
    public static <E> Page<E> startPage(int pageNum, int pageSize, boolean count, Boolean reasonable, Boolean pageSizeZero) {
        Page<E> page = new Page<E>(pageNum, pageSize, count);
        page.setReasonable(reasonable);
        page.setPageSizeZero(pageSizeZero);
        //當已經執行過orderBy的時候
        Page<E> oldPage = SqlUtil.getLocalPage();
        if (oldPage != null && oldPage.isOrderByOnly()) {
            page.setOrderBy(oldPage.getOrderBy());
        }
        SqlUtil.setLocalPage(page);
        return page;
    }

    /**
     * 開始分頁
     *
     * @param params
     */
    public static <E> Page<E> startPage(Object params) {
        Page<E> page = SqlUtil.getPageFromObject(params);
        //當已經執行過orderBy的時候
        Page<E> oldPage = SqlUtil.getLocalPage();
        if (oldPage != null && oldPage.isOrderByOnly()) {
            page.setOrderBy(oldPage.getOrderBy());
        }
        SqlUtil.setLocalPage(page);
        return page;
    }

    /**
     * 排序
     *
     * @param orderBy
     */
    public static void orderBy(String orderBy) {
        Page<?> page = SqlUtil.getLocalPage();
        if (page != null) {
            page.setOrderBy(orderBy);
        } else {
            page = new Page();
            page.setOrderBy(orderBy);
            page.setOrderByOnly(true);
            SqlUtil.setLocalPage(page);
        }
    }

    /**
     * 擷取orderBy
     *
     * @return
     */
    public static String getOrderBy() {
        Page<?> page = SqlUtil.getLocalPage();
        if (page != null) {
            String orderBy = page.getOrderBy();
            if (StringUtil.isEmpty(orderBy)) {
                return null;
            } else {
                return orderBy;
            }
        }
        return null;
    }

    /**
     * Mybatis攔截器方法
     *
     * @param invocation 攔截器入參
     * @return 傳回執行結果
     * @throws Throwable 抛出異常
     */
    public Object intercept(Invocation invocation) throws Throwable {
        if (autoRuntimeDialect) {
            SqlUtil sqlUtil = getSqlUtil(invocation);
            return sqlUtil.processPage(invocation);
        } else {
            if (autoDialect) {
                initSqlUtil(invocation);
            }
            return sqlUtil.processPage(invocation);
        }
    }

    /**
     * 初始化sqlUtil
     *
     * @param invocation
     */
    public synchronized void initSqlUtil(Invocation invocation) {
        if (this.sqlUtil == null) {
            this.sqlUtil = getSqlUtil(invocation);
            if (!autoRuntimeDialect) {
                properties = null;
                sqlUtilConfig = null;
            }
            autoDialect = false;
        }
    }

    /**
     * 擷取url
     *
     * @param dataSource
     * @return
     */
    public String getUrl(DataSource dataSource){
        Connection conn = null;
        try {
            conn = dataSource.getConnection();
            return conn.getMetaData().getURL();
        } catch (SQLException e) {
            throw new RuntimeException(e);
        } finally {
            if(conn != null){
                try {
                    if(closeConn){
                        conn.close();
                    }
                } catch (SQLException e) {
                    //ignore
                }
            }
        }
    }

    /**
     * 根據datasource建立對應的sqlUtil
     *
     * @param invocation
     */
    public SqlUtil getSqlUtil(Invocation invocation) {
        MappedStatement ms = (MappedStatement) invocation.getArgs()[0];
        //改為對dataSource做緩存
        DataSource dataSource = ms.getConfiguration().getEnvironment().getDataSource();
        String url = getUrl(dataSource);
        if (urlSqlUtilMap.containsKey(url)) {
            return urlSqlUtilMap.get(url);
        }
        try {
            lock.lock();
            if (urlSqlUtilMap.containsKey(url)) {
                return urlSqlUtilMap.get(url);
            }
            if (StringUtil.isEmpty(url)) {
                throw new RuntimeException("無法自動擷取jdbcUrl，請在分頁插件中配置dialect參數!");
            }
            String dialect = Dialect.fromJdbcUrl(url);
            if (dialect == null) {
                throw new RuntimeException("無法自動擷取資料庫類型，請通過dialect參數指定!");
            }
            SqlUtil sqlUtil = null;
            if (properties != null) {
                sqlUtil = new SqlUtil(dialect, properties.getProperty("sqlCacheClass"));
                sqlUtil.setProperties(properties);
            } else if (sqlUtilConfig != null) {
                sqlUtil = new SqlUtil(dialect, sqlUtilConfig.getSqlCacheClass());
                sqlUtil.setSqlUtilConfig(sqlUtilConfig);
            }
            urlSqlUtilMap.put(url, sqlUtil);
            return sqlUtil;
        } finally {
            lock.unlock();
        }
    }

    /**
     * 隻攔截Executor
     *
     * @param target
     * @return
     */
    public Object plugin(Object target) {
        if (target instanceof Executor) {
            return Plugin.wrap(target, this);
        } else {
            return target;
        }
    }

    private void checkVersion() {
        //MyBatis3.2.0版本校驗
        try {
            Class.forName("org.apache.ibatis.scripting.xmltags.SqlNode");//SqlNode是3.2.0之後新增的類
        } catch (ClassNotFoundException e) {
            throw new RuntimeException("您使用的MyBatis版本太低，MyBatis分頁插件PageHelper支援MyBatis3.2.0及以上版本!");
        }
    }

    /**
     * 設定屬性值
     *
     * @param p 屬性值
     */
    public void setProperties(Properties p) {
        checkVersion();
        //多資料源時，擷取jdbcurl後是否關閉資料源
        String closeConn = p.getProperty("closeConn");
        //解決#97
        if(StringUtil.isNotEmpty(closeConn)){
            this.closeConn = Boolean.parseBoolean(closeConn);
        }
        //初始化SqlUtil的PARAMS
        SqlUtil.setParams(p.getProperty("params"));
        //資料庫方言
        String dialect = p.getProperty("dialect");
        String runtimeDialect = p.getProperty("autoRuntimeDialect");
        if (StringUtil.isNotEmpty(runtimeDialect) && runtimeDialect.equalsIgnoreCase("TRUE")) {
            this.autoRuntimeDialect = true;
            this.autoDialect = false;
            this.properties = p;
        } else if (StringUtil.isEmpty(dialect)) {
            autoDialect = true;
            this.properties = p;
        } else {
            autoDialect = false;
            sqlUtil = new SqlUtil(dialect, p.getProperty("sqlCacheClass"));
            sqlUtil.setProperties(p);
        }
    }

    /**
     * 設定屬性值
     *
     * @param config
     */
    public void setSqlUtilConfig(SqlUtilConfig config) {
        checkVersion();
        //初始化SqlUtil的PARAMS
        SqlUtil.setParams(config.getParams());
        //多資料源時，擷取jdbcurl後是否關閉資料源
        this.closeConn = config.isCloseConn();
        if (config.isAutoRuntimeDialect()) {
            this.autoRuntimeDialect = true;
            this.autoDialect = false;
            this.sqlUtilConfig = config;
        } else if (StringUtil.isEmpty(config.getDialect())) {
            autoDialect = true;
            this.sqlUtilConfig = config;
        } else {
            autoDialect = false;
            sqlUtil = new SqlUtil(config.getDialect(), config.getSqlCacheClass());
            sqlUtil.setSqlUtilConfig(config);
        }
    }
}