001/**
002 * Copyright (c) 2022-2023, Mybatis-Flex (fuhai999@gmail.com).
003 * <p>
004 * Licensed under the Apache License, Version 2.0 (the "License");
005 * you may not use this file except in compliance with the License.
006 * You may obtain a copy of the License at
007 * <p>
008 * http://www.apache.org/licenses/LICENSE-2.0
009 * <p>
010 * Unless required by applicable law or agreed to in writing, software
011 * distributed under the License is distributed on an "AS IS" BASIS,
012 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
013 * See the License for the specific language governing permissions and
014 * limitations under the License.
015 */
016package com.mybatisflex.core.audit;
017
018import com.mybatisflex.core.FlexConsts;
019import org.apache.ibatis.mapping.BoundSql;
020import org.apache.ibatis.mapping.ParameterMapping;
021import org.apache.ibatis.mapping.ParameterMode;
022import org.apache.ibatis.reflection.MetaObject;
023import org.apache.ibatis.reflection.ParamNameResolver;
024import org.apache.ibatis.session.Configuration;
025import org.apache.ibatis.type.TypeHandlerRegistry;
026
027import java.sql.SQLException;
028import java.util.Collection;
029import java.util.List;
030import java.util.Map;
031
032/**
033 * 审计管理器，统一执行如何和配置入口
034 */
035public class AuditManager {
036
037    private static boolean auditEnable = false;
038    private static Clock clock = System::currentTimeMillis;
039    private static MessageFactory MessageFactory = new DefaultMessageFactory();
040    private static MessageCollector messageCollector = new ScheduledMessageCollector();
041
042    public static boolean isAuditEnable() {
043        return auditEnable;
044    }
045
046    public static void setAuditEnable(boolean auditEnable) {
047        AuditManager.auditEnable = auditEnable;
048    }
049
050    public static Clock getClock() {
051        return clock;
052    }
053
054    public static void setClock(Clock clock) {
055        AuditManager.clock = clock;
056    }
057
058    public static MessageFactory getMessageFactory() {
059        return MessageFactory;
060    }
061
062    public static void setMessageFactory(MessageFactory MessageFactory) {
063        AuditManager.MessageFactory = MessageFactory;
064    }
065
066    public static MessageCollector getMessageCollector() {
067        return messageCollector;
068    }
069
070
071    public static void setMessageReporter(MessageReporter messageReporter) {
072        MessageCollector newMessageCollector = new ScheduledMessageCollector(10, messageReporter);
073        setMessageCollector(newMessageCollector);
074    }
075
076    public static void setMessageCollector(MessageCollector messageCollector) {
077        MessageCollector temp = AuditManager.messageCollector;
078        AuditManager.messageCollector = messageCollector;
079        releaseScheduledMessageCollector(temp);
080
081    }
082
083    private static void releaseScheduledMessageCollector(MessageCollector messageCollector) {
084        if (messageCollector instanceof ScheduledMessageCollector) {
085            ((ScheduledMessageCollector) messageCollector).release();
086        }
087    }
088
089    public static <T> T startAudit(AuditRunnable<T> supplier, BoundSql boundSql, Configuration configuration) throws SQLException {
090        AuditMessage auditMessage = MessageFactory.create();
091        if (auditMessage == null) {
092            return supplier.execute();
093        }
094        auditMessage.setQueryTime(clock.getTick());
095        try {
096            T result = supplier.execute();
097            if (result instanceof Collection) {
098                auditMessage.setQueryCount(((Collection) result).size());
099            } else if (result != null) {
100                auditMessage.setQueryCount(1);
101            }
102            return result;
103        } finally {
104            auditMessage.setElapsedTime(clock.getTick() - auditMessage.getQueryTime());
105            auditMessage.setQuery(boundSql.getSql());
106            Object parameter = boundSql.getParameterObject();
107
108            /** parameter 的组装请查看 getNamedParams 方法
109             * @see ParamNameResolver#getNamedParams(Object[])
110             */
111            if (parameter instanceof Map) {
112                TypeHandlerRegistry typeHandlerRegistry = configuration.getTypeHandlerRegistry();
113                if (((Map<?, ?>) parameter).containsKey(FlexConsts.SQL_ARGS)) {
114                    auditMessage.addParams(((Map<?, ?>) parameter).get(FlexConsts.SQL_ARGS));
115                } else if (((Map<?, ?>) parameter).containsKey("collection")) {
116                    Collection collection = (Collection) ((Map<?, ?>) parameter).get("collection");
117                    auditMessage.addParams(collection.toArray());
118                } else if (((Map<?, ?>) parameter).containsKey("array")) {
119                    auditMessage.addParams(((Map<?, ?>) parameter).get("array"));
120                } else {
121                    List<ParameterMapping> parameterMappings = boundSql.getParameterMappings();
122                    for (ParameterMapping parameterMapping : parameterMappings) {
123                        if (parameterMapping.getMode() != ParameterMode.OUT) {
124                            Object value;
125                            String propertyName = parameterMapping.getProperty();
126                            if (boundSql.hasAdditionalParameter(propertyName)) {
127                                value = boundSql.getAdditionalParameter(propertyName);
128                            } else if (typeHandlerRegistry.hasTypeHandler(parameter.getClass())) {
129                                value = parameter;
130                            } else {
131                                MetaObject metaObject = configuration.newMetaObject(parameter);
132                                value = metaObject.getValue(propertyName);
133                            }
134                            auditMessage.addParams(value);
135                        }
136                    }
137                }
138            }
139            messageCollector.collect(auditMessage);
140        }
141    }
142
143
144    @FunctionalInterface
145    public interface AuditRunnable<T> {
146        T execute() throws SQLException;
147    }
148
149}