customized mybatis mapper proxy

package com.raybow.mybatis; import com.raybow.core.Universe; import com.raybow.error.AppErrors; import com.raybow.model.base.*; import com.raybow.model.result.Results; import com.raybow.util.Action; import com.raybow.util.CloneUtils; import com.raybow.util.Utils; import org.apache.ibatis.binding.MapperMethod; import org.apache.ibatis.mapping.SqlCommandType; import org.apache.ibatis.reflection.ExceptionUtil; import org.apache.ibatis.session.SqlSession; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import java.io.Serializable; import java.lang.reflect.Field; import java.lang.reflect.InvocationHandler; import java.lang.reflect.Method; import java.lang.reflect.Proxy; import java.util.*; import static com.raybow.util.Utils.*; public class MagicMapperProxy<T> implements InvocationHandler, Serializable { protected final Logger LOGGER = LoggerFactory.getLogger(this.getClass()); private static final String[] IGNORE_PROPS = new String[]{ "deleted", "ownerId", "ownerType", "createdBy", "createdTime", "updatedBy", "updatedTime", "resultsTotalCount", "slug"}; private final SqlSession sqlSession; private final SqlSession slaveSqlSession; private final Class<T> mapperInterface; private final Map<Method, MagicMapperMethod> methodCache; private final Map<Method, Method> hackedMethodCache = new HashMap<>(); private final Map<Class, Method> getByIdMethodCache = new HashMap<>(); public static <T> T newMapperProxy(Class<T> mapperInterface, SqlSession sqlSession, SqlSession slaveSqlSession, Map<Method, MapperMethod> methodCache) { ClassLoader classLoader = mapperInterface.getClassLoader(); Class<?>[] interfaces = new Class[]{mapperInterface}; MagicMapperProxy proxy = new MagicMapperProxy( mapperInterface, sqlSession, slaveSqlSession, methodCache); return (T) Proxy.newProxyInstance(classLoader, interfaces, proxy); } public MagicMapperProxy(Class<T> mapperInterface, SqlSession sqlSession, SqlSession slaveSqlSession, Map<Method, MagicMapperMethod> methodCache) { this.sqlSession = sqlSession; this.slaveSqlSession = slaveSqlSession; this.mapperInterface = mapperInterface; this.methodCache = methodCache; } public Object invoke(Object proxy, Method method, Object[] args) throws Throwable { if (Object.class.equals(method.getDeclaringClass())) { try { return method.invoke(this, args); } catch (Throwable t) { throw ExceptionUtil.unwrapThrowable(t); } } if (method.getReturnType() == Results.class) { return processResultsReturn(proxy, method, args); } String methodName = method.getName(); final MagicMapperMethod mapperMethod = cachedMapperMethod(method); SqlCommandType commandType = getCommandType(methodName, getSqlSession()); // check slave query checkSlaveQuery(commandType); // pre-process model preHijackModel(proxy, args, commandType); Object result = mapperMethod.execute(getSqlSession(), args); // post-process model postHijackModel(args); if (result == null) return null; // deserialize single model if (BaseModel.class.isAssignableFrom(result.getClass())) { return buildResultModel(result); } // deserialize model list if (List.class.isAssignableFrom(result.getClass())) { Iterator iter = ((Collection) result).iterator(); if (!iter.hasNext()) return result; List<BaseModel> finalResults = new ArrayList<>(); while (iter.hasNext()) { Object element = iter.next(); if (element == null) continue; if (!BaseModel.class.isAssignableFrom(element.getClass())) return result; finalResults.add(buildResultModel(element)); } return finalResults; } return result; } public Object processResultsReturn(Object proxy, Method method, Object[] args) throws Throwable { // initialize results Results results = new Results(); try { Method hackedMethod = hackedMethodCache.get(method); if (hackedMethod == null) { hackedMethod = CloneUtils.clone(method); // hack mapper method return type Field returnTypeField = hackedMethod.getClass().getDeclaredField("returnType"); returnTypeField.setAccessible(true); returnTypeField.set(hackedMethod, List.class); hackedMethodCache.put(method, hackedMethod); } // get result items results.setItems((List) invoke(proxy, hackedMethod, args)); // get result total count // enter counting mode Universe.current().setCountingMode(true); Object result = invoke(proxy, hackedMethod, args); // use model id to to store total count for now results.setTotal(0); if (((Iterable) result).iterator().hasNext()) { Object value = ((Iterable) result).iterator().next(); if (value instanceof Countable) { results.setTotal(((Countable) value).getResultsTotalCount()); } } } finally { // exit counting mode Universe.current().setCountingMode(false); } return results; } private MagicMapperMethod cachedMapperMethod(Method method) { MagicMapperMethod mapperMethod = methodCache.get(method); // don't worry about concurrent if (mapperMethod == null) { mapperMethod = new MagicMapperMethod(mapperInterface, method, getSqlSession().getConfiguration()); methodCache.put(method, mapperMethod); } return mapperMethod; } private void preHijackModel(final Object proxy, Object args[], final SqlCommandType commandType) { processArguments(args, new Action<Object>() { public void apply(Object arg) { casLock((BaseModel) arg, commandType, proxy); processModel((BaseModel) arg, commandType); } }); } private void postHijackModel(Object args[]) { processArguments(args, new Action<Object>() { public void apply(Object arg) { ((BaseModel) arg).setPayload(null); } }); } private void processModel(BaseModel model, SqlCommandType commandType) { // set default value setDefaultValue(model, commandType); // erase 'payload' field in toJSON result model.setPayload(null); model.setPayload(toJson(model, IGNORE_PROPS)); // validate model if (model instanceof Validatable) { ((Validatable) model).validate(); } // check securable model if (commandType == SqlCommandType.UPDATE && (model instanceof Securable)) { ((Securable) model).checkOwner(); } } private void casLock(BaseModel model, SqlCommandType commandType, Object proxy) { if (commandType != SqlCommandType.UPDATE || !(model instanceof Lockable)) { return; } Method getByIdMethod = getByIdMethodCache.get(model.getClass()); if (getByIdMethod == null) { try { getByIdMethod = mapperInterface.getMethod("get", Long.class); } catch (NoSuchMethodException e) { // ignore } getByIdMethodCache.put(model.getClass(), getByIdMethod); } Object result; try { result = invoke(proxy, getByIdMethod, new Object[]{model.getId()}); } catch (Throwable e) { throw new RuntimeException(e); } String currentStamp = ((Lockable) result).getStamp(); // throw exception if update stale object if (currentStamp != null && !Utils.equals(currentStamp, ((Lockable) model).getStamp())) { throw AppErrors.INSTANCE.common("Entity was updated or soft deleted by another transaction.").exception(); } // refresh lock stamp ((Lockable) model).setStamp(UUID.randomUUID().toString()); } private void processArguments(Object args[], Action<Object> callback) { if (args == null) return; for (Object arg : args) { if (arg == null) continue; if (arg instanceof BaseModel) { callback.apply(arg); } if (arg instanceof List) { List listArg = (List) arg; Class<?> componentType = guessComponentType(listArg); if (componentType != null && BaseModel.class.isAssignableFrom(componentType)) { for (Object element : listArg) { callback.apply(element); } } } } } private void setDefaultValue(BaseModel model, SqlCommandType commandType) { if (commandType == SqlCommandType.INSERT) { if (model.getSlug() == null) { model.setSlug(UUID.randomUUID().toString()); } if (model.getCreatedTime() == null) { model.setCreatedTime(Utils.now()); } if (model.getCreatedBy() == null) { model.setCreatedBy(str(Universe.current().getAccountId())); } } if (commandType == SqlCommandType.UPDATE || commandType == SqlCommandType.INSERT) { model.setUpdatedTime(Utils.now()); model.setUpdatedBy(str(Universe.current().getAccountId())); } if (commandType == SqlCommandType.INSERT && (model instanceof Lockable)) { ((Lockable) model).setStamp(UUID.randomUUID().toString()); } if (model.getOwnerId() == null) { model.setOwnerId(Universe.current().getAccountId()); } if (model.getOwnerType() == null) { model.setOwnerType(Universe.current().getAccountType()); } if (model.getDeleted() == null) { model.setDeleted(Boolean.FALSE); } } private SqlCommandType getCommandType(String methodName, SqlSession session) { return session.getConfiguration().getMappedStatement( this.mapperInterface.getName() + "." + methodName).getSqlCommandType(); } private BaseModel buildResultModel(Object result) { // special logic for counting mode if (Universe.current().isCountingMode()) return (BaseModel) result; BaseModel finalResult = (BaseModel) (Utils.fromJson(((BaseModel) result).getPayload(), result.getClass())); finalResult.setId(((BaseModel) result).getId()); finalResult.setSlug(((BaseModel) result).getSlug()); finalResult.setOwnerId(((BaseModel) result).getOwnerId()); finalResult.setOwnerType(((BaseModel) result).getOwnerType()); finalResult.setCreatedBy(((BaseModel) result).getCreatedBy()); finalResult.setCreatedTime(((BaseModel) result).getCreatedTime()); finalResult.setUpdatedBy(((BaseModel) result).getUpdatedBy()); finalResult.setUpdatedTime(((BaseModel) result).getUpdatedTime()); finalResult.setDeleted(((BaseModel) result).getDeleted()); // clear payload finalResult.setPayload(null); // check securable model if (finalResult instanceof Securable) { ((Securable) finalResult).checkOwner(); } // check whether in compact mode or not if (Universe.current().isCompactMode()) { finalResult.compact(); return finalResult; } return finalResult; } private SqlSession getSqlSession() { return Universe.current().isSlaveQuery() ? slaveSqlSession : sqlSession; } private void checkSlaveQuery(SqlCommandType commandType) { if (Universe.current().isSlaveQuery() && commandType != SqlCommandType.SELECT) { throw new IllegalStateException("Only [SELECT] query can be executed against slave database."); } } }