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