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;
030import org.slf4j.Logger;
031import org.slf4j.LoggerFactory;
032
033import java.util.ArrayList;
034import java.util.Collections;
035import java.util.List;
036import java.util.Queue;
037import java.util.concurrent.atomic.AtomicInteger;
038
039import static org.apache.commons.lang3.StringUtils.trim;
040
041public abstract class BaseCaptureQueriesListener
042                implements ProxyDataSourceBuilder.SingleQueryExecution, ProxyDataSourceBuilder.SingleMethodExecution {
043        private static final Logger ourLog = LoggerFactory.getLogger(BaseCaptureQueriesListener.class);
044        private boolean myCaptureQueryStackTrace = false;
045
046        /**
047         * This has an impact on performance! Use with caution.
048         */
049        public boolean isCaptureQueryStackTrace() {
050                return myCaptureQueryStackTrace;
051        }
052
053        /**
054         * This has an impact on performance! Use with caution.
055         */
056        public void setCaptureQueryStackTrace(boolean theCaptureQueryStackTrace) {
057                myCaptureQueryStackTrace = theCaptureQueryStackTrace;
058        }
059
060        @Override
061        public void execute(ExecutionInfo theExecutionInfo, List<QueryInfo> theQueryInfoList) {
062                final Queue<SqlQuery> queryList = provideQueryList();
063                if (queryList == null) {
064                        return;
065                }
066
067                RequestPartitionId requestPartitionId = HapiTransactionService.getRequestPartitionAssociatedWithThread();
068
069                /*
070                 * Note on how this works:
071                 *
072                 * Hibernate batches queries where the SQL is identical other than parameter values.
073                 *
074                 * So for example, if we call "select PID from SOMETABLE where NAME = ?" twice with
075                 * two different name values, this method will be called only once with 2 parameter
076                 * lists. In this case, we say size=2 and we capture the first parameters in case
077                 * we want to log them. We don't capture subsequent parameters currently because
078                 * this is only used for troubleshooting logging anyhow so the values are used
079                 * only as a representative example only.
080                 */
081
082                for (QueryInfo next : theQueryInfoList) {
083                        String sql = trim(next.getQuery());
084                        List<String> params;
085                        int size;
086                        if (next.getParametersList().size() > 0
087                                        && next.getParametersList().get(0).size() > 0) {
088                                size = next.getParametersList().size();
089                                List<ParameterSetOperation> values = next.getParametersList().get(0);
090                                params = new ArrayList<>();
091                                for (ParameterSetOperation t : values) {
092                                        if (t.getMethod().getName().equals("setNull")) {
093                                                params.add(null);
094                                        } else {
095                                                Object arg = t.getArgs()[1];
096                                                String s = arg != null ? arg.toString() : null;
097                                                params.add(s);
098                                        }
099                                }
100                        } else {
101                                params = Collections.emptyList();
102                                size = next.getParametersList().size();
103                        }
104
105                        StackTraceElement[] stackTraceElements = null;
106                        if (isCaptureQueryStackTrace()) {
107                                stackTraceElements = Thread.currentThread().getStackTrace();
108                        }
109
110                        long elapsedTime = theExecutionInfo.getElapsedTime();
111                        long startTime = System.currentTimeMillis() - elapsedTime;
112                        SqlQuery sqlQuery =
113                                        new SqlQuery(sql, params, startTime, elapsedTime, stackTraceElements, size, requestPartitionId);
114                        queryList.add(sqlQuery);
115                }
116        }
117
118        protected abstract Queue<SqlQuery> provideQueryList();
119
120        @Nullable
121        protected abstract AtomicInteger provideCommitCounter();
122
123        @Nullable
124        protected abstract AtomicInteger provideGetConnectionCounter();
125
126        @Nullable
127        protected abstract AtomicInteger provideRollbackCounter();
128
129        @Override
130        public void execute(MethodExecutionContext executionContext) {
131                AtomicInteger counter = null;
132                switch (executionContext.getMethod().getName()) {
133                        case "commit":
134                                counter = provideCommitCounter();
135                                break;
136                        case "rollback":
137                                counter = provideRollbackCounter();
138                                break;
139                        case "getConnection":
140                                counter = provideGetConnectionCounter();
141                                break;
142                }
143
144                if (counter != null) {
145                        counter.incrementAndGet();
146                }
147        }
148
149        /**
150         * @return Returns the number of times the connection pool was asked for a new connection
151         */
152        public int countGetConnections() {
153                return provideGetConnectionCounter().get();
154        }
155
156        /**
157         * @return Returns the number of DB commits which have happened against connections from the pool
158         */
159        public int countCommits() {
160                return provideCommitCounter().get();
161        }
162
163        /**
164         * @return Returns the number of DB rollbacks which have happened against connections from the pool
165         */
166        public int countRollbacks() {
167                return provideRollbackCounter().get();
168        }
169}