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}