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