Skip to content

Fix generated keys retrieval for PostgreSQL SINGLE tables (#37485) #38024

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
aviralgarg05 wants to merge 4 commits into
base: master
Choose a base branch
from
+410 −78
Open

Fix generated keys retrieval for PostgreSQL SINGLE tables (#37485) #38024

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
6d28e8d
Fix generated keys retrieval for PostgreSQL SINGLE tables (#37485)
aviralgarg05 Feb 12, 2026
632d613
Fix code style: wrap long lines and apply spotless formatting
aviralgarg05 Feb 12, 2026
bd4ecca
ci: validate e2e-agent image archives
aviralgarg05 Feb 12, 2026
82382bd
jdbc: fix generated key fallback and CI scope
aviralgarg05 Feb 15, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
113 changes: 80 additions & 33 deletions ...org/apache/shardingsphere/driver/jdbc/core/statement/ShardingSpherePreparedStatement.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -111,51 +111,69 @@ public final class ShardingSpherePreparedStatement extends AbstractPreparedState

private QueryContext queryContext;

public ShardingSpherePreparedStatement(final ShardingSphereConnection connection, final String sql) throws SQLException {
this(connection, sql, ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY, ResultSet.HOLD_CURSORS_OVER_COMMIT, false, null);
public ShardingSpherePreparedStatement(final ShardingSphereConnection connection, final String sql)
throws SQLException {
this(connection, sql, ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY,
ResultSet.HOLD_CURSORS_OVER_COMMIT, false, null);
}

public ShardingSpherePreparedStatement(final ShardingSphereConnection connection, final String sql, final int resultSetType, final int resultSetConcurrency) throws SQLException {
public ShardingSpherePreparedStatement(final ShardingSphereConnection connection, final String sql,
final int resultSetType, final int resultSetConcurrency) throws SQLException {
this(connection, sql, resultSetType, resultSetConcurrency, ResultSet.HOLD_CURSORS_OVER_COMMIT, false, null);
}

public ShardingSpherePreparedStatement(final ShardingSphereConnection connection, final String sql, final int autoGeneratedKeys) throws SQLException {
this(connection, sql, ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY, ResultSet.HOLD_CURSORS_OVER_COMMIT, RETURN_GENERATED_KEYS == autoGeneratedKeys, null);
public ShardingSpherePreparedStatement(final ShardingSphereConnection connection, final String sql,
final int autoGeneratedKeys) throws SQLException {
this(connection, sql, ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY,
ResultSet.HOLD_CURSORS_OVER_COMMIT, RETURN_GENERATED_KEYS == autoGeneratedKeys, null);
}

public ShardingSpherePreparedStatement(final ShardingSphereConnection connection, final String sql, final String[] columns) throws SQLException {
this(connection, sql, ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY, ResultSet.HOLD_CURSORS_OVER_COMMIT, true, columns);
public ShardingSpherePreparedStatement(final ShardingSphereConnection connection, final String sql,
final String[] columns) throws SQLException {
this(connection, sql, ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY,
ResultSet.HOLD_CURSORS_OVER_COMMIT, true, columns);
}

public ShardingSpherePreparedStatement(final ShardingSphereConnection connection, final String sql, final int resultSetType, final int resultSetConcurrency,
public ShardingSpherePreparedStatement(final ShardingSphereConnection connection, final String sql,
final int resultSetType, final int resultSetConcurrency,
final int resultSetHoldability) throws SQLException {
this(connection, sql, resultSetType, resultSetConcurrency, resultSetHoldability, false, null);
}

private ShardingSpherePreparedStatement(final ShardingSphereConnection connection, final String originSQL, final int resultSetType, final int resultSetConcurrency,
final int resultSetHoldability, final boolean returnGeneratedKeys, final String[] columns) throws SQLException {
private ShardingSpherePreparedStatement(final ShardingSphereConnection connection, final String originSQL,
final int resultSetType, final int resultSetConcurrency,
final int resultSetHoldability, final boolean returnGeneratedKeys, final String[] columns)
throws SQLException {
ShardingSpherePreconditions.checkNotEmpty(originSQL, () -> new EmptySQLException().toSQLException());
this.connection = connection;
metaData = connection.getContextManager().getMetaDataContexts().getMetaData();
sql = SQLHintUtils.removeHint(originSQL);
hintValueContext = SQLHintUtils.extractHint(originSQL);
ShardingSphereDatabase currentDatabase = metaData.getDatabase(connection.getCurrentDatabaseName());
SQLStatement sqlStatement = metaData.getGlobalRuleMetaData().getSingleRule(SQLParserRule.class).getSQLParserEngine(currentDatabase.getProtocolType()).parse(sql, true);
sqlStatementContext = new SQLBindEngine(metaData, connection.getCurrentDatabaseName(), hintValueContext).bind(sqlStatement);
String usedDatabaseName = sqlStatementContext.getTablesContext().getDatabaseName().orElse(connection.getCurrentDatabaseName());
connection.getDatabaseConnectionManager().getConnectionContext().setCurrentDatabaseName(connection.getCurrentDatabaseName());
SQLStatement sqlStatement = metaData.getGlobalRuleMetaData().getSingleRule(SQLParserRule.class)
.getSQLParserEngine(currentDatabase.getProtocolType()).parse(sql, true);
sqlStatementContext = new SQLBindEngine(metaData, connection.getCurrentDatabaseName(), hintValueContext)
.bind(sqlStatement);
String usedDatabaseName = sqlStatementContext.getTablesContext().getDatabaseName()
.orElse(connection.getCurrentDatabaseName());
connection.getDatabaseConnectionManager().getConnectionContext()
.setCurrentDatabaseName(connection.getCurrentDatabaseName());
usedDatabase = metaData.getDatabase(usedDatabaseName);
statementOption = returnGeneratedKeys ? new StatementOption(true, columns) : new StatementOption(resultSetType, resultSetConcurrency, resultSetHoldability);
statementOption = returnGeneratedKeys ? new StatementOption(true, columns)
: new StatementOption(resultSetType, resultSetConcurrency, resultSetHoldability);
statementManager = new StatementManager();
connection.getStatementManagers().add(statementManager);
parameterMetaData = new ShardingSphereParameterMetaData(sqlStatement);
driverExecutorFacade = new DriverExecutorFacade(connection, statementOption, statementManager, JDBCDriverType.PREPARED_STATEMENT, currentDatabase);
executeBatchExecutor = new DriverExecuteBatchExecutor(connection, metaData, statementOption, statementManager, usedDatabase);
driverExecutorFacade = new DriverExecutorFacade(connection, statementOption, statementManager,
JDBCDriverType.PREPARED_STATEMENT, currentDatabase);
executeBatchExecutor = new DriverExecuteBatchExecutor(connection, metaData, statementOption, statementManager,
usedDatabase);
statementsCacheable = isStatementsCacheable();
}

private boolean isStatementsCacheable() {
return usedDatabase.getRuleMetaData().getAttributes(StorageConnectorReusableRuleAttribute.class).size() == usedDatabase.getRuleMetaData().getRules().size()
return usedDatabase.getRuleMetaData().getAttributes(StorageConnectorReusableRuleAttribute.class)
.size() == usedDatabase.getRuleMetaData().getRules().size()
&& !HintManager.isInstantiated();
}

Expand All @@ -171,8 +189,9 @@ public ResultSet executeQuery() throws SQLException {
this.queryContext = queryContext;
handleAutoCommitBeforeExecution(queryContext.getSqlStatementContext().getSqlStatement(), connection);
findGeneratedKey().ifPresent(optional -> generatedValues.addAll(optional.getGeneratedValues()));
currentResultSet =
driverExecutorFacade.executeQuery(usedDatabase, metaData, queryContext, this, columnLabelAndIndexMap, (StatementAddCallback<PreparedStatement>) this::addStatements, this::replay);
currentResultSet = driverExecutorFacade.executeQuery(usedDatabase, metaData, queryContext, this,
columnLabelAndIndexMap, (StatementAddCallback<PreparedStatement>) this::addStatements,
this::replay);
if (currentResultSet instanceof ShardingSphereResultSet) {
columnLabelAndIndexMap = ((ShardingSphereResultSet) currentResultSet).getColumnLabelAndIndexMap();
}
Expand All @@ -188,7 +207,8 @@ public ResultSet executeQuery() throws SQLException {
}
}

private void addStatements(final Collection<PreparedStatement> statements, final Collection<List<Object>> parameterSets) {
private void addStatements(final Collection<PreparedStatement> statements,
final Collection<List<Object>> parameterSets) {
this.statements.addAll(statements);
this.parameterSets.addAll(parameterSets);
}
Expand All @@ -209,7 +229,8 @@ public int executeUpdate() throws SQLException {
this.queryContext = queryContext;
handleAutoCommitBeforeExecution(queryContext.getSqlStatementContext().getSqlStatement(), connection);
int result = driverExecutorFacade.executeUpdate(usedDatabase, metaData, queryContext,
(sql, statement) -> ((PreparedStatement) statement).executeUpdate(), (StatementAddCallback<PreparedStatement>) this::addStatements, this::replay);
(sql, statement) -> ((PreparedStatement) statement).executeUpdate(),
(StatementAddCallback<PreparedStatement>) this::addStatements, this::replay);
findGeneratedKey().ifPresent(optional -> generatedValues.addAll(optional.getGeneratedValues()));
return result;
// CHECKSTYLE:OFF
Expand All @@ -233,7 +254,8 @@ public boolean execute() throws SQLException {
QueryContext queryContext = createQueryContext();
this.queryContext = queryContext;
handleAutoCommitBeforeExecution(queryContext.getSqlStatementContext().getSqlStatement(), connection);
boolean result = driverExecutorFacade.execute(usedDatabase, metaData, queryContext, (sql, statement) -> ((PreparedStatement) statement).execute(),
boolean result = driverExecutorFacade.execute(usedDatabase, metaData, queryContext,
(sql, statement) -> ((PreparedStatement) statement).execute(),
(StatementAddCallback<PreparedStatement>) this::addStatements, this::replay);
findGeneratedKey().ifPresent(optional -> generatedValues.addAll(optional.getGeneratedValues()));
return result;
Expand All @@ -252,7 +274,8 @@ public ResultSet getResultSet() throws SQLException {
if (null != currentResultSet) {
return currentResultSet;
}
driverExecutorFacade.getResultSet(usedDatabase, queryContext, this, statements).ifPresent(optional -> currentResultSet = optional);
driverExecutorFacade.getResultSet(usedDatabase, queryContext, this, statements)
.ifPresent(optional -> currentResultSet = optional);
if (null == columnLabelAndIndexMap && currentResultSet instanceof ShardingSphereResultSet) {
columnLabelAndIndexMap = ((ShardingSphereResultSet) currentResultSet).getColumnLabelAndIndexMap();
}
Expand All @@ -264,7 +287,8 @@ private QueryContext createQueryContext() {
if (sqlStatementContext instanceof ParameterAware) {
((ParameterAware) sqlStatementContext).bindParameters(params);
}
return new QueryContext(sqlStatementContext, sql, params, hintValueContext, connection.getDatabaseConnectionManager().getConnectionContext(), metaData, true);
return new QueryContext(sqlStatementContext, sql, params, hintValueContext,
connection.getDatabaseConnectionManager().getConnectionContext(), metaData, true);
}

private void replay() throws SQLException {
Expand All @@ -288,7 +312,9 @@ private void clearPrevious() {
}

private Optional<GeneratedKeyContext> findGeneratedKey() {
return sqlStatementContext instanceof InsertStatementContext ? ((InsertStatementContext) sqlStatementContext).getGeneratedKeyContext() : Optional.empty();
return sqlStatementContext instanceof InsertStatementContext
? ((InsertStatementContext) sqlStatementContext).getGeneratedKeyContext()
: Optional.empty();
}

@Override
Expand All @@ -298,20 +324,39 @@ public ResultSet getGeneratedKeys() throws SQLException {
}
Optional<GeneratedKeyContext> generatedKey = findGeneratedKey();
if (generatedKey.isPresent() && statementOption.isReturnGeneratedKeys() && !generatedValues.isEmpty()) {
return new GeneratedKeysResultSet(getGeneratedKeysColumnName(generatedKey.get().getColumnName()), generatedValues.iterator(), this);
return new GeneratedKeysResultSet(getGeneratedKeysColumnName(generatedKey.get().getColumnName()),
generatedValues.iterator(), this);
}
String columnName = generatedKey.map(GeneratedKeyContext::getColumnName).orElse(null);
String generatedKeysColumnName = getGeneratedKeysColumnName(columnName);
for (PreparedStatement each : statements) {
ResultSet resultSet = each.getGeneratedKeys();
while (resultSet.next()) {
generatedValues.add((Comparable<?>) resultSet.getObject(1));
generatedValues.add(getGeneratedValue(resultSet, generatedKeysColumnName, columnName));
}
}
String columnName = generatedKey.map(GeneratedKeyContext::getColumnName).orElse(null);
return new GeneratedKeysResultSet(getGeneratedKeysColumnName(columnName), generatedValues.iterator(), this);
return new GeneratedKeysResultSet(generatedKeysColumnName, generatedValues.iterator(), this);
}

private Comparable<?> getGeneratedValue(final ResultSet resultSet, final String generatedKeysColumnName, final String columnName) throws SQLException {
if (null != generatedKeysColumnName) {
try {
return (Comparable<?>) resultSet.getObject(generatedKeysColumnName);
} catch (final SQLException ignored) {
}
}
if (null != columnName && !columnName.equals(generatedKeysColumnName)) {
try {
return (Comparable<?>) resultSet.getObject(columnName);
} catch (final SQLException ignored) {
}
}
return (Comparable<?>) resultSet.getObject(1);
}

private String getGeneratedKeysColumnName(final String columnName) {
Optional<DialectGeneratedKeyOption> generatedKeyOption = new DatabaseTypeRegistry(usedDatabase.getProtocolType()).getDialectDatabaseMetaData().getGeneratedKeyOption();
Optional<DialectGeneratedKeyOption> generatedKeyOption = new DatabaseTypeRegistry(
usedDatabase.getProtocolType()).getDialectDatabaseMetaData().getGeneratedKeyOption();
return generatedKeyOption.isPresent() ? generatedKeyOption.get().getColumnName() : columnName;
}

Expand All @@ -328,8 +373,10 @@ public void addBatch() {
@Override
public int[] executeBatch() throws SQLException {
try {
return executeBatchExecutor.executeBatch(usedDatabase, sqlStatementContext, generatedValues, statementOption,
(StatementAddCallback<PreparedStatement>) (statements, parameterSets) -> this.statements.addAll(statements),
return executeBatchExecutor.executeBatch(usedDatabase, sqlStatementContext, generatedValues,
statementOption,
(StatementAddCallback<PreparedStatement>) (statements, parameterSets) -> this.statements
.addAll(statements),
this::replaySetParameter,
() -> {
currentBatchGeneratedKeysResultSet = getGeneratedKeys();
Expand Down
Loading