001/*-
002 * #%L
003 * HAPI FHIR Storage api
004 * %%
005 * Copyright (C) 2014 - 2024 Smile CDR, Inc.
006 * %%
007 * Licensed under the Apache License, Version 2.0 (the "License");
008 * you may not use this file except in compliance with the License.
009 * You may obtain a copy of the License at
010 *
011 *      http://www.apache.org/licenses/LICENSE-2.0
012 *
013 * Unless required by applicable law or agreed to in writing, software
014 * distributed under the License is distributed on an "AS IS" BASIS,
015 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
016 * See the License for the specific language governing permissions and
017 * limitations under the License.
018 * #L%
019 */
020package ca.uhn.fhir.jpa.util;
021
022import ca.uhn.fhir.interceptor.model.RequestPartitionId;
023import ca.uhn.fhir.jpa.dao.tx.HapiTransactionService;
024import jakarta.annotation.Nullable;
025import net.ttddyy.dsproxy.ExecutionInfo;
026import net.ttddyy.dsproxy.QueryInfo;
027import net.ttddyy.dsproxy.listener.MethodExecutionContext;
028import net.ttddyy.dsproxy.proxy.ParameterSetOperation;
029import net.ttddyy.dsproxy.support.ProxyDataSourceBuilder;
030
031import java.util.Collections;
032import java.util.List;
033import java.util.Queue;
034import java.util.concurrent.atomic.AtomicInteger;
035import java.util.stream.Collectors;
036
037import static org.apache.commons.lang3.StringUtils.trim;
038
039public abstract class BaseCaptureQueriesListener
040                implements ProxyDataSourceBuilder.SingleQueryExecution, ProxyDataSourceBuilder.SingleMethodExecution {
041
042        private boolean myCaptureQueryStackTrace = false;
043
044        /**
045         * This has an impact on performance! Use with caution.
046         */
047        public boolean isCaptureQueryStackTrace() {
048                return myCaptureQueryStackTrace;
049        }
050
051        /**
052         * This has an impact on performance! Use with caution.
053         */
054        public void setCaptureQueryStackTrace(boolean theCaptureQueryStackTrace) {
055                myCaptureQueryStackTrace = theCaptureQueryStackTrace;
056        }
057
058        @Override
059        public void execute(ExecutionInfo theExecutionInfo, List<QueryInfo> theQueryInfoList) {
060                final Queue<SqlQuery> queryList = provideQueryList();
061                if (queryList == null) {
062                        return;
063                }
064
065                RequestPartitionId requestPartitionId = HapiTransactionService.getRequestPartitionAssociatedWithThread();
066
067                /*
068                 * Note on how this works:
069                 *
070                 * Hibernate batches queries where the SQL is identical other than parameter values.
071                 *
072                 * So for example, if we call "select PID from SOMETABLE where NAME = ?" twice with
073                 * two different name values, this method will be called only once with 2 parameter
074                 * lists. In this case, we say size=2 and we capture the first parameters in case
075                 * we want to log them. We don't capture subsequent parameters currently because
076                 * this is only used for troubleshooting logging anyhow so the values are used
077                 * only as a representative example only.
078                 */
079
080                for (QueryInfo next : theQueryInfoList) {
081                        String sql = trim(next.getQuery());
082                        List<String> params;
083                        int size;
084                        if (next.getParametersList().size() > 0
085                                        && next.getParametersList().get(0).size() > 0) {
086                                size = next.getParametersList().size();
087                                List<ParameterSetOperation> values = next.getParametersList().get(0);
088                                params = values.stream()
089                                                .map(t -> t.getArgs()[1])
090                                                .map(t -> t != null ? t.toString() : "NULL")
091                                                .collect(Collectors.toList());
092                        } else {
093                                params = Collections.emptyList();
094                                size = next.getParametersList().size();
095                        }
096
097                        StackTraceElement[] stackTraceElements = null;
098                        if (isCaptureQueryStackTrace()) {
099                                stackTraceElements = Thread.currentThread().getStackTrace();
100                        }
101
102                        long elapsedTime = theExecutionInfo.getElapsedTime();
103                        long startTime = System.currentTimeMillis() - elapsedTime;
104                        SqlQuery sqlQuery =
105                                        new SqlQuery(sql, params, startTime, elapsedTime, stackTraceElements, size, requestPartitionId);
106                        queryList.add(sqlQuery);
107                }
108        }
109
110        protected abstract Queue<SqlQuery> provideQueryList();
111
112        @Nullable
113        protected abstract AtomicInteger provideCommitCounter();
114
115        @Nullable
116        protected abstract AtomicInteger provideRollbackCounter();
117
118        @Override
119        public void execute(MethodExecutionContext executionContext) {
120                AtomicInteger counter = null;
121                switch (executionContext.getMethod().getName()) {
122                        case "commit":
123                                counter = provideCommitCounter();
124                                break;
125                        case "rollback":
126                                counter = provideRollbackCounter();
127                                break;
128                }
129
130                if (counter != null) {
131                        counter.incrementAndGet();
132                }
133        }
134
135        public int countCommits() {
136                return provideCommitCounter().get();
137        }
138
139        public int countRollbacks() {
140                return provideRollbackCounter().get();
141        }
142}