001/*-
002 * #%L
003 * HAPI FHIR Server - SQL Migration
004 * %%
005 * Copyright (C) 2014 - 2026 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.jpa.migrate.JdbcUtils;
023import ca.uhn.fhir.util.VersionEnum;
024import org.apache.commons.lang3.StringUtils;
025import org.apache.commons.lang3.builder.EqualsBuilder;
026import org.apache.commons.lang3.builder.HashCodeBuilder;
027import org.slf4j.Logger;
028import org.slf4j.LoggerFactory;
029import org.springframework.jdbc.core.ColumnMapRowMapper;
030import org.springframework.jdbc.core.JdbcTemplate;
031
032import java.sql.SQLException;
033import java.util.ArrayList;
034import java.util.List;
035import java.util.Map;
036import java.util.Set;
037import java.util.function.Consumer;
038
039public class ArbitrarySqlTask extends BaseTask {
040
041        private static final Logger ourLog = LoggerFactory.getLogger(ArbitrarySqlTask.class);
042        private final String myDescription;
043        private final String myTableName;
044        private List<BaseTask> myTask = new ArrayList<>();
045        private int myBatchSize = 1000;
046        private String myExecuteOnlyIfTableExists;
047        private List<TableAndColumn> myConditionalOnExistenceOf = new ArrayList<>();
048
049        /**
050         * Constructor
051         */
052        public ArbitrarySqlTask(VersionEnum theRelease, String theVersion, String theTableName, String theDescription) {
053                super(theRelease.toString(), theVersion);
054                myTableName = theTableName;
055                myDescription = theDescription;
056        }
057
058        public void addQuery(String theSql, QueryModeEnum theMode, Consumer<Map<String, Object>> theConsumer) {
059                myTask.add(new QueryTask(theSql, theMode, theConsumer));
060        }
061
062        @Override
063        public void validate() {
064                // nothing
065        }
066
067        @Override
068        public void doExecute() throws SQLException {
069                logInfo(ourLog, "Starting: {}", myDescription);
070
071                if (StringUtils.isNotBlank(myExecuteOnlyIfTableExists)) {
072                        Set<String> tableNames = JdbcUtils.getTableNames(getConnectionProperties());
073                        if (!tableNames.contains(myExecuteOnlyIfTableExists.toUpperCase())) {
074                                logInfo(ourLog, "Table {} does not exist - No action performed", myExecuteOnlyIfTableExists);
075                                return;
076                        }
077                }
078
079                for (TableAndColumn next : myConditionalOnExistenceOf) {
080                        JdbcUtils.ColumnType columnType =
081                                        JdbcUtils.getColumnType(getConnectionProperties(), next.getTable(), next.getColumn());
082                        if (columnType == null) {
083                                logInfo(
084                                                ourLog,
085                                                "Table {} does not have column {} - No action performed",
086                                                next.getTable(),
087                                                next.getColumn());
088                                return;
089                        }
090                }
091
092                for (BaseTask next : myTask) {
093                        next.execute();
094                }
095        }
096
097        public void setBatchSize(int theBatchSize) {
098                myBatchSize = theBatchSize;
099        }
100
101        public void setExecuteOnlyIfTableExists(String theExecuteOnlyIfTableExists) {
102                myExecuteOnlyIfTableExists = theExecuteOnlyIfTableExists;
103        }
104
105        /**
106         * This task will only execute if the following column exists
107         */
108        public void addExecuteOnlyIfColumnExists(String theTableName, String theColumnName) {
109                myConditionalOnExistenceOf.add(new TableAndColumn(theTableName, theColumnName));
110        }
111
112        @Override
113        protected void generateEquals(EqualsBuilder theBuilder, ca.uhn.fhir.jpa.migrate.taskdef.BaseTask theOtherObject) {
114                ArbitrarySqlTask otherObject = (ArbitrarySqlTask) theOtherObject;
115                theBuilder.append(myTableName, otherObject.myTableName);
116        }
117
118        @Override
119        protected void generateHashCode(HashCodeBuilder theBuilder) {
120                theBuilder.append(myTableName);
121        }
122
123        public enum QueryModeEnum {
124                BATCH_UNTIL_NO_MORE
125        }
126
127        private static class TableAndColumn {
128                private final String myTable;
129                private final String myColumn;
130
131                private TableAndColumn(String theTable, String theColumn) {
132                        myTable = theTable;
133                        myColumn = theColumn;
134                }
135
136                public String getTable() {
137                        return myTable;
138                }
139
140                public String getColumn() {
141                        return myColumn;
142                }
143        }
144
145        private abstract class BaseTask {
146                public abstract void execute();
147        }
148
149        private class QueryTask extends BaseTask {
150                private final String mySql;
151                private final Consumer<Map<String, Object>> myConsumer;
152
153                public QueryTask(String theSql, QueryModeEnum theMode, Consumer<Map<String, Object>> theConsumer) {
154                        mySql = theSql;
155                        myConsumer = theConsumer;
156                        setDescription("Execute raw sql");
157                }
158
159                @Override
160                public void execute() {
161                        if (isDryRun()) {
162                                return;
163                        }
164
165                        List<Map<String, Object>> rows;
166                        do {
167                                logInfo(ourLog, "Querying for up to {} rows", myBatchSize);
168                                rows = getTxTemplate().execute(t -> {
169                                        JdbcTemplate jdbcTemplate = newJdbcTemplate();
170                                        jdbcTemplate.setMaxRows(myBatchSize);
171                                        return jdbcTemplate.query(mySql, new ColumnMapRowMapper());
172                                });
173
174                                logInfo(ourLog, "Processing {} rows", rows.size());
175                                List<Map<String, Object>> finalRows = rows;
176                                getTxTemplate().execute(t -> {
177                                        for (Map<String, Object> nextRow : finalRows) {
178                                                myConsumer.accept(nextRow);
179                                        }
180                                        return null;
181                                });
182                        } while (rows.size() > 0);
183                }
184        }
185}