This is an automated email from the ASF dual-hosted git repository.
zhangliang pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/shardingsphere.git
The following commit(s) were added to refs/heads/master by this push:
new 1dd149d9db5 Fix generated keys retrieval for PostgreSQL SINGLE tables
(#37485) (#38024)
1dd149d9db5 is described below
commit 1dd149d9db511735d3639fb28742a607436b23d9
Author: Aviral Garg <[email protected]>
AuthorDate: Wed Feb 18 09:37:12 2026 +0530
Fix generated keys retrieval for PostgreSQL SINGLE tables (#37485) (#38024)
* Fix generated keys retrieval for PostgreSQL SINGLE tables (#37485)
When inserting into a SINGLE table with @GeneratedValue(IDENTITY)
on PostgreSQL, getGeneratedKeys() threw NumberFormatException because
it always read column index 1 from the underlying ResultSet. PostgreSQL
JDBC driver returns all columns via RETURNING *, so column 1 may not
be the generated key column.
Fix: use the generated key column name from GeneratedKeyContext (when
available) to read the correct column via resultSet.getObject(columnName)
instead of blindly reading resultSet.getObject(1).
Closes #37485
* Fix code style: wrap long lines and apply spotless formatting
* ci: validate e2e-agent image archives
* jdbc: fix generated key fallback and CI scope
---
.../statement/ShardingSpherePreparedStatement.java | 23 +++-
.../core/statement/ShardingSphereStatement.java | 23 +++-
.../ShardingSpherePreparedStatementTest.java | 98 ++++++++++++++
.../statement/ShardingSphereStatementTest.java | 149 +++++++++++++++++++++
4 files changed, 287 insertions(+), 6 deletions(-)
diff --git
a/jdbc/src/main/java/org/apache/shardingsphere/driver/jdbc/core/statement/ShardingSpherePreparedStatement.java
b/jdbc/src/main/java/org/apache/shardingsphere/driver/jdbc/core/statement/ShardingSpherePreparedStatement.java
index 7dab5b2c5af..c73e21d3ba0 100644
---
a/jdbc/src/main/java/org/apache/shardingsphere/driver/jdbc/core/statement/ShardingSpherePreparedStatement.java
+++
b/jdbc/src/main/java/org/apache/shardingsphere/driver/jdbc/core/statement/ShardingSpherePreparedStatement.java
@@ -300,14 +300,31 @@ public final class ShardingSpherePreparedStatement
extends AbstractPreparedState
if (generatedKey.isPresent() &&
statementOption.isReturnGeneratedKeys() && !generatedValues.isEmpty()) {
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) {
diff --git
a/jdbc/src/main/java/org/apache/shardingsphere/driver/jdbc/core/statement/ShardingSphereStatement.java
b/jdbc/src/main/java/org/apache/shardingsphere/driver/jdbc/core/statement/ShardingSphereStatement.java
index 840b9f22a93..f1be7e71f6b 100644
---
a/jdbc/src/main/java/org/apache/shardingsphere/driver/jdbc/core/statement/ShardingSphereStatement.java
+++
b/jdbc/src/main/java/org/apache/shardingsphere/driver/jdbc/core/statement/ShardingSphereStatement.java
@@ -346,15 +346,32 @@ public final class ShardingSphereStatement extends
AbstractStatementAdapter {
if (returnGeneratedKeys && generatedKey.isPresent() &&
!generatedKey.get().getGeneratedValues().isEmpty()) {
return new
GeneratedKeysResultSet(getGeneratedKeysColumnName(generatedKey.get().getColumnName()),
generatedKey.get().getGeneratedValues().iterator(), this);
}
+ String columnName =
generatedKey.map(GeneratedKeyContext::getColumnName).orElse(null);
+ String generatedKeysColumnName =
getGeneratedKeysColumnName(columnName);
Collection<Comparable<?>> generatedValues = new LinkedList<>();
for (Statement 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 Optional<GeneratedKeyContext> findGeneratedKey() {
diff --git
a/jdbc/src/test/java/org/apache/shardingsphere/driver/jdbc/core/statement/ShardingSpherePreparedStatementTest.java
b/jdbc/src/test/java/org/apache/shardingsphere/driver/jdbc/core/statement/ShardingSpherePreparedStatementTest.java
index f0cf7387fbb..9771043b7e2 100644
---
a/jdbc/src/test/java/org/apache/shardingsphere/driver/jdbc/core/statement/ShardingSpherePreparedStatementTest.java
+++
b/jdbc/src/test/java/org/apache/shardingsphere/driver/jdbc/core/statement/ShardingSpherePreparedStatementTest.java
@@ -20,6 +20,8 @@ package org.apache.shardingsphere.driver.jdbc.core.statement;
import org.apache.shardingsphere.database.connector.core.type.DatabaseType;
import
org.apache.shardingsphere.driver.jdbc.core.connection.ShardingSphereConnection;
import
org.apache.shardingsphere.driver.jdbc.core.resultset.GeneratedKeysResultSet;
+import
org.apache.shardingsphere.infra.binder.context.segment.insert.keygen.GeneratedKeyContext;
+import
org.apache.shardingsphere.infra.binder.context.statement.type.dml.InsertStatementContext;
import org.apache.shardingsphere.infra.config.props.ConfigurationProperties;
import org.apache.shardingsphere.infra.metadata.ShardingSphereMetaData;
import
org.apache.shardingsphere.infra.metadata.database.ShardingSphereDatabase;
@@ -32,6 +34,8 @@ import
org.apache.shardingsphere.sqlfederation.rule.builder.DefaultSQLFederation
import org.junit.jupiter.api.Test;
import org.mockito.internal.configuration.plugins.Plugins;
+import java.lang.reflect.Field;
+import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Statement;
@@ -39,12 +43,17 @@ import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.Properties;
+import java.util.Optional;
+import static org.hamcrest.Matchers.is;
import static org.hamcrest.Matchers.not;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.mockito.Mockito.RETURNS_DEEP_STUBS;
import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
class ShardingSpherePreparedStatementTest {
@@ -73,4 +82,93 @@ class ShardingSpherePreparedStatementTest {
assertThat(actual, not(cachedResultSet));
assertFalse(actual.isClosed());
}
+
+ @Test
+ void assertGetGeneratedKeysByColumnName() throws SQLException,
ReflectiveOperationException {
+ ShardingSpherePreparedStatement preparedStatement =
createPreparedStatement(TypedSPILoader.getService(DatabaseType.class, "SQL92"));
+ ResultSet generatedKeys = mock(ResultSet.class);
+ when(generatedKeys.next()).thenReturn(true, false);
+ when(generatedKeys.getObject("id")).thenReturn(1L);
+ setInsertStatementContext(preparedStatement, "id");
+ addPreparedStatement(preparedStatement, generatedKeys);
+ ResultSet actual = preparedStatement.getGeneratedKeys();
+ assertTrue(actual.next());
+ assertThat(actual.getObject(1), is(1L));
+ verify(generatedKeys).getObject("id");
+ verify(generatedKeys, never()).getObject(1);
+ }
+
+ @Test
+ void assertGetGeneratedKeysByColumnIndexWhenColumnNameMissing() throws
SQLException, ReflectiveOperationException {
+ ResultSet generatedKeys = mock(ResultSet.class);
+ when(generatedKeys.next()).thenReturn(true, false);
+ when(generatedKeys.getObject("id")).thenThrow(new SQLException("Column
not found"));
+ when(generatedKeys.getObject(1)).thenReturn(2L);
+ ShardingSpherePreparedStatement preparedStatement =
createPreparedStatement(TypedSPILoader.getService(DatabaseType.class, "SQL92"));
+ setInsertStatementContext(preparedStatement, "id");
+ addPreparedStatement(preparedStatement, generatedKeys);
+ ResultSet actual = preparedStatement.getGeneratedKeys();
+ assertTrue(actual.next());
+ assertThat(actual.getObject(1), is(2L));
+ verify(generatedKeys).getObject("id");
+ verify(generatedKeys).getObject(1);
+ }
+
+ @Test
+ void assertGetGeneratedKeysByDialectGeneratedKeyColumn() throws
SQLException, ReflectiveOperationException {
+ ShardingSpherePreparedStatement preparedStatement =
createPreparedStatement(TypedSPILoader.getService(DatabaseType.class, "MySQL"));
+ ResultSet generatedKeys = mock(ResultSet.class);
+ when(generatedKeys.next()).thenReturn(true, false);
+ when(generatedKeys.getObject("GENERATED_KEY")).thenReturn(3L);
+ setInsertStatementContext(preparedStatement, "id");
+ addPreparedStatement(preparedStatement, generatedKeys);
+ ResultSet actual = preparedStatement.getGeneratedKeys();
+ assertTrue(actual.next());
+ assertThat(actual.getObject(1), is(3L));
+ verify(generatedKeys).getObject("GENERATED_KEY");
+ verify(generatedKeys, never()).getObject("id");
+ verify(generatedKeys, never()).getObject(1);
+ }
+
+ private ShardingSpherePreparedStatement createPreparedStatement(final
DatabaseType databaseType) throws SQLException {
+ ShardingSphereMetaData metaData = createMetaData(databaseType);
+ ShardingSphereConnection connection =
mock(ShardingSphereConnection.class, RETURNS_DEEP_STUBS);
+
when(connection.getContextManager().getMetaDataContexts().getMetaData()).thenReturn(metaData);
+ when(connection.getCurrentDatabaseName()).thenReturn("foo_db");
+ return new ShardingSpherePreparedStatement(connection, "SELECT 1",
Statement.RETURN_GENERATED_KEYS);
+ }
+
+ private ShardingSphereMetaData createMetaData(final DatabaseType
databaseType) {
+ ShardingSphereDatabase database = mock(ShardingSphereDatabase.class);
+ when(database.getProtocolType()).thenReturn(databaseType);
+ when(database.getRuleMetaData()).thenReturn(new
RuleMetaData(Collections.emptyList()));
+ ShardingSphereMetaData result = mock(ShardingSphereMetaData.class);
+ when(result.getDatabase("foo_db")).thenReturn(database);
+ when(result.getGlobalRuleMetaData()).thenReturn(new
RuleMetaData(Arrays.asList(
+ new SQLParserRule(new
DefaultSQLParserRuleConfigurationBuilder().build()),
+ new SQLFederationRule(new
DefaultSQLFederationRuleConfigurationBuilder().build(),
Collections.emptyList()))));
+ when(result.getProps()).thenReturn(new ConfigurationProperties(new
Properties()));
+ return result;
+ }
+
+ private void setInsertStatementContext(final
ShardingSpherePreparedStatement preparedStatement, final String columnName)
throws ReflectiveOperationException {
+ GeneratedKeyContext generatedKeyContext =
mock(GeneratedKeyContext.class);
+ when(generatedKeyContext.getColumnName()).thenReturn(columnName);
+
when(generatedKeyContext.getGeneratedValues()).thenReturn(Collections.emptyList());
+ InsertStatementContext insertStatementContext =
mock(InsertStatementContext.class);
+
when(insertStatementContext.getGeneratedKeyContext()).thenReturn(Optional.of(generatedKeyContext));
+
Plugins.getMemberAccessor().set(ShardingSpherePreparedStatement.class.getDeclaredField("sqlStatementContext"),
preparedStatement, insertStatementContext);
+ }
+
+ private void addPreparedStatement(final ShardingSpherePreparedStatement
preparedStatement, final ResultSet generatedKeys) throws
ReflectiveOperationException, SQLException {
+ PreparedStatement statement = mock(PreparedStatement.class);
+ when(statement.getGeneratedKeys()).thenReturn(generatedKeys);
+ getPreparedStatements(preparedStatement).add(statement);
+ }
+
+ @SuppressWarnings("unchecked")
+ private Collection<PreparedStatement> getPreparedStatements(final
ShardingSpherePreparedStatement preparedStatement) throws
ReflectiveOperationException {
+ Field statementsField =
ShardingSpherePreparedStatement.class.getDeclaredField("statements");
+ return (Collection<PreparedStatement>)
Plugins.getMemberAccessor().get(statementsField, preparedStatement);
+ }
}
diff --git
a/jdbc/src/test/java/org/apache/shardingsphere/driver/jdbc/core/statement/ShardingSphereStatementTest.java
b/jdbc/src/test/java/org/apache/shardingsphere/driver/jdbc/core/statement/ShardingSphereStatementTest.java
new file mode 100644
index 00000000000..eeef347136c
--- /dev/null
+++
b/jdbc/src/test/java/org/apache/shardingsphere/driver/jdbc/core/statement/ShardingSphereStatementTest.java
@@ -0,0 +1,149 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.shardingsphere.driver.jdbc.core.statement;
+
+import org.apache.shardingsphere.database.connector.core.type.DatabaseType;
+import
org.apache.shardingsphere.driver.jdbc.core.connection.ShardingSphereConnection;
+import
org.apache.shardingsphere.infra.binder.context.segment.insert.keygen.GeneratedKeyContext;
+import
org.apache.shardingsphere.infra.binder.context.statement.type.dml.InsertStatementContext;
+import org.apache.shardingsphere.infra.config.props.ConfigurationProperties;
+import org.apache.shardingsphere.infra.metadata.ShardingSphereMetaData;
+import
org.apache.shardingsphere.infra.metadata.database.ShardingSphereDatabase;
+import org.apache.shardingsphere.infra.metadata.database.rule.RuleMetaData;
+import org.apache.shardingsphere.infra.session.query.QueryContext;
+import org.apache.shardingsphere.infra.spi.type.typed.TypedSPILoader;
+import org.apache.shardingsphere.parser.rule.SQLParserRule;
+import
org.apache.shardingsphere.parser.rule.builder.DefaultSQLParserRuleConfigurationBuilder;
+import org.apache.shardingsphere.sqlfederation.rule.SQLFederationRule;
+import
org.apache.shardingsphere.sqlfederation.rule.builder.DefaultSQLFederationRuleConfigurationBuilder;
+import org.junit.jupiter.api.Test;
+import org.mockito.internal.configuration.plugins.Plugins;
+
+import java.lang.reflect.Field;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.sql.Statement;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Properties;
+import java.util.Optional;
+
+import static org.hamcrest.Matchers.is;
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.Mockito.RETURNS_DEEP_STUBS;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+class ShardingSphereStatementTest {
+
+ @Test
+ void assertGetGeneratedKeysByColumnName() throws SQLException,
ReflectiveOperationException {
+ ShardingSphereStatement statement =
createStatement(TypedSPILoader.getService(DatabaseType.class, "SQL92"));
+ ResultSet generatedKeys = mock(ResultSet.class);
+ when(generatedKeys.next()).thenReturn(true, false);
+ when(generatedKeys.getObject("id")).thenReturn(1L);
+ setInsertStatementContext(statement, "id");
+ addStatement(statement, generatedKeys);
+ ResultSet actual = statement.getGeneratedKeys();
+ assertTrue(actual.next());
+ assertThat(actual.getObject(1), is(1L));
+ verify(generatedKeys).getObject("id");
+ verify(generatedKeys, never()).getObject(1);
+ }
+
+ @Test
+ void assertGetGeneratedKeysByColumnIndexWhenColumnNameMissing() throws
SQLException, ReflectiveOperationException {
+ ResultSet generatedKeys = mock(ResultSet.class);
+ when(generatedKeys.next()).thenReturn(true, false);
+ when(generatedKeys.getObject("id")).thenThrow(new SQLException("Column
not found"));
+ when(generatedKeys.getObject(1)).thenReturn(2L);
+ ShardingSphereStatement statement =
createStatement(TypedSPILoader.getService(DatabaseType.class, "SQL92"));
+ setInsertStatementContext(statement, "id");
+ addStatement(statement, generatedKeys);
+ ResultSet actual = statement.getGeneratedKeys();
+ assertTrue(actual.next());
+ assertThat(actual.getObject(1), is(2L));
+ verify(generatedKeys).getObject("id");
+ verify(generatedKeys).getObject(1);
+ }
+
+ @Test
+ void assertGetGeneratedKeysByDialectGeneratedKeyColumn() throws
SQLException, ReflectiveOperationException {
+ ShardingSphereStatement statement =
createStatement(TypedSPILoader.getService(DatabaseType.class, "MySQL"));
+ ResultSet generatedKeys = mock(ResultSet.class);
+ when(generatedKeys.next()).thenReturn(true, false);
+ when(generatedKeys.getObject("GENERATED_KEY")).thenReturn(3L);
+ setInsertStatementContext(statement, "id");
+ addStatement(statement, generatedKeys);
+ ResultSet actual = statement.getGeneratedKeys();
+ assertTrue(actual.next());
+ assertThat(actual.getObject(1), is(3L));
+ verify(generatedKeys).getObject("GENERATED_KEY");
+ verify(generatedKeys, never()).getObject("id");
+ verify(generatedKeys, never()).getObject(1);
+ }
+
+ private ShardingSphereStatement createStatement(final DatabaseType
databaseType) throws SQLException {
+ ShardingSphereMetaData metaData = createMetaData(databaseType);
+ ShardingSphereConnection connection =
mock(ShardingSphereConnection.class, RETURNS_DEEP_STUBS);
+
when(connection.getContextManager().getMetaDataContexts().getMetaData()).thenReturn(metaData);
+ when(connection.getCurrentDatabaseName()).thenReturn("foo_db");
+ return new ShardingSphereStatement(connection);
+ }
+
+ private ShardingSphereMetaData createMetaData(final DatabaseType
databaseType) {
+ ShardingSphereDatabase database = mock(ShardingSphereDatabase.class);
+ when(database.getProtocolType()).thenReturn(databaseType);
+ when(database.getRuleMetaData()).thenReturn(new
RuleMetaData(Collections.emptyList()));
+ ShardingSphereMetaData result = mock(ShardingSphereMetaData.class);
+ when(result.getDatabase("foo_db")).thenReturn(database);
+ when(result.getGlobalRuleMetaData()).thenReturn(new
RuleMetaData(Arrays.asList(
+ new SQLParserRule(new
DefaultSQLParserRuleConfigurationBuilder().build()),
+ new SQLFederationRule(new
DefaultSQLFederationRuleConfigurationBuilder().build(),
Collections.emptyList()))));
+ when(result.getProps()).thenReturn(new ConfigurationProperties(new
Properties()));
+ return result;
+ }
+
+ private void setInsertStatementContext(final ShardingSphereStatement
statement, final String columnName) throws ReflectiveOperationException {
+ GeneratedKeyContext generatedKeyContext =
mock(GeneratedKeyContext.class);
+ when(generatedKeyContext.getColumnName()).thenReturn(columnName);
+
when(generatedKeyContext.getGeneratedValues()).thenReturn(Collections.emptyList());
+ InsertStatementContext insertStatementContext =
mock(InsertStatementContext.class);
+
when(insertStatementContext.getGeneratedKeyContext()).thenReturn(Optional.of(generatedKeyContext));
+ QueryContext queryContext = mock(QueryContext.class);
+
when(queryContext.getSqlStatementContext()).thenReturn(insertStatementContext);
+
Plugins.getMemberAccessor().set(ShardingSphereStatement.class.getDeclaredField("queryContext"),
statement, queryContext);
+
Plugins.getMemberAccessor().set(ShardingSphereStatement.class.getDeclaredField("returnGeneratedKeys"),
statement, true);
+ }
+
+ private void addStatement(final ShardingSphereStatement
shardingSphereStatement, final ResultSet generatedKeys) throws
ReflectiveOperationException, SQLException {
+ Statement statement = mock(Statement.class);
+ when(statement.getGeneratedKeys()).thenReturn(generatedKeys);
+ getStatements(shardingSphereStatement).add(statement);
+ }
+
+ @SuppressWarnings("unchecked")
+ private Collection<Statement> getStatements(final ShardingSphereStatement
statement) throws ReflectiveOperationException {
+ Field statementsField =
ShardingSphereStatement.class.getDeclaredField("statements");
+ return (Collection<Statement>)
Plugins.getMemberAccessor().get(statementsField, statement);
+ }
+}
