001/*
002 *  Copyright (c) 2022-2025, Mybatis-Flex (fuhai999@gmail.com).
003 *  <p>
004 *  Licensed under the Apache License, Version 2.0 (the "License");
005 *  you may not use this file except in compliance with the License.
006 *  You may obtain a copy of the License at
007 *  <p>
008 *  http://www.apache.org/licenses/LICENSE-2.0
009 *  <p>
010 *  Unless required by applicable law or agreed to in writing, software
011 *  distributed under the License is distributed on an "AS IS" BASIS,
012 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
013 *  See the License for the specific language governing permissions and
014 *  limitations under the License.
015 */
016package com.mybatisflex.core.row;
017
018import com.mybatisflex.core.mybatis.MappedStatementTypes;
019import com.mybatisflex.core.paginate.Page;
020import com.mybatisflex.core.query.QueryWrapper;
021import org.apache.ibatis.executor.BatchResult;
022import org.apache.ibatis.session.ExecutorType;
023import org.apache.ibatis.session.SqlSession;
024import org.apache.ibatis.session.SqlSessionFactory;
025
026import java.util.Collection;
027import java.util.List;
028import java.util.Map;
029import java.util.function.BiConsumer;
030import java.util.function.Function;
031
032public class RowMapperInvoker {
033
034    private final SqlSessionFactory sqlSessionFactory;
035
036    public RowMapperInvoker(SqlSessionFactory sqlSessionFactory) {
037        this.sqlSessionFactory = sqlSessionFactory;
038    }
039
040    protected <R> R execute(Function<RowMapper, R> function) {
041        Class<?> currentType = MappedStatementTypes.getCurrentType();
042        if (currentType == null) {
043            try (SqlSession sqlSession = sqlSessionFactory.openSession(true)) {
044                RowMapper mapper = sqlSession.getMapper(RowMapper.class);
045                return function.apply(mapper);
046            }
047        } else {
048            MappedStatementTypes.clear();
049            try (SqlSession sqlSession = sqlSessionFactory.openSession(true)) {
050                RowMapper mapper = sqlSession.getMapper(RowMapper.class);
051                return function.apply(mapper);
052            } finally {
053                MappedStatementTypes.setCurrentType(currentType);
054            }
055        }
056    }
057
058    public int insert(String schema, String tableName, Row row) {
059        return execute(mapper -> mapper.insert(schema, tableName, row));
060    }
061
062
063    public int insertBySql(String sql, Object... args) {
064        return execute(mapper -> mapper.insertBySql(sql, args));
065    }
066
067    public int insertBatchWithFirstRowColumns(String schema, String tableName, List<Row> rows) {
068        return execute(mapper -> mapper.insertBatchWithFirstRowColumns(schema, tableName, rows));
069    }
070
071    public int deleteBySql(String sql, Object... args) {
072        return execute(mapper -> mapper.deleteBySql(sql, args));
073    }
074
075    public int deleteById(String schema, String tableName, Row row) {
076        return execute(mapper -> mapper.deleteById(schema, tableName, row));
077    }
078
079    public int deleteById(String schema, String tableName, String primaryKey, Object id) {
080        return execute(mapper -> mapper.deleteById(schema, tableName, primaryKey, id));
081    }
082
083    public int deleteBatchByIds(String schema, String tableName, String primaryKey, Collection<?> ids) {
084        return execute(mapper -> mapper.deleteBatchByIds(schema, tableName, primaryKey, ids));
085    }
086
087
088    public int deleteByQuery(String schema, String tableName, QueryWrapper queryWrapper) {
089        return execute(mapper -> mapper.deleteByQuery(schema, tableName, queryWrapper));
090    }
091
092    public int updateBySql(String sql, Object... args) {
093        return execute(mapper -> mapper.updateBySql(sql, args));
094    }
095
096
097    public <M, E> int[] executeBatch(Collection<E> datas, int batchSize, Class<M> mapperClass, BiConsumer<M, E> consumer) {
098        int[] results = new int[datas.size()];
099        try (SqlSession sqlSession = sqlSessionFactory.openSession(ExecutorType.BATCH, true)) {
100            M mapper = sqlSession.getMapper(mapperClass);
101            int counter = 0;
102            int resultsPos = 0;
103            for (E data : datas) {
104                consumer.accept(mapper, data);
105                if (++counter == batchSize) {
106                    counter = 0;
107                    List<BatchResult> batchResults = sqlSession.flushStatements();
108                    for (BatchResult batchResult : batchResults) {
109                        int[] updateCounts = batchResult.getUpdateCounts();
110                        for (int updateCount : updateCounts) {
111                            results[resultsPos++] = updateCount;
112                        }
113                    }
114                }
115            }
116            if (counter != 0) {
117                List<BatchResult> batchResults = sqlSession.flushStatements();
118                for (BatchResult batchResult : batchResults) {
119                    int[] updateCounts = batchResult.getUpdateCounts();
120                    for (int updateCount : updateCounts) {
121                        results[resultsPos++] = updateCount;
122                    }
123                }
124            }
125        }
126        return results;
127    }
128
129
130    public <M> int[] executeBatch(int totalSize, int batchSize, Class<M> mapperClass, BiConsumer<M, Integer> consumer) {
131        int[] results = new int[totalSize];
132        try (SqlSession sqlSession = sqlSessionFactory.openSession(ExecutorType.BATCH, true)) {
133            M mapper = sqlSession.getMapper(mapperClass);
134            int counter = 0;
135            int resultsPos = 0;
136            for (int i = 0; i < totalSize; i++) {
137                consumer.accept(mapper, i);
138                if (++counter == batchSize) {
139                    counter = 0;
140                    List<BatchResult> batchResults = sqlSession.flushStatements();
141                    for (BatchResult batchResult : batchResults) {
142                        int[] updateCounts = batchResult.getUpdateCounts();
143                        for (int updateCount : updateCounts) {
144                            results[resultsPos++] = updateCount;
145                        }
146                    }
147                }
148
149            }
150
151            if (counter != 0) {
152                List<BatchResult> batchResults = sqlSession.flushStatements();
153                for (BatchResult batchResult : batchResults) {
154                    int[] updateCounts = batchResult.getUpdateCounts();
155                    for (int updateCount : updateCounts) {
156                        results[resultsPos++] = updateCount;
157                    }
158                }
159            }
160        }
161        return results;
162    }
163
164    public int updateById(String schema, String tableName, Row row) {
165        return execute(mapper -> mapper.updateById(schema, tableName, row));
166    }
167
168    public int updateByQuery(String schema, String tableName, Row data, QueryWrapper queryWrapper) {
169        return execute(mapper -> mapper.updateByQuery(schema, tableName, data, queryWrapper));
170    }
171
172    public int updateBatchById(String schema, String tableName, List<Row> rows) {
173        return execute(mapper -> mapper.updateBatchById(schema, tableName, rows));
174    }
175
176    public Row selectOneBySql(String sql, Object... args) {
177        return execute(mapper -> mapper.selectOneBySql(sql, args));
178    }
179
180    public Row selectOneById(String schema, String tableName, Row row) {
181        return execute(mapper -> mapper.selectOneById(schema, tableName, row));
182    }
183
184    public Row selectOneById(String schema, String tableName, String primaryKey, Object id) {
185        return execute(mapper -> mapper.selectOneById(schema, tableName, primaryKey, id));
186    }
187
188    public Row selectOneByQuery(String schema, String tableName, QueryWrapper queryWrapper) {
189        return execute(mapper -> mapper.selectOneByQuery(schema, tableName, queryWrapper));
190    }
191
192    public List<Row> selectListBySql(String sql, Object... args) {
193        return execute(mapper -> mapper.selectListBySql(sql, args));
194    }
195
196    public List<Row> selectListByQuery(String schema, String tableName, QueryWrapper queryWrapper) {
197        return execute(mapper -> mapper.selectListByQuery(schema, tableName, queryWrapper));
198    }
199
200    public List<Row> selectAll(String schema, String tableName) {
201        return execute(mapper -> mapper.selectAll(schema, tableName));
202    }
203
204    public Map selectFirstAndSecondColumnsAsMapByQuery(String schema, String tableName, QueryWrapper queryWrapper) {
205        return execute(mapper -> mapper.selectFirstAndSecondColumnsAsMapByQuery(schema, tableName, queryWrapper));
206    }
207
208    public Map selectFirstAndSecondColumnsAsMap(String sql, Object... args) {
209        return execute(mapper -> mapper.selectFirstAndSecondColumnsAsMap(sql, args));
210    }
211
212    public Object selectObjectByQuery(String schema, String tableName, QueryWrapper queryWrapper) {
213        return execute(mapper -> mapper.selectObjectByQuery(schema, tableName, queryWrapper));
214    }
215
216    public List<Object> selectObjectListByQuery(String schema, String tableName, QueryWrapper queryWrapper) {
217        return execute(mapper -> mapper.selectObjectListByQuery(schema, tableName, queryWrapper));
218    }
219
220    public Object selectObject(String sql, Object... args) {
221        return execute(mapper -> mapper.selectObject(sql, args));
222    }
223
224    public List<Object> selectObjectList(String sql, Object... args) {
225        return execute(mapper -> mapper.selectObjectList(sql, args));
226    }
227
228    public long selectCount(String sql, Object... args) {
229        return execute(mapper -> mapper.selectCount(sql, args));
230    }
231
232
233    public long selectCountByQuery(String schema, String tableName, QueryWrapper queryWrapper) {
234        return execute(mapper -> mapper.selectCountByQuery(schema, tableName, queryWrapper));
235    }
236
237    public Page<Row> paginate(String schema, String tableName, Page<Row> page, QueryWrapper queryWrapper) {
238        return execute(mapper -> mapper.paginate(schema, tableName, page, queryWrapper));
239    }
240
241}