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