001/*-
002 * #%L
003 * HAPI FHIR Storage api
004 * %%
005 * Copyright (C) 2014 - 2025 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 jakarta.annotation.Nullable;
023import org.slf4j.Logger;
024import org.slf4j.LoggerFactory;
025
026import java.util.ArrayDeque;
027import java.util.ArrayList;
028import java.util.List;
029import java.util.Queue;
030import java.util.concurrent.atomic.AtomicInteger;
031import java.util.stream.Collectors;
032
033public class CurrentThreadCaptureQueriesListener extends BaseCaptureQueriesListener {
034
035        private static final ThreadLocal<Queue<SqlQuery>> ourQueues = new ThreadLocal<>();
036        private static final ThreadLocal<AtomicInteger> ourGetConnections = new ThreadLocal<>();
037        private static final ThreadLocal<AtomicInteger> ourCommits = new ThreadLocal<>();
038        private static final ThreadLocal<AtomicInteger> ourRollbacks = new ThreadLocal<>();
039        private static final Logger ourLog = LoggerFactory.getLogger(CurrentThreadCaptureQueriesListener.class);
040
041        @Override
042        protected Queue<SqlQuery> provideQueryList() {
043                return ourQueues.get();
044        }
045
046        @Override
047        protected AtomicInteger provideCommitCounter() {
048                return ourCommits.get();
049        }
050
051        @Nullable
052        @Override
053        protected AtomicInteger provideGetConnectionCounter() {
054                return ourGetConnections.get();
055        }
056
057        @Override
058        protected AtomicInteger provideRollbackCounter() {
059                return ourRollbacks.get();
060        }
061
062        /**
063         * Get the current queue of items and stop collecting
064         */
065        public static SqlQueryList getCurrentQueueAndStopCapturing() {
066                Queue<SqlQuery> retVal = ourQueues.get();
067                ourQueues.remove();
068                ourGetConnections.remove();
069                ourCommits.remove();
070                ourRollbacks.remove();
071                if (retVal == null) {
072                        return new SqlQueryList();
073                }
074                return new SqlQueryList(retVal);
075        }
076
077        /**
078         * Starts capturing queries for the current thread.
079         * <p>
080         * Note that you should strongly consider calling this in a
081         * try-finally block to ensure that you also call
082         * {@link #getCurrentQueueAndStopCapturing()} afterward. Otherwise
083         * this method is a potential memory leak!
084         * </p>
085         */
086        public static void startCapturing() {
087                ourQueues.set(new ArrayDeque<>());
088                ourGetConnections.set(new AtomicInteger(0));
089                ourCommits.set(new AtomicInteger(0));
090                ourRollbacks.set(new AtomicInteger(0));
091        }
092
093        /**
094         * Log all captured SELECT queries
095         *
096         * @return Returns the number of queries captured
097         */
098        public static int logQueriesForCurrentThreadAndStopCapturing(int... theIndexes) {
099                List<String> queries = getCurrentQueueAndStopCapturing().stream()
100                                .map(CircularQueueCaptureQueriesListener::formatQueryAsSql)
101                                .collect(Collectors.toList());
102
103                if (theIndexes != null && theIndexes.length > 0) {
104                        List<String> newList = new ArrayList<>();
105                        for (int i = 0; i < theIndexes.length; i++) {
106                                newList.add(queries.get(theIndexes[i]));
107                        }
108                        queries = newList;
109                }
110
111                ourLog.info("Select Queries:\n{}", String.join("\n", queries));
112
113                return queries.size();
114        }
115}