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 com.mybatisflex.core.FlexGlobalConfig;
020import com.mybatisflex.core.datasource.DataSourceKey;
021import com.mybatisflex.core.util.StringUtil;
022import org.apache.ibatis.mapping.BoundSql;
023import org.apache.ibatis.mapping.ParameterMapping;
024import org.apache.ibatis.mapping.ParameterMode;
025import org.apache.ibatis.reflection.MetaObject;
026import org.apache.ibatis.reflection.ParamNameResolver;
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            Object parameter = boundSql.getParameterObject();
127
128            /** parameter 的组装请查看 getNamedParams 方法
129             * @see ParamNameResolver#getNamedParams(Object[])
130             */
131            if (parameter instanceof Map) {
132                TypeHandlerRegistry typeHandlerRegistry = configuration.getTypeHandlerRegistry();
133                if (((Map<?, ?>) parameter).containsKey(FlexConsts.SQL_ARGS)) {
134                    auditMessage.addParams(statement, ((Map<?, ?>) parameter).get(FlexConsts.SQL_ARGS));
135                } else if (((Map<?, ?>) parameter).containsKey("collection")) {
136                    Collection collection = (Collection) ((Map<?, ?>) parameter).get("collection");
137                    auditMessage.addParams(statement, collection.toArray());
138                } else if (((Map<?, ?>) parameter).containsKey("array")) {
139                    auditMessage.addParams(statement, ((Map<?, ?>) parameter).get("array"));
140                } else {
141                    List<ParameterMapping> parameterMappings = boundSql.getParameterMappings();
142                    for (ParameterMapping parameterMapping : parameterMappings) {
143                        if (parameterMapping.getMode() != ParameterMode.OUT) {
144                            Object value;
145                            String propertyName = parameterMapping.getProperty();
146                            if (boundSql.hasAdditionalParameter(propertyName)) {
147                                value = boundSql.getAdditionalParameter(propertyName);
148                            } else if (typeHandlerRegistry.hasTypeHandler(parameter.getClass())) {
149                                value = parameter;
150                            } else {
151                                MetaObject metaObject = configuration.newMetaObject(parameter);
152                                value = metaObject.getValue(propertyName);
153                            }
154                            auditMessage.addParams(statement, value);
155                        }
156                    }
157                }
158            }
159            messageCollector.collect(auditMessage);
160        }
161    }
162
163
164    @FunctionalInterface
165    public interface AuditRunnable<T> {
166        T execute() throws SQLException;
167    }
168
169}