001package ca.uhn.fhir.jpa.migrate.taskdef;
002
003/*-
004 * #%L
005 * HAPI FHIR Server - SQL Migration
006 * %%
007 * Copyright (C) 2014 - 2022 Smile CDR, Inc.
008 * %%
009 * Licensed under the Apache License, Version 2.0 (the "License");
010 * you may not use this file except in compliance with the License.
011 * You may obtain a copy of the License at
012 *
013 *      http://www.apache.org/licenses/LICENSE-2.0
014 *
015 * Unless required by applicable law or agreed to in writing, software
016 * distributed under the License is distributed on an "AS IS" BASIS,
017 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
018 * See the License for the specific language governing permissions and
019 * limitations under the License.
020 * #L%
021 */
022
023import ca.uhn.fhir.i18n.Msg;
024import ca.uhn.fhir.util.StopWatch;
025import ca.uhn.fhir.util.VersionEnum;
026import com.google.common.collect.ForwardingMap;
027import org.apache.commons.lang3.concurrent.BasicThreadFactory;
028import org.slf4j.Logger;
029import org.slf4j.LoggerFactory;
030import org.springframework.jdbc.core.ColumnMapRowMapper;
031import org.springframework.jdbc.core.JdbcTemplate;
032import org.springframework.jdbc.core.RowCallbackHandler;
033
034import java.sql.ResultSet;
035import java.sql.SQLException;
036import java.util.ArrayList;
037import java.util.Date;
038import java.util.HashMap;
039import java.util.List;
040import java.util.Map;
041import java.util.concurrent.Future;
042import java.util.concurrent.LinkedBlockingQueue;
043import java.util.concurrent.RejectedExecutionException;
044import java.util.concurrent.RejectedExecutionHandler;
045import java.util.concurrent.ThreadPoolExecutor;
046import java.util.concurrent.TimeUnit;
047import java.util.function.Function;
048
049public abstract class BaseColumnCalculatorTask extends BaseTableColumnTask {
050
051        protected static final Logger ourLog = LoggerFactory.getLogger(BaseColumnCalculatorTask.class);
052        private int myBatchSize = 10000;
053        private ThreadPoolExecutor myExecutor;
054        private String myPidColumnName;
055
056        /**
057         * Constructor
058         */
059        public BaseColumnCalculatorTask(VersionEnum theRelease, String theVersion) {
060                this(theRelease.toString(), theVersion);
061        }
062
063        /**
064         * Constructor
065         */
066        public BaseColumnCalculatorTask(String theRelease, String theVersion) {
067                super(theRelease, theVersion);
068        }
069
070        public void setBatchSize(int theBatchSize) {
071                myBatchSize = theBatchSize;
072        }
073
074        /**
075         * Allows concrete implementations to decide if they should be skipped.
076         *
077         * @return a boolean indicating whether or not to skip execution of the task.
078         */
079        protected abstract boolean shouldSkipTask();
080
081        @Override
082        public synchronized void doExecute() throws SQLException {
083                if (isDryRun() || shouldSkipTask()) {
084                        return;
085                }
086
087                initializeExecutor();
088
089                try {
090
091                        while (true) {
092                                MyRowCallbackHandler rch = new MyRowCallbackHandler();
093                                getTxTemplate().execute(t -> {
094                                        JdbcTemplate jdbcTemplate = newJdbcTemplate();
095                                        jdbcTemplate.setMaxRows(100000);
096
097                                        String sql = "SELECT * FROM " + getTableName() + " WHERE " + getWhereClause();
098                                        logInfo(ourLog, "Finding up to {} rows in {} that requires calculations, using query: {}", myBatchSize, getTableName(), sql);
099
100                                        jdbcTemplate.query(sql, rch);
101                                        rch.done();
102
103                                        return null;
104                                });
105
106                                rch.submitNext();
107                                List<Future<?>> futures = rch.getFutures();
108                                if (futures.isEmpty()) {
109                                        break;
110                                }
111
112                                logInfo(ourLog, "Waiting for {} tasks to complete", futures.size());
113                                for (Future<?> next : futures) {
114                                        try {
115                                                next.get();
116                                        } catch (Exception e) {
117                                                throw new SQLException(Msg.code(69) + e);
118                                        }
119                                }
120
121                        }
122
123                } finally {
124                        destroyExecutor();
125                }
126        }
127
128        private void destroyExecutor() {
129                myExecutor.shutdownNow();
130        }
131
132        private void initializeExecutor() {
133                int maximumPoolSize = Runtime.getRuntime().availableProcessors();
134
135                LinkedBlockingQueue<Runnable> executorQueue = new LinkedBlockingQueue<>(maximumPoolSize);
136                BasicThreadFactory threadFactory = new BasicThreadFactory.Builder()
137                        .namingPattern("worker-" + "-%d")
138                        .daemon(false)
139                        .priority(Thread.NORM_PRIORITY)
140                        .build();
141                RejectedExecutionHandler rejectedExecutionHandler = new RejectedExecutionHandler() {
142                        @Override
143                        public void rejectedExecution(Runnable theRunnable, ThreadPoolExecutor theExecutor) {
144                                logInfo(ourLog, "Note: Executor queue is full ({} elements), waiting for a slot to become available!", executorQueue.size());
145                                StopWatch sw = new StopWatch();
146                                try {
147                                        executorQueue.put(theRunnable);
148                                } catch (InterruptedException theE) {
149                                        throw new RejectedExecutionException(Msg.code(70) + "Task " + theRunnable.toString() +
150                                                " rejected from " + theE.toString());
151                                }
152                                logInfo(ourLog, "Slot become available after {}ms", sw.getMillis());
153                        }
154                };
155                myExecutor = new ThreadPoolExecutor(
156                        1,
157                        maximumPoolSize,
158                        0L,
159                        TimeUnit.MILLISECONDS,
160                        executorQueue,
161                        threadFactory,
162                        rejectedExecutionHandler);
163        }
164
165        public void setPidColumnName(String thePidColumnName) {
166                myPidColumnName = thePidColumnName;
167        }
168
169        private Future<?> updateRows(List<Map<String, Object>> theRows) {
170                Runnable task = () -> {
171                        StopWatch sw = new StopWatch();
172                        getTxTemplate().execute(t -> {
173
174                                // Loop through rows
175                                assert theRows != null;
176                                for (Map<String, Object> nextRow : theRows) {
177
178                                        Map<String, Object> newValues = new HashMap<>();
179                                        MandatoryKeyMap<String, Object> nextRowMandatoryKeyMap = new MandatoryKeyMap<>(nextRow);
180
181                                        // Apply calculators
182                                        for (Map.Entry<String, Function<MandatoryKeyMap<String, Object>, Object>> nextCalculatorEntry : myCalculators.entrySet()) {
183                                                String nextColumn = nextCalculatorEntry.getKey();
184                                                Function<MandatoryKeyMap<String, Object>, Object> nextCalculator = nextCalculatorEntry.getValue();
185                                                Object value = nextCalculator.apply(nextRowMandatoryKeyMap);
186                                                newValues.put(nextColumn, value);
187                                        }
188
189                                        // Generate update SQL
190                                        StringBuilder sqlBuilder = new StringBuilder();
191                                        List<Object> arguments = new ArrayList<>();
192                                        sqlBuilder.append("UPDATE ");
193                                        sqlBuilder.append(getTableName());
194                                        sqlBuilder.append(" SET ");
195                                        for (Map.Entry<String, Object> nextNewValueEntry : newValues.entrySet()) {
196                                                if (arguments.size() > 0) {
197                                                        sqlBuilder.append(", ");
198                                                }
199                                                sqlBuilder.append(nextNewValueEntry.getKey()).append(" = ?");
200                                                arguments.add(nextNewValueEntry.getValue());
201                                        }
202                                        sqlBuilder.append(" WHERE " + myPidColumnName + " = ?");
203                                        arguments.add((Number) nextRow.get(myPidColumnName));
204
205                                        // Apply update SQL
206                                        newJdbcTemplate().update(sqlBuilder.toString(), arguments.toArray());
207                                }
208                                return theRows.size();
209                        });
210                        logInfo(ourLog, "Updated {} rows on {} in {}", theRows.size(), getTableName(), sw.toString());
211                };
212                return myExecutor.submit(task);
213        }
214
215        public static class MandatoryKeyMap<K, V> extends ForwardingMap<K, V> {
216
217                private final Map<K, V> myWrap;
218
219                public MandatoryKeyMap(Map<K, V> theWrap) {
220                        myWrap = theWrap;
221                }
222
223                @Override
224                public V get(Object theKey) {
225                        if (!containsKey(theKey)) {
226                                throw new IllegalArgumentException(Msg.code(71) + "No key: " + theKey);
227                        }
228                        return super.get(theKey);
229                }
230
231                public String getString(String theKey) {
232                        return (String) get(theKey);
233                }
234
235                public Date getDate(String theKey) {
236                        return (Date) get(theKey);
237                }
238
239                @Override
240                protected Map<K, V> delegate() {
241                        return myWrap;
242                }
243
244                public String getResourceType() {
245                        return getString("RES_TYPE");
246                }
247
248                public String getParamName() {
249                        return getString("SP_NAME");
250                }
251        }
252
253        private class MyRowCallbackHandler implements RowCallbackHandler {
254
255                private List<Map<String, Object>> myRows = new ArrayList<>();
256                private List<Future<?>> myFutures = new ArrayList<>();
257
258                @Override
259                public void processRow(ResultSet rs) throws SQLException {
260                        Map<String, Object> row = new ColumnMapRowMapper().mapRow(rs, 0);
261                        myRows.add(row);
262
263                        if (myRows.size() >= myBatchSize) {
264                                submitNext();
265                        }
266                }
267
268                private void submitNext() {
269                        if (myRows.size() > 0) {
270                                myFutures.add(updateRows(myRows));
271                                myRows = new ArrayList<>();
272                        }
273                }
274
275                public List<Future<?>> getFutures() {
276                        return myFutures;
277                }
278
279                public void done() {
280                        if (myRows.size() > 0) {
281                                submitNext();
282                        }
283                }
284        }
285}