001/*
002 *  Copyright (c) 2022-2024, 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 com.mybatisflex.core.FlexGlobalConfig;
020import com.mybatisflex.core.datasource.DataSourceKey;
021import com.mybatisflex.core.util.CollectionUtil;
022import com.mybatisflex.core.util.StringUtil;
023import org.apache.ibatis.mapping.BoundSql;
024import org.apache.ibatis.mapping.ParameterMapping;
025import org.apache.ibatis.mapping.ParameterMode;
026import org.apache.ibatis.reflection.MetaObject;
027import org.apache.ibatis.session.Configuration;
028import org.apache.ibatis.type.TypeHandlerRegistry;
029
030import java.sql.SQLException;
031import java.sql.Statement;
032import java.util.Collection;
033import java.util.List;
034import java.util.Map;
035
036/**
037 * 审计管理器，统一执行如何和配置入口
038 */
039public class AuditManager {
040
041
042    private AuditManager() {
043    }
044
045
046    private static MessageFactory messageFactory = new DefaultMessageFactory();
047
048    private static boolean auditEnable = false;
049    private static Clock clock = System::currentTimeMillis;
050    private static MessageCollector messageCollector = new ScheduledMessageCollector();
051
052    public static boolean isAuditEnable() {
053        return auditEnable;
054    }
055
056    public static void setAuditEnable(boolean auditEnable) {
057        AuditManager.auditEnable = auditEnable;
058    }
059
060    public static Clock getClock() {
061        return clock;
062    }
063
064    public static void setClock(Clock clock) {
065        AuditManager.clock = clock;
066    }
067
068    public static MessageFactory getMessageFactory() {
069        return messageFactory;
070    }
071
072    public static void setMessageFactory(MessageFactory messageFactory) {
073        AuditManager.messageFactory = messageFactory;
074    }
075
076    public static MessageCollector getMessageCollector() {
077        return messageCollector;
078    }
079
080
081    public static void setMessageReporter(MessageReporter messageReporter) {
082        MessageCollector newMessageCollector = new ScheduledMessageCollector(10, messageReporter);
083        setMessageCollector(newMessageCollector);
084    }
085
086    public static void setMessageCollector(MessageCollector messageCollector) {
087        MessageCollector temp = AuditManager.messageCollector;
088        AuditManager.messageCollector = messageCollector;
089        releaseScheduledMessageCollector(temp);
090
091    }
092
093    private static void releaseScheduledMessageCollector(MessageCollector messageCollector) {
094        if (messageCollector instanceof ScheduledMessageCollector) {
095            ((ScheduledMessageCollector) messageCollector).release();
096        }
097    }
098
099    @SuppressWarnings("rawtypes")
100    public static <T> T startAudit(AuditRunnable<T> supplier, String stmtId, Statement statement, BoundSql boundSql, Configuration configuration) throws SQLException {
101        AuditMessage auditMessage = messageFactory.create();
102        if (auditMessage == null) {
103            return supplier.execute();
104        }
105        auditMessage.setStmtId(stmtId);
106        String key = DataSourceKey.get();
107        if (StringUtil.noText(key)) {
108            key = FlexGlobalConfig.getDefaultConfig()
109                .getDataSource()
110                .getDefaultDataSourceKey();
111        }
112        auditMessage.setDsName(key);
113        auditMessage.setQueryTime(clock.getTick());
114        try {
115            T result = supplier.execute();
116            if (result instanceof Collection) {
117                auditMessage.setQueryCount(((Collection) result).size());
118            } else if (result instanceof Number) {
119                auditMessage.setQueryCount(((Number) result).intValue());
120            } else if (result != null) {
121                auditMessage.setQueryCount(1);
122            }
123            return result;
124        } finally {
125            auditMessage.setElapsedTime(clock.getTick() - auditMessage.getQueryTime());
126            auditMessage.setQuery(boundSql.getSql());
127
128            TypeHandlerRegistry typeHandlerRegistry = configuration.getTypeHandlerRegistry();
129            List<ParameterMapping> parameterMappings = boundSql.getParameterMappings();
130            Object parameter = boundSql.getParameterObject();
131
132            // 实现 XML 与 QueryWrapper 参数解析互不干涉
133
134            if (CollectionUtil.isNotEmpty(parameterMappings)) {
135                // 组装 XML 中的 #{user.age} 参数
136                for (ParameterMapping parameterMapping : parameterMappings) {
137                    if (parameterMapping.getMode() != ParameterMode.OUT) {
138                        Object value;
139                        String propertyName = parameterMapping.getProperty();
140                        if (boundSql.hasAdditionalParameter(propertyName)) {
141                            value = boundSql.getAdditionalParameter(propertyName);
142                        } else if (typeHandlerRegistry.hasTypeHandler(parameter.getClass())) {
143                            value = parameter;
144                        } else {
145                            MetaObject metaObject = configuration.newMetaObject(parameter);
146                            value = metaObject.getValue(propertyName);
147                        }
148                        auditMessage.addParams(statement, value);
149                    }
150                }
151            } else {
152                // 组装 QueryWrapper 里面的 age = ? 参数
153                // parameter 的组装请查看 ParamNameResolver#getNamedParams(Object[]) 方法
154                if (parameter instanceof Map) {
155                    Map<?, ?> map = (Map<?, ?>) parameter;
156                    if (map.containsKey(FlexConsts.SQL_ARGS)) {
157                        auditMessage.addParams(statement, map.get(FlexConsts.SQL_ARGS));
158                    } else if (map.containsKey("collection")) {
159                        Collection collection = (Collection) map.get("collection");
160                        auditMessage.addParams(statement, collection.toArray());
161                    } else if (map.containsKey("array")) {
162                        auditMessage.addParams(statement, map.get("array"));
163                    }
164                }
165            }
166
167            messageCollector.collect(auditMessage);
168        }
169    }
170
171
172    @FunctionalInterface
173    public interface AuditRunnable<T> {
174        T execute() throws SQLException;
175    }
176
177}