ca.uhn.fhir.jpa.migrate.taskdef.BaseColumnCalculatorTask Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of hapi-fhir-sql-migrate Show documentation
Tooling for migrating SQL schemas.
The newest version!
/*-
 * #%L
 * HAPI FHIR Server - SQL Migration
 * %%
 * Copyright (C) 2014 - 2024 Smile CDR, Inc.
 * %%
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 * #L%
 */
package ca.uhn.fhir.jpa.migrate.taskdef;

import ca.uhn.fhir.i18n.Msg;
import ca.uhn.fhir.util.StopWatch;
import ca.uhn.fhir.util.VersionEnum;
import com.google.common.collect.ForwardingMap;
import org.apache.commons.lang3.concurrent.BasicThreadFactory;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.jdbc.core.ColumnMapRowMapper;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.jdbc.core.RowCallbackHandler;

import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.Date;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.Future;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.RejectedExecutionException;
import java.util.concurrent.RejectedExecutionHandler;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.function.Function;

public abstract class BaseColumnCalculatorTask extends BaseTableColumnTask {

	protected static final Logger ourLog = LoggerFactory.getLogger(BaseColumnCalculatorTask.class);
	private int myBatchSize = 10000;
	private ThreadPoolExecutor myExecutor;
	private String myPidColumnName;

	/**
	 * Constructor
	 */
	public BaseColumnCalculatorTask(VersionEnum theRelease, String theVersion) {
		this(theRelease.toString(), theVersion);
	}

	/**
	 * Constructor
	 */
	public BaseColumnCalculatorTask(String theRelease, String theVersion) {
		super(theRelease, theVersion);
	}

	public void setBatchSize(int theBatchSize) {
		myBatchSize = theBatchSize;
	}

	/**
	 * Allows concrete implementations to decide if they should be skipped.
	 *
	 * @return a boolean indicating whether or not to skip execution of the task.
	 */
	protected abstract boolean shouldSkipTask();

	@Override
	public synchronized void doExecute() throws SQLException {
		if (isDryRun() || shouldSkipTask()) {
			return;
		}

		initializeExecutor();

		try {

			while (true) {
				MyRowCallbackHandler rch = new MyRowCallbackHandler();
				getTxTemplate().execute(t -> {
					JdbcTemplate jdbcTemplate = newJdbcTemplate();
					jdbcTemplate.setMaxRows(100000);

					String sql = "SELECT * FROM " + getTableName() + " WHERE " + getWhereClause();
					logInfo(
							ourLog,
							"Finding up to {} rows in {} that requires calculations, using query: {}",
							myBatchSize,
							getTableName(),
							sql);

					jdbcTemplate.query(sql, rch);
					rch.done();

					return null;
				});

				rch.submitNext();
				List> futures = rch.getFutures();
				if (futures.isEmpty()) {
					break;
				}

				logInfo(ourLog, "Waiting for {} tasks to complete", futures.size());
				for (Future next : futures) {
					try {
						next.get();
					} catch (Exception e) {
						throw new SQLException(Msg.code(69) + e, e);
					}
				}
			}

		} finally {
			destroyExecutor();
		}
	}

	private void destroyExecutor() {
		myExecutor.shutdownNow();
	}

	private void initializeExecutor() {
		int maximumPoolSize = Runtime.getRuntime().availableProcessors();

		LinkedBlockingQueue executorQueue = new LinkedBlockingQueue<>(maximumPoolSize);
		BasicThreadFactory threadFactory = new BasicThreadFactory.Builder()
				.namingPattern("worker-" + "-%d")
				.daemon(false)
				.priority(Thread.NORM_PRIORITY)
				.build();
		RejectedExecutionHandler rejectedExecutionHandler = new RejectedExecutionHandler() {
			@Override
			public void rejectedExecution(Runnable theRunnable, ThreadPoolExecutor theExecutor) {
				logInfo(
						ourLog,
						"Note: Executor queue is full ({} elements), waiting for a slot to become available!",
						executorQueue.size());
				StopWatch sw = new StopWatch();
				try {
					executorQueue.put(theRunnable);
				} catch (InterruptedException theE) {
					throw new RejectedExecutionException(
							Msg.code(70) + "Task " + theRunnable.toString() + " rejected from " + theE.toString());
				}
				logInfo(ourLog, "Slot become available after {}ms", sw.getMillis());
			}
		};
		myExecutor = new ThreadPoolExecutor(
				maximumPoolSize,
				maximumPoolSize,
				0L,
				TimeUnit.MILLISECONDS,
				executorQueue,
				threadFactory,
				rejectedExecutionHandler);
	}

	public BaseColumnCalculatorTask setPidColumnName(String thePidColumnName) {
		myPidColumnName = thePidColumnName;
		return this;
	}

	private Future updateRows(List> theRows) {
		Runnable task = () -> {
			StopWatch sw = new StopWatch();
			getTxTemplate().execute(t -> {

				// Loop through rows
				assert theRows != null;
				for (Map nextRow : theRows) {

					Map newValues = new HashMap<>();
					MandatoryKeyMap nextRowMandatoryKeyMap = new MandatoryKeyMap<>(nextRow);

					// Apply calculators
					for (Map.Entry, Object>> nextCalculatorEntry :
							myCalculators.entrySet()) {
						String nextColumn = nextCalculatorEntry.getKey();
						Function, Object> nextCalculator =
								nextCalculatorEntry.getValue();
						Object value = nextCalculator.apply(nextRowMandatoryKeyMap);
						newValues.put(nextColumn, value);
					}

					// Generate update SQL
					StringBuilder sqlBuilder = new StringBuilder();
					List