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