ca.uhn.fhir.jpa.migrate.taskdef.ArbitrarySqlTask 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.jpa.migrate.JdbcUtils;
import ca.uhn.fhir.util.VersionEnum;
import org.apache.commons.lang3.StringUtils;
import org.apache.commons.lang3.builder.EqualsBuilder;
import org.apache.commons.lang3.builder.HashCodeBuilder;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.jdbc.core.ColumnMapRowMapper;
import org.springframework.jdbc.core.JdbcTemplate;

import java.sql.SQLException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.function.Consumer;

public class ArbitrarySqlTask extends BaseTask {

	private static final Logger ourLog = LoggerFactory.getLogger(ArbitrarySqlTask.class);
	private final String myDescription;
	private final String myTableName;
	private List myTask = new ArrayList<>();
	private int myBatchSize = 1000;
	private String myExecuteOnlyIfTableExists;
	private List myConditionalOnExistenceOf = new ArrayList<>();

	/**
	 * Constructor
	 */
	public ArbitrarySqlTask(VersionEnum theRelease, String theVersion, String theTableName, String theDescription) {
		super(theRelease.toString(), theVersion);
		myTableName = theTableName;
		myDescription = theDescription;
	}

	public void addQuery(String theSql, QueryModeEnum theMode, Consumer> theConsumer) {
		myTask.add(new QueryTask(theSql, theMode, theConsumer));
	}

	@Override
	public void validate() {
		// nothing
	}

	@Override
	public void doExecute() throws SQLException {
		logInfo(ourLog, "Starting: {}", myDescription);

		if (StringUtils.isNotBlank(myExecuteOnlyIfTableExists)) {
			Set tableNames = JdbcUtils.getTableNames(getConnectionProperties());
			if (!tableNames.contains(myExecuteOnlyIfTableExists.toUpperCase())) {
				logInfo(ourLog, "Table {} does not exist - No action performed", myExecuteOnlyIfTableExists);
				return;
			}
		}

		for (TableAndColumn next : myConditionalOnExistenceOf) {
			JdbcUtils.ColumnType columnType =
					JdbcUtils.getColumnType(getConnectionProperties(), next.getTable(), next.getColumn());
			if (columnType == null) {
				logInfo(
						ourLog,
						"Table {} does not have column {} - No action performed",
						next.getTable(),
						next.getColumn());
				return;
			}
		}

		for (BaseTask next : myTask) {
			next.execute();
		}
	}

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

	public void setExecuteOnlyIfTableExists(String theExecuteOnlyIfTableExists) {
		myExecuteOnlyIfTableExists = theExecuteOnlyIfTableExists;
	}

	/**
	 * This task will only execute if the following column exists
	 */
	public void addExecuteOnlyIfColumnExists(String theTableName, String theColumnName) {
		myConditionalOnExistenceOf.add(new TableAndColumn(theTableName, theColumnName));
	}

	@Override
	protected void generateEquals(EqualsBuilder theBuilder, ca.uhn.fhir.jpa.migrate.taskdef.BaseTask theOtherObject) {
		ArbitrarySqlTask otherObject = (ArbitrarySqlTask) theOtherObject;
		theBuilder.append(myTableName, otherObject.myTableName);
	}

	@Override
	protected void generateHashCode(HashCodeBuilder theBuilder) {
		theBuilder.append(myTableName);
	}

	public enum QueryModeEnum {
		BATCH_UNTIL_NO_MORE
	}

	private static class TableAndColumn {
		private final String myTable;
		private final String myColumn;

		private TableAndColumn(String theTable, String theColumn) {
			myTable = theTable;
			myColumn = theColumn;
		}

		public String getTable() {
			return myTable;
		}

		public String getColumn() {
			return myColumn;
		}
	}

	private abstract class BaseTask {
		public abstract void execute();
	}

	private class QueryTask extends BaseTask {
		private final String mySql;
		private final Consumer> myConsumer;

		public QueryTask(String theSql, QueryModeEnum theMode, Consumer> theConsumer) {
			mySql = theSql;
			myConsumer = theConsumer;
			setDescription("Execute raw sql");
		}

		@Override
		public void execute() {
			if (isDryRun()) {
				return;
			}

			List> rows;
			do {
				logInfo(ourLog, "Querying for up to {} rows", myBatchSize);
				rows = getTxTemplate().execute(t -> {
					JdbcTemplate jdbcTemplate = newJdbcTemplate();
					jdbcTemplate.setMaxRows(myBatchSize);
					return jdbcTemplate.query(mySql, new ColumnMapRowMapper());
				});

				logInfo(ourLog, "Processing {} rows", rows.size());
				List> finalRows = rows;
				getTxTemplate().execute(t -> {
					for (Map nextRow : finalRows) {
						myConsumer.accept(nextRow);
					}
					return null;
				});
			} while (rows.size() > 0);
		}
	}
}