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.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 = JdbcUtils.getColumnType(getConnectionProperties(), next.getTable(), next.getColumn());
078                        if (columnType == null) {
079                                logInfo(ourLog, "Table {} does not have column {} - No action performed", next.getTable(), next.getColumn());
080                                return;
081                        }
082                }
083
084                for (BaseTask next : myTask) {
085                        next.execute();
086                }
087
088        }
089
090        public void setBatchSize(int theBatchSize) {
091                myBatchSize = theBatchSize;
092        }
093
094        public void setExecuteOnlyIfTableExists(String theExecuteOnlyIfTableExists) {
095                myExecuteOnlyIfTableExists = theExecuteOnlyIfTableExists;
096        }
097
098        /**
099         * This task will only execute if the following column exists
100         */
101        public void addExecuteOnlyIfColumnExists(String theTableName, String theColumnName) {
102                myConditionalOnExistenceOf.add(new TableAndColumn(theTableName, theColumnName));
103        }
104
105        @Override
106        protected void generateEquals(EqualsBuilder theBuilder, ca.uhn.fhir.jpa.migrate.taskdef.BaseTask theOtherObject) {
107                ArbitrarySqlTask otherObject = (ArbitrarySqlTask) theOtherObject;
108                theBuilder.append(myTableName, otherObject.myTableName);
109        }
110
111        @Override
112        protected void generateHashCode(HashCodeBuilder theBuilder) {
113                theBuilder.append(myTableName);
114        }
115
116        public enum QueryModeEnum {
117                BATCH_UNTIL_NO_MORE
118        }
119
120        private static class TableAndColumn {
121                private final String myTable;
122                private final String myColumn;
123
124                private TableAndColumn(String theTable, String theColumn) {
125                        myTable = theTable;
126                        myColumn = theColumn;
127                }
128
129                public String getTable() {
130                        return myTable;
131                }
132
133                public String getColumn() {
134                        return myColumn;
135                }
136        }
137
138        private abstract class BaseTask {
139                public abstract void execute();
140        }
141
142        private class QueryTask extends BaseTask {
143                private final String mySql;
144                private final Consumer<Map<String, Object>> myConsumer;
145
146                public QueryTask(String theSql, QueryModeEnum theMode, Consumer<Map<String, Object>> theConsumer) {
147                        mySql = theSql;
148                        myConsumer = theConsumer;
149                        setDescription("Execute raw sql");
150                }
151
152
153                @Override
154                public void execute() {
155                        if (isDryRun()) {
156                                return;
157                        }
158
159                        List<Map<String, Object>> rows;
160                        do {
161                                logInfo(ourLog, "Querying for up to {} rows", myBatchSize);
162                                rows = getTxTemplate().execute(t -> {
163                                        JdbcTemplate jdbcTemplate = newJdbcTemplate();
164                                        jdbcTemplate.setMaxRows(myBatchSize);
165                                        return jdbcTemplate.query(mySql, new ColumnMapRowMapper());
166                                });
167
168                                logInfo(ourLog, "Processing {} rows", rows.size());
169                                List<Map<String, Object>> finalRows = rows;
170                                getTxTemplate().execute(t -> {
171                                        for (Map<String, Object> nextRow : finalRows) {
172                                                myConsumer.accept(nextRow);
173                                        }
174                                        return null;
175                                });
176                        } while (rows.size() > 0);
177                }
178        }
179}