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.reflection.ParamNameResolver;
021
022import java.sql.SQLException;
023import java.util.Collection;
024import java.util.Map;
025
026/**
027 * 审计管理器，统一执行如何和配置入口
028 */
029public class AuditManager {
030
031    private static boolean auditEnable = false;
032    private static Clock clock = System::currentTimeMillis;
033    private static MessageFactory MessageFactory = new DefaultMessageFactory();
034    private static MessageCollector messageCollector = new ScheduledMessageCollector();
035
036    public static boolean isAuditEnable() {
037        return auditEnable;
038    }
039
040    public static void setAuditEnable(boolean auditEnable) {
041        AuditManager.auditEnable = auditEnable;
042    }
043
044    public static Clock getClock() {
045        return clock;
046    }
047
048    public static void setClock(Clock clock) {
049        AuditManager.clock = clock;
050    }
051
052    public static MessageFactory getMessageFactory() {
053        return MessageFactory;
054    }
055
056    public static void setMessageFactory(MessageFactory MessageFactory) {
057        AuditManager.MessageFactory = MessageFactory;
058    }
059
060    public static MessageCollector getMessageCollector() {
061        return messageCollector;
062    }
063
064
065    public static void setMessageReporter(MessageReporter messageReporter) {
066        MessageCollector newMessageCollector = new ScheduledMessageCollector(10, messageReporter);
067        setMessageCollector(newMessageCollector);
068    }
069
070    public static void setMessageCollector(MessageCollector messageCollector) {
071        MessageCollector temp = AuditManager.messageCollector;
072        AuditManager.messageCollector = messageCollector;
073        releaseScheduledMessageCollector(temp);
074
075    }
076
077    private static void releaseScheduledMessageCollector(MessageCollector messageCollector) {
078        if (messageCollector instanceof ScheduledMessageCollector) {
079            ((ScheduledMessageCollector) messageCollector).release();
080        }
081    }
082
083    public static <T> T startAudit(AuditRunnable<T> supplier, BoundSql boundSql) throws SQLException {
084        AuditMessage auditMessage = MessageFactory.create();
085        if (auditMessage == null) {
086            return supplier.execute();
087        }
088        auditMessage.setQueryTime(clock.getTick());
089        try {
090            T result = supplier.execute();
091            if (result instanceof Collection) {
092                auditMessage.setQueryCount(((Collection) result).size());
093            } else if (result != null) {
094                auditMessage.setQueryCount(1);
095            }
096            return result;
097        } finally {
098            auditMessage.setElapsedTime(clock.getTick() - auditMessage.getQueryTime());
099            auditMessage.setQuery(boundSql.getSql());
100            Object parameter = boundSql.getParameterObject();
101
102
103            /** parameter 的组装请查看 getNamedParams 方法
104             * @see ParamNameResolver#getNamedParams(Object[])
105             */
106            if (parameter instanceof Map) {
107                if (((Map<?, ?>) parameter).containsKey(FlexConsts.SQL_ARGS)) {
108                    auditMessage.addParams(((Map<?, ?>) parameter).get(FlexConsts.SQL_ARGS));
109                } else if (((Map<?, ?>) parameter).containsKey("collection")) {
110                    Collection collection = (Collection) ((Map<?, ?>) parameter).get("collection");
111                    auditMessage.addParams(collection.toArray());
112                } else if (((Map<?, ?>) parameter).containsKey("array")) {
113                    auditMessage.addParams(((Map<?, ?>) parameter).get("array"));
114                } else {
115                    for (int i = 1; i <= 100; i++) {
116                        if (((Map<?, ?>) parameter).containsKey(ParamNameResolver.GENERIC_NAME_PREFIX + i)) {
117                            auditMessage.addParams(((Map<?, ?>) parameter).get(ParamNameResolver.GENERIC_NAME_PREFIX + i));
118                        }
119                    }
120                }
121            }
122            messageCollector.collect(auditMessage);
123        }
124    }
125
126
127    @FunctionalInterface
128    public interface AuditRunnable<T> {
129        T execute() throws SQLException;
130    }
131
132}