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, Statement statement, BoundSql boundSql, Configuration configuration) throws SQLException {
101        AuditMessage auditMessage = messageFactory.create();
102        if (auditMessage == null) {
103            return supplier.execute();
104        }
105        String key = DataSourceKey.get();
106        if (StringUtil.isBlank(key)) {
107            key = FlexGlobalConfig.getDefaultConfig()
108                .getDataSource()
109                .getDefaultDataSourceKey();
110        }
111        auditMessage.setDsName(key);
112        auditMessage.setQueryTime(clock.getTick());
113        try {
114            T result = supplier.execute();
115            if (result instanceof Collection) {
116                auditMessage.setQueryCount(((Collection) result).size());
117            } else if (result instanceof Number) {
118                auditMessage.setQueryCount(((Number) result).intValue());
119            } else if (result != null) {
120                auditMessage.setQueryCount(1);
121            }
122            return result;
123        } finally {
124            auditMessage.setElapsedTime(clock.getTick() - auditMessage.getQueryTime());
125            auditMessage.setQuery(boundSql.getSql());
126
127            TypeHandlerRegistry typeHandlerRegistry = configuration.getTypeHandlerRegistry();
128            List<ParameterMapping> parameterMappings = boundSql.getParameterMappings();
129            Object parameter = boundSql.getParameterObject();
130
131            // 实现 XML 与 QueryWrapper 参数解析互不干涉
132
133            if (CollectionUtil.isNotEmpty(parameterMappings)) {
134                // 组装 XML 中的 #{user.age} 参数
135                for (ParameterMapping parameterMapping : parameterMappings) {
136                    if (parameterMapping.getMode() != ParameterMode.OUT) {
137                        Object value;
138                        String propertyName = parameterMapping.getProperty();
139                        if (boundSql.hasAdditionalParameter(propertyName)) {
140                            value = boundSql.getAdditionalParameter(propertyName);
141                        } else if (typeHandlerRegistry.hasTypeHandler(parameter.getClass())) {
142                            value = parameter;
143                        } else {
144                            MetaObject metaObject = configuration.newMetaObject(parameter);
145                            value = metaObject.getValue(propertyName);
146                        }
147                        auditMessage.addParams(statement, value);
148                    }
149                }
150            } else {
151                // 组装 QueryWrapper 里面的 age = ? 参数
152                // parameter 的组装请查看 ParamNameResolver#getNamedParams(Object[]) 方法
153                if (parameter instanceof Map) {
154                    Map<?, ?> map = (Map<?, ?>) parameter;
155                    if (map.containsKey(FlexConsts.SQL_ARGS)) {
156                        auditMessage.addParams(statement, map.get(FlexConsts.SQL_ARGS));
157                    } else if (map.containsKey("collection")) {
158                        Collection collection = (Collection) map.get("collection");
159                        auditMessage.addParams(statement, collection.toArray());
160                    } else if (map.containsKey("array")) {
161                        auditMessage.addParams(statement, map.get("array"));
162                    }
163                }
164            }
165
166            messageCollector.collect(auditMessage);
167        }
168    }
169
170
171    @FunctionalInterface
172    public interface AuditRunnable<T> {
173        T execute() throws SQLException;
174    }
175
176}