1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18 package org.apache.shardingsphere.driver.jdbc.core.statement;
19
20 import lombok.AccessLevel;
21 import lombok.Getter;
22 import org.apache.shardingsphere.driver.executor.callback.add.StatementAddCallback;
23 import org.apache.shardingsphere.driver.executor.engine.batch.preparedstatement.DriverExecuteBatchExecutor;
24 import org.apache.shardingsphere.driver.executor.engine.facade.DriverExecutorFacade;
25 import org.apache.shardingsphere.driver.jdbc.adapter.AbstractPreparedStatementAdapter;
26 import org.apache.shardingsphere.driver.jdbc.core.connection.ShardingSphereConnection;
27 import org.apache.shardingsphere.driver.jdbc.core.resultset.GeneratedKeysResultSet;
28 import org.apache.shardingsphere.driver.jdbc.core.resultset.ShardingSphereResultSet;
29 import org.apache.shardingsphere.driver.jdbc.core.statement.metadata.ShardingSphereParameterMetaData;
30 import org.apache.shardingsphere.infra.annotation.HighFrequencyInvocation;
31 import org.apache.shardingsphere.infra.binder.context.aware.ParameterAware;
32 import org.apache.shardingsphere.infra.binder.context.segment.insert.keygen.GeneratedKeyContext;
33 import org.apache.shardingsphere.infra.binder.context.statement.SQLStatementContext;
34 import org.apache.shardingsphere.infra.binder.context.statement.type.dml.InsertStatementContext;
35 import org.apache.shardingsphere.infra.binder.engine.SQLBindEngine;
36 import org.apache.shardingsphere.infra.database.core.keygen.GeneratedKeyColumnProvider;
37 import org.apache.shardingsphere.infra.database.core.spi.DatabaseTypedSPILoader;
38 import org.apache.shardingsphere.infra.database.core.type.DatabaseType;
39 import org.apache.shardingsphere.infra.exception.core.ShardingSpherePreconditions;
40 import org.apache.shardingsphere.infra.exception.dialect.SQLExceptionTransformEngine;
41 import org.apache.shardingsphere.infra.exception.kernel.syntax.EmptySQLException;
42 import org.apache.shardingsphere.infra.executor.sql.prepare.driver.jdbc.JDBCDriverType;
43 import org.apache.shardingsphere.infra.executor.sql.prepare.driver.jdbc.StatementOption;
44 import org.apache.shardingsphere.infra.hint.HintManager;
45 import org.apache.shardingsphere.infra.hint.HintValueContext;
46 import org.apache.shardingsphere.infra.hint.SQLHintUtils;
47 import org.apache.shardingsphere.infra.metadata.ShardingSphereMetaData;
48 import org.apache.shardingsphere.infra.metadata.database.ShardingSphereDatabase;
49 import org.apache.shardingsphere.infra.rule.attribute.datanode.DataNodeRuleAttribute;
50 import org.apache.shardingsphere.infra.rule.attribute.resoure.StorageConnectorReusableRuleAttribute;
51 import org.apache.shardingsphere.infra.session.query.QueryContext;
52 import org.apache.shardingsphere.parser.rule.SQLParserRule;
53 import org.apache.shardingsphere.sql.parser.statement.core.statement.SQLStatement;
54
55 import java.sql.ParameterMetaData;
56 import java.sql.PreparedStatement;
57 import java.sql.ResultSet;
58 import java.sql.SQLException;
59 import java.sql.Statement;
60 import java.util.ArrayList;
61 import java.util.Collection;
62 import java.util.Collections;
63 import java.util.LinkedList;
64 import java.util.List;
65 import java.util.Map;
66 import java.util.Optional;
67
68
69
70
71 @HighFrequencyInvocation
72 public final class ShardingSpherePreparedStatement extends AbstractPreparedStatementAdapter {
73
74 @Getter
75 private final ShardingSphereConnection connection;
76
77 private final ShardingSphereMetaData metaData;
78
79 private final String sql;
80
81 private final HintValueContext hintValueContext;
82
83 private final SQLStatementContext sqlStatementContext;
84
85 private final ShardingSphereDatabase usedDatabase;
86
87 private final StatementOption statementOption;
88
89 @Getter(AccessLevel.PROTECTED)
90 private final StatementManager statementManager;
91
92 @Getter
93 private final ParameterMetaData parameterMetaData;
94
95 private final DriverExecutorFacade driverExecutorFacade;
96
97 private final DriverExecuteBatchExecutor executeBatchExecutor;
98
99 private final List<PreparedStatement> statements = new ArrayList<>();
100
101 private final List<List<Object>> parameterSets = new ArrayList<>();
102
103 private final Collection<Comparable<?>> generatedValues = new LinkedList<>();
104
105 private final boolean statementsCacheable;
106
107 private Map<String, Integer> columnLabelAndIndexMap;
108
109 private ResultSet currentResultSet;
110
111 private ResultSet currentBatchGeneratedKeysResultSet;
112
113 public ShardingSpherePreparedStatement(final ShardingSphereConnection connection, final String sql) throws SQLException {
114 this(connection, sql, ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY, ResultSet.HOLD_CURSORS_OVER_COMMIT, false, null);
115 }
116
117 public ShardingSpherePreparedStatement(final ShardingSphereConnection connection, final String sql, final int resultSetType, final int resultSetConcurrency) throws SQLException {
118 this(connection, sql, resultSetType, resultSetConcurrency, ResultSet.HOLD_CURSORS_OVER_COMMIT, false, null);
119 }
120
121 public ShardingSpherePreparedStatement(final ShardingSphereConnection connection, final String sql, final int autoGeneratedKeys) throws SQLException {
122 this(connection, sql, ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY, ResultSet.HOLD_CURSORS_OVER_COMMIT, RETURN_GENERATED_KEYS == autoGeneratedKeys, null);
123 }
124
125 public ShardingSpherePreparedStatement(final ShardingSphereConnection connection, final String sql, final String[] columns) throws SQLException {
126 this(connection, sql, ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY, ResultSet.HOLD_CURSORS_OVER_COMMIT, true, columns);
127 }
128
129 public ShardingSpherePreparedStatement(final ShardingSphereConnection connection, final String sql, final int resultSetType, final int resultSetConcurrency,
130 final int resultSetHoldability) throws SQLException {
131 this(connection, sql, resultSetType, resultSetConcurrency, resultSetHoldability, false, null);
132 }
133
134 private ShardingSpherePreparedStatement(final ShardingSphereConnection connection, final String originSQL, final int resultSetType, final int resultSetConcurrency,
135 final int resultSetHoldability, final boolean returnGeneratedKeys, final String[] columns) throws SQLException {
136 ShardingSpherePreconditions.checkNotEmpty(originSQL, () -> new EmptySQLException().toSQLException());
137 this.connection = connection;
138 metaData = connection.getContextManager().getMetaDataContexts().getMetaData();
139 sql = SQLHintUtils.removeHint(originSQL);
140 hintValueContext = SQLHintUtils.extractHint(originSQL);
141 DatabaseType databaseType = metaData.getDatabase(connection.getCurrentDatabaseName()).getProtocolType();
142 SQLStatement sqlStatement = metaData.getGlobalRuleMetaData().getSingleRule(SQLParserRule.class).getSQLParserEngine(databaseType).parse(sql, true);
143 sqlStatementContext = new SQLBindEngine(metaData, connection.getCurrentDatabaseName(), hintValueContext).bind(databaseType, sqlStatement, Collections.emptyList());
144 String usedDatabaseName = sqlStatementContext.getTablesContext().getDatabaseName().orElse(connection.getCurrentDatabaseName());
145 connection.getDatabaseConnectionManager().getConnectionContext().setCurrentDatabaseName(connection.getCurrentDatabaseName());
146 usedDatabase = metaData.getDatabase(usedDatabaseName);
147 statementOption = returnGeneratedKeys ? new StatementOption(true, columns) : new StatementOption(resultSetType, resultSetConcurrency, resultSetHoldability);
148 statementManager = new StatementManager();
149 connection.getStatementManagers().add(statementManager);
150 parameterMetaData = new ShardingSphereParameterMetaData(sqlStatement);
151 driverExecutorFacade = new DriverExecutorFacade(connection, statementOption, statementManager, JDBCDriverType.PREPARED_STATEMENT);
152 executeBatchExecutor = new DriverExecuteBatchExecutor(connection, metaData, statementOption, statementManager, usedDatabase);
153 statementsCacheable = isStatementsCacheable();
154 }
155
156 private boolean isStatementsCacheable() {
157 return usedDatabase.getRuleMetaData().getAttributes(StorageConnectorReusableRuleAttribute.class).size() == usedDatabase.getRuleMetaData().getRules().size()
158 && !HintManager.isInstantiated();
159 }
160
161 @Override
162 public ResultSet executeQuery() throws SQLException {
163 try {
164 if (statementsCacheable && !statements.isEmpty()) {
165 resetParameters();
166 return statements.iterator().next().executeQuery();
167 }
168 clearPrevious();
169 QueryContext queryContext = createQueryContext();
170 handleAutoCommitBeforeExecution(queryContext.getSqlStatementContext().getSqlStatement(), connection);
171 findGeneratedKey().ifPresent(optional -> generatedValues.addAll(optional.getGeneratedValues()));
172 currentResultSet =
173 driverExecutorFacade.executeQuery(usedDatabase, metaData, queryContext, this, columnLabelAndIndexMap, (StatementAddCallback<PreparedStatement>) this::addStatements, this::replay);
174 if (currentResultSet instanceof ShardingSphereResultSet) {
175 columnLabelAndIndexMap = ((ShardingSphereResultSet) currentResultSet).getColumnLabelAndIndexMap();
176 }
177 return currentResultSet;
178
179 } catch (final RuntimeException | SQLException ex) {
180
181 handleExceptionInTransaction(connection, metaData);
182 throw SQLExceptionTransformEngine.toSQLException(ex, usedDatabase.getProtocolType());
183 } finally {
184 executeBatchExecutor.clear();
185 clearParameters();
186 }
187 }
188
189 private void addStatements(final Collection<PreparedStatement> statements, final Collection<List<Object>> parameterSets) {
190 this.statements.addAll(statements);
191 this.parameterSets.addAll(parameterSets);
192 }
193
194 private void resetParameters() throws SQLException {
195 replaySetParameter(statements, Collections.singletonList(getParameters()));
196 }
197
198 @Override
199 public int executeUpdate() throws SQLException {
200 try {
201 if (statementsCacheable && !statements.isEmpty()) {
202 resetParameters();
203 return statements.iterator().next().executeUpdate();
204 }
205 clearPrevious();
206 QueryContext queryContext = createQueryContext();
207 handleAutoCommitBeforeExecution(queryContext.getSqlStatementContext().getSqlStatement(), connection);
208 int result = driverExecutorFacade.executeUpdate(usedDatabase, metaData, queryContext,
209 (sql, statement) -> ((PreparedStatement) statement).executeUpdate(), (StatementAddCallback<PreparedStatement>) this::addStatements, this::replay);
210 findGeneratedKey().ifPresent(optional -> generatedValues.addAll(optional.getGeneratedValues()));
211 return result;
212
213 } catch (final RuntimeException | SQLException ex) {
214
215 handleExceptionInTransaction(connection, metaData);
216 throw SQLExceptionTransformEngine.toSQLException(ex, usedDatabase.getProtocolType());
217 } finally {
218 clearBatch();
219 }
220 }
221
222 @Override
223 public boolean execute() throws SQLException {
224 try {
225 if (statementsCacheable && !statements.isEmpty()) {
226 resetParameters();
227 return statements.iterator().next().execute();
228 }
229 clearPrevious();
230 QueryContext queryContext = createQueryContext();
231 handleAutoCommitBeforeExecution(queryContext.getSqlStatementContext().getSqlStatement(), connection);
232 boolean result = driverExecutorFacade.execute(usedDatabase, metaData, queryContext, (sql, statement) -> ((PreparedStatement) statement).execute(),
233 (StatementAddCallback<PreparedStatement>) this::addStatements, this::replay);
234 findGeneratedKey().ifPresent(optional -> generatedValues.addAll(optional.getGeneratedValues()));
235 return result;
236
237 } catch (final RuntimeException | SQLException ex) {
238
239 handleExceptionInTransaction(connection, metaData);
240 throw SQLExceptionTransformEngine.toSQLException(ex, usedDatabase.getProtocolType());
241 } finally {
242 clearBatch();
243 }
244 }
245
246 @Override
247 public ResultSet getResultSet() throws SQLException {
248 if (null != currentResultSet) {
249 return currentResultSet;
250 }
251 driverExecutorFacade.getResultSet(usedDatabase, sqlStatementContext, this, statements).ifPresent(optional -> currentResultSet = optional);
252 if (null == columnLabelAndIndexMap && currentResultSet instanceof ShardingSphereResultSet) {
253 columnLabelAndIndexMap = ((ShardingSphereResultSet) currentResultSet).getColumnLabelAndIndexMap();
254 }
255 return currentResultSet;
256 }
257
258 private QueryContext createQueryContext() {
259 List<Object> params = new ArrayList<>(getParameters());
260 if (sqlStatementContext instanceof ParameterAware) {
261 ((ParameterAware) sqlStatementContext).setUpParameters(params);
262 }
263 return new QueryContext(sqlStatementContext, sql, params, hintValueContext, connection.getDatabaseConnectionManager().getConnectionContext(), metaData, true);
264 }
265
266 private void replay() throws SQLException {
267 replaySetParameter(statements, parameterSets);
268 for (Statement each : statements) {
269 getMethodInvocationRecorder().replay(each);
270 }
271 }
272
273 private void replaySetParameter(final List<PreparedStatement> statements, final List<List<Object>> parameterSets) throws SQLException {
274 for (int i = 0; i < statements.size(); i++) {
275 replaySetParameter(statements.get(i), parameterSets.get(i));
276 }
277 }
278
279 private void clearPrevious() {
280 currentResultSet = null;
281 statements.clear();
282 parameterSets.clear();
283 generatedValues.clear();
284 }
285
286 private Optional<GeneratedKeyContext> findGeneratedKey() {
287 return sqlStatementContext instanceof InsertStatementContext ? ((InsertStatementContext) sqlStatementContext).getGeneratedKeyContext() : Optional.empty();
288 }
289
290 @Override
291 public ResultSet getGeneratedKeys() throws SQLException {
292 if (null != currentBatchGeneratedKeysResultSet) {
293 return currentBatchGeneratedKeysResultSet;
294 }
295 Optional<GeneratedKeyContext> generatedKey = findGeneratedKey();
296 if (generatedKey.isPresent() && statementOption.isReturnGeneratedKeys() && !generatedValues.isEmpty()) {
297 return new GeneratedKeysResultSet(getGeneratedKeysColumnName(generatedKey.get().getColumnName()), generatedValues.iterator(), this);
298 }
299 for (PreparedStatement each : statements) {
300 ResultSet resultSet = each.getGeneratedKeys();
301 while (resultSet.next()) {
302 generatedValues.add((Comparable<?>) resultSet.getObject(1));
303 }
304 }
305 String columnName = generatedKey.map(GeneratedKeyContext::getColumnName).orElse(null);
306 return new GeneratedKeysResultSet(getGeneratedKeysColumnName(columnName), generatedValues.iterator(), this);
307 }
308
309 private String getGeneratedKeysColumnName(final String columnName) {
310 return DatabaseTypedSPILoader.findService(GeneratedKeyColumnProvider.class, usedDatabase.getProtocolType())
311 .map(GeneratedKeyColumnProvider::getColumnName).orElse(columnName);
312 }
313
314 @Override
315 public void addBatch() {
316 currentResultSet = null;
317 QueryContext queryContext = createQueryContext();
318 executeBatchExecutor.addBatch(queryContext, usedDatabase);
319 findGeneratedKey().ifPresent(optional -> generatedValues.addAll(optional.getGeneratedValues()));
320 clearParameters();
321 }
322
323 @Override
324 public int[] executeBatch() throws SQLException {
325 try {
326 return executeBatchExecutor.executeBatch(usedDatabase, sqlStatementContext, generatedValues, statementOption,
327 (StatementAddCallback<PreparedStatement>) (statements, parameterSets) -> this.statements.addAll(statements),
328 this::replaySetParameter,
329 () -> {
330 currentBatchGeneratedKeysResultSet = getGeneratedKeys();
331 statements.clear();
332 });
333
334 } catch (final RuntimeException ex) {
335
336 handleExceptionInTransaction(connection, metaData);
337 throw SQLExceptionTransformEngine.toSQLException(ex, usedDatabase.getProtocolType());
338 } finally {
339 clearBatch();
340 }
341 }
342
343 @Override
344 public void clearBatch() {
345 currentResultSet = null;
346 executeBatchExecutor.clear();
347 clearParameters();
348 }
349
350 @SuppressWarnings("MagicConstant")
351 @Override
352 public int getResultSetType() {
353 return statementOption.getResultSetType();
354 }
355
356 @SuppressWarnings("MagicConstant")
357 @Override
358 public int getResultSetConcurrency() {
359 return statementOption.getResultSetConcurrency();
360 }
361
362 @Override
363 public int getResultSetHoldability() {
364 return statementOption.getResultSetHoldability();
365 }
366
367 @Override
368 public boolean isAccumulate() {
369 for (DataNodeRuleAttribute each : usedDatabase.getRuleMetaData().getAttributes(DataNodeRuleAttribute.class)) {
370 if (each.isNeedAccumulate(sqlStatementContext.getTablesContext().getTableNames())) {
371 return true;
372 }
373 }
374 return false;
375 }
376
377 @Override
378 public Collection<PreparedStatement> getRoutedStatements() {
379 return statements;
380 }
381
382 @Override
383 protected void closeExecutor() throws SQLException {
384 driverExecutorFacade.close();
385 }
386 }