mybatis分页练手

最近碰到个需求，要做个透明的mybatis分页功能，描述如下：
目标：搜索列表的Controller action要和原先保持一样，并且返回的json需要有分页信息，如：

@ResponseBody
@RequestMapping(value="/search", method={RequestMethod.POST})
public List<ProjectInfo> search(@RequestBody SearchProjectCommand command)
{
        List<ProjectInfo> projects=projectFetcher.search(command.getKey(), command.getFrom(), command.getTo());

        return projects;
}

返回信息：

{
  "successful": true,
  "message": null,
  "messages": null,
  "dateTime": 1505651777350,
  "body": {
    "totalCount": 2,
    "totalPage": 1,
    "records": [
      {
        "projectId": "1111",
        "projectName": "11111111111111",
        "title": "11111111111111"
      },
      {
        "projectId": "22222",
        "projectName": "222222",
        "title": "222222"
      }
    ]
  }
}

　　

关键点：

1. 针对Controller方法的aop
2. Mybatis interceptor && PagingContext保存分页信息
3. ResponseBodyAdvice（用于在输出json之前加入通用格式）

开始之前，先来看看消息格式，以及某些限制，主要是针对分页pageIndex这种参数的传递：

public abstract class PagingCommand {
    private int pageSize;
    private int pageIndex;

    public PagingCommand getPagingInfo()
    {
        return this;
    }

    public int getPageSize() {
        if(pageSize<=0)
            return Integer.MAX_VALUE;

        return pageSize;
    }

    public void setPageSize(int pageSize) {
        this.pageSize = pageSize;
    }

    public int getPageIndex() {
        if(pageIndex<0)
            return 0;

        return pageIndex;
    }

    public void setPageIndex(int pageIndex) {
        this.pageIndex = pageIndex;
    }
}




public class PagingResponse {

    private int totalCount;
    private int totalPage;
    private List<Object> records;

    public int getTotalCount() {
        return totalCount;
    }

    public void setTotalCount(int totalCount) {
        this.totalCount = totalCount;
    }

    public List<Object> getRecords() {
        return records;
    }

    public int getTotalPage() {
        return totalPage;
    }

    public void setTotalPage(int totalPage) {
        this.totalPage = totalPage;
    }

    public void setRecords(List<Object> records) {
        this.records = records;
    }
}

PagingCommand是抽象类，所有的具体Command必须继承这个Command
PagingResponse是分页结果

　　

先来看看横切入口AOP类：

@Aspect
@Component
public class PagingAop {
    private static final Logger logger = LoggerFactory.getLogger(PagingAop.class);

    @Pointcut("@annotation(org.springframework.web.bind.annotation.RequestMapping)")
    public void controllerMethodPointcut() {
    }

    @Around("controllerMethodPointcut()")
    public Object Interceptor(ProceedingJoinPoint pjp) throws Throwable {

        logger.info("Paging...");

        //找到是否具有PagingCommand的class作为输入参数
        //有，则放入PagingContext中
        for(Object arg:pjp.getArgs())
        {
            if(arg==null)
                continue;

            logger.info(arg.getClass().toString());
            if(PagingCommand.class.isAssignableFrom(arg.getClass()))
            {
                logger.info("需要分页行为");
                PagingContext.setPagingCommand((PagingCommand)arg);
            }
            else
            {
                logger.info("不需要分页行为");
            }
        }

        return pjp.proceed();
    }
}

代码很容易识别，判断参数是否是继承自PagingCommand，只要有1个继承自PagingCommand就会设置相应参数到PagingContext来标识需要分页处理，下面看看这个Context类：

public final class PagingContext {
    private static ThreadLocal<PagingCommand> pagingCommand=new ThreadLocal<PagingCommand>();
    private static ThreadLocal<Integer> totalCount=new ThreadLocal<Integer>();
    private static ThreadLocal<Integer> totalPage=new ThreadLocal<Integer>();

    public static void setPagingCommand(PagingCommand cmd)
    {
        pagingCommand.set(cmd);
    }

    public static PagingCommand getPagingCommand()
    {
        return pagingCommand.get();
    }

    public static boolean isPagingCommandEmpty()
    {
        if(pagingCommand.get()==null)
            return true;

        return  false;
    }


    public static int getTotalCount() {
        return totalCount.get();
    }

    public static void setTotalCount(int count) {
        totalCount.set(count);
    }

    public static boolean isTotalCountEmpty()
    {
        if(totalCount.get()==null)
            return true;

        return false;
    }


    public static int getTotalPage() {
        return totalPage.get();
    }

    public static void setTotalPage(int pages) {
        totalPage.set(pages);
    }
}

针对各个线程的ThreadLocal变量，但是目前只支持普通的httprequest线程才能正常工作，ThreadPool的有问题，等以后再解决。

下面是核心的mybatis分页插件了：

@Intercepts({@Signature(type=Executor.class,method="query",args={ MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class/*, CacheKey.class, BoundSql.class*/})})
public class PagingInterceptor implements Interceptor {
    private static final Logger logger = LoggerFactory.getLogger(PagingInterceptor.class);

    @Override
    public Object intercept(Invocation invocation) throws Throwable {


        logger.info("intercept.............");

        //判断是否需要分页行为, from PagingContext中
        if(PagingContext.isPagingCommandEmpty())
            return invocation.proceed();

        MappedStatement mappedStatement=(MappedStatement)invocation.getArgs()[0];
        Object parameter = invocation.getArgs()[1];
        BoundSql boundSql = mappedStatement.getBoundSql(parameter);
        String originalSql = boundSql.getSql().trim();

        //生成count sql，然后执行
        int totalCount = getTotalCount(mappedStatement, boundSql, originalSql);
        //set totalCount value to context
        PagingContext.setTotalCount(totalCount);

        int totalPages=calculateTotalPagesCount(totalCount, PagingContext.getPagingCommand().getPageSize());
        PagingContext.setTotalPage(totalPages);

        //生成分页limit sql，然后执行
        MappedStatement newMs = wrapPagedMappedStatement(mappedStatement, boundSql, originalSql);
        invocation.getArgs()[0]= newMs;

        return invocation.proceed();
    }

    private int calculateTotalPagesCount(int totalCount, int pageSize) {
        int pageCount=totalCount/pageSize;

        if(pageCount==0)
            return 1;

        if(pageCount*pageSize<=totalCount)
            return pageCount;

        return pageCount+1;
    }

    private MappedStatement wrapPagedMappedStatement(MappedStatement mappedStatement, BoundSql boundSql, String originalSql) {
        PagingCommand page= PagingContext.getPagingCommand();
        int offset = (page.getPageIndex()) * page.getPageSize();
        StringBuffer sb = new StringBuffer();
        sb.append(originalSql).append(" limit ").append(offset).append(",").append(page.getPageSize());
        BoundSql newBoundSql = MyBatisUtils.copyFromBoundSql(mappedStatement, boundSql, sb.toString());
        return MyBatisUtils.copyFromMappedStatement(mappedStatement,new BoundSqlSqlSource(newBoundSql));
    }

    private int getTotalCount(MappedStatement mappedStatement, BoundSql boundSql, String originalSql) throws SQLException {
        Object parameterObject = boundSql.getParameterObject();
        String countSql = getCountSql(originalSql);
        Connection connection=mappedStatement.getConfiguration().getEnvironment().getDataSource().getConnection()  ;
        PreparedStatement countStmt = connection.prepareStatement(countSql);
        BoundSql countBS = MyBatisUtils.copyFromBoundSql(mappedStatement, boundSql, countSql);
        DefaultParameterHandler parameterHandler = new DefaultParameterHandler(mappedStatement, parameterObject, countBS);
        parameterHandler.setParameters(countStmt);
        ResultSet rs = countStmt.executeQuery();
        int totalCount=0;
        if (rs.next()) {
            totalCount = rs.getInt(1);
        }
        rs.close();
        countStmt.close();
        connection.close();
        return totalCount;
    }

    private String getCountSql(String sql) {
        return "SELECT COUNT(1) FROM (" + sql + ") Mybatis_Pager_TBL_ALIAS";
    }

    @Override
    public Object plugin(Object o) {
        return Plugin.wrap(o, this);
    }

    @Override
    public void setProperties(Properties properties) {

    }
}

最后就一步了，就是写一个ResponseBodyAdvice来根据判断是否分页输出，来返回json：

@ControllerAdvice
public class GlobalMessageResponseBodyAdvice implements ResponseBodyAdvice {

    @Override
    public boolean supports(MethodParameter methodParameter, Class aClass) {
        return true;
    }

    @Override
    public Object beforeBodyWrite(Object o, MethodParameter methodParameter, MediaType mediaType, Class aClass, ServerHttpRequest serverHttpRequest, ServerHttpResponse serverHttpResponse) {
        
            Object payload = o;

            //判断是否需要分页
            if (isNeedPagingResponse()) {
                PagingResponse response = new PagingResponse();

                response.setTotalCount(PagingContext.getTotalCount());
                response.setTotalPage(PagingContext.getTotalPage());
                response.setRecords((List<Object>) payload);

                payload = response;
            }

            NormalMessage msg = new NormalMessage();
            msg.setSuccessful(true);
            msg.setMessage(null);
            msg.setBody(payload);
            return msg;
        
    }

    public boolean isNeedPagingResponse() {
        if(PagingContext.isPagingCommandEmpty())
            return false;

        return true;
    }
}

完成。