Skip to content

Commit 0c34532

Browse files
olavloiteolavloite
authored
feat: return generated keys (googleapis#1310)
Adds support for returning generated keys.
1 parent ad15e89 commit 0c34532

9 files changed

Lines changed: 888 additions & 153 deletions

File tree

‎java-spanner-jdbc/src/main/java/com/google/cloud/spanner/jdbc/AbstractJdbcStatement.java‎

Lines changed: 0 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,6 @@
2424
import com.google.cloud.spanner.connection.Connection;
2525
import com.google.cloud.spanner.connection.StatementResult;
2626
import com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType;
27-
import com.google.rpc.Code;
2827
import java.sql.ResultSet;
2928
import java.sql.SQLException;
3029
import java.sql.SQLWarning;
@@ -201,43 +200,6 @@ private ResultSet executeQuery(
201200
}
202201
}
203202

204-
/**
205-
* Executes a SQL statement on the connection of this {@link Statement} as an update (DML)
206-
* statement.
207-
*
208-
* @param statement The SQL statement to execute
209-
* @return the number of rows that was inserted/updated/deleted
210-
* @throws SQLException if a database error occurs, or if the number of rows affected is larger
211-
* than {@link Integer#MAX_VALUE}
212-
*/
213-
int executeUpdate(com.google.cloud.spanner.Statement statement) throws SQLException {
214-
long count = executeLargeUpdate(statement);
215-
if (count > Integer.MAX_VALUE) {
216-
throw JdbcSqlExceptionFactory.of(
217-
"update count too large for executeUpdate: " + count, Code.OUT_OF_RANGE);
218-
}
219-
return (int) count;
220-
}
221-
222-
/**
223-
* Executes a SQL statement on the connection of this {@link Statement} as an update (DML)
224-
* statement.
225-
*
226-
* @param statement The SQL statement to execute
227-
* @return the number of rows that was inserted/updated/deleted
228-
* @throws SQLException if a database error occurs
229-
*/
230-
long executeLargeUpdate(com.google.cloud.spanner.Statement statement) throws SQLException {
231-
StatementTimeout originalTimeout = setTemporaryStatementTimeout();
232-
try {
233-
return connection.getSpannerConnection().executeUpdate(statement);
234-
} catch (SpannerException e) {
235-
throw JdbcSqlExceptionFactory.of(e);
236-
} finally {
237-
resetStatementTimeout(originalTimeout);
238-
}
239-
}
240-
241203
/**
242204
* Executes a SQL statement on the connection of this {@link Statement}. The SQL statement can be
243205
* any supported SQL statement, including client side statements such as SET AUTOCOMMIT ON|OFF.

‎java-spanner-jdbc/src/main/java/com/google/cloud/spanner/jdbc/JdbcConnection.java‎

Lines changed: 24 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,9 @@
1616

1717
package com.google.cloud.spanner.jdbc;
1818

19+
import static com.google.cloud.spanner.jdbc.JdbcStatement.ALL_COLUMNS;
20+
import static com.google.cloud.spanner.jdbc.JdbcStatement.isNullOrEmpty;
21+
1922
import com.google.api.client.util.Preconditions;
2023
import com.google.cloud.spanner.CommitResponse;
2124
import com.google.cloud.spanner.Mutation;
@@ -25,6 +28,7 @@
2528
import com.google.cloud.spanner.connection.ConnectionOptions;
2629
import com.google.cloud.spanner.connection.SavepointSupport;
2730
import com.google.cloud.spanner.connection.TransactionMode;
31+
import com.google.common.collect.ImmutableList;
2832
import com.google.common.collect.Iterators;
2933
import java.sql.Array;
3034
import java.sql.Blob;
@@ -49,9 +53,10 @@ class JdbcConnection extends AbstractJdbcConnection {
4953
"Only result sets with concurrency CONCUR_READ_ONLY are supported";
5054
private static final String ONLY_CLOSE_CURSORS_AT_COMMIT =
5155
"Only result sets with holdability CLOSE_CURSORS_AT_COMMIT are supported";
52-
static final String ONLY_NO_GENERATED_KEYS = "Only NO_GENERATED_KEYS are supported";
5356
static final String IS_VALID_QUERY = "SELECT 1";
5457

58+
static final ImmutableList<String> NO_GENERATED_KEY_COLUMNS = ImmutableList.of();
59+
5560
private Map<String, Class<?>> typeMap = new HashMap<>();
5661

5762
JdbcConnection(String connectionUrl, ConnectionOptions options) throws SQLException {
@@ -66,8 +71,13 @@ public Statement createStatement() throws SQLException {
6671

6772
@Override
6873
public JdbcPreparedStatement prepareStatement(String sql) throws SQLException {
74+
return prepareStatement(sql, NO_GENERATED_KEY_COLUMNS);
75+
}
76+
77+
private JdbcPreparedStatement prepareStatement(
78+
String sql, ImmutableList<String> generatedKeyColumns) throws SQLException {
6979
checkClosed();
70-
return new JdbcPreparedStatement(this, sql);
80+
return new JdbcPreparedStatement(this, sql, generatedKeyColumns);
7181
}
7282

7383
@Override
@@ -299,22 +309,26 @@ public PreparedStatement prepareStatement(
299309

300310
@Override
301311
public PreparedStatement prepareStatement(String sql, int autoGeneratedKeys) throws SQLException {
302-
checkClosed();
303-
JdbcPreconditions.checkSqlFeatureSupported(
304-
autoGeneratedKeys == Statement.NO_GENERATED_KEYS, ONLY_NO_GENERATED_KEYS);
305-
return prepareStatement(sql);
312+
return prepareStatement(
313+
sql,
314+
autoGeneratedKeys == Statement.RETURN_GENERATED_KEYS
315+
? ALL_COLUMNS
316+
: NO_GENERATED_KEY_COLUMNS);
306317
}
307318

308319
@Override
309320
public PreparedStatement prepareStatement(String sql, int[] columnIndexes) throws SQLException {
310-
checkClosed();
311-
return prepareStatement(sql);
321+
// This should preferably have returned an error, but the initial version of the driver just
322+
// accepted and ignored this. Starting to throw an error now would be a breaking change.
323+
// TODO: Consider throwing an Unsupported error for the next major version bump.
324+
return prepareStatement(sql, NO_GENERATED_KEY_COLUMNS);
312325
}
313326

314327
@Override
315328
public PreparedStatement prepareStatement(String sql, String[] columnNames) throws SQLException {
316-
checkClosed();
317-
return prepareStatement(sql);
329+
return prepareStatement(
330+
sql,
331+
isNullOrEmpty(columnNames) ? NO_GENERATED_KEY_COLUMNS : ImmutableList.copyOf(columnNames));
318332
}
319333

320334
@Override

‎java-spanner-jdbc/src/main/java/com/google/cloud/spanner/jdbc/JdbcPreparedStatement.java‎

Lines changed: 16 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,9 @@
2424
import com.google.cloud.spanner.Type;
2525
import com.google.cloud.spanner.connection.AbstractStatementParser.ParametersInfo;
2626
import com.google.common.annotations.VisibleForTesting;
27+
import com.google.common.base.Preconditions;
2728
import com.google.common.collect.ImmutableList;
29+
import com.google.rpc.Code;
2830
import java.sql.PreparedStatement;
2931
import java.sql.ResultSet;
3032
import java.sql.ResultSetMetaData;
@@ -36,8 +38,11 @@ class JdbcPreparedStatement extends AbstractJdbcPreparedStatement {
3638
private final String sql;
3739
private final String sqlWithoutComments;
3840
private final ParametersInfo parameters;
41+
private final ImmutableList<String> generatedKeysColumns;
3942

40-
JdbcPreparedStatement(JdbcConnection connection, String sql) throws SQLException {
43+
JdbcPreparedStatement(
44+
JdbcConnection connection, String sql, ImmutableList<String> generatedKeysColumns)
45+
throws SQLException {
4146
super(connection);
4247
this.sql = sql;
4348
try {
@@ -47,6 +52,7 @@ class JdbcPreparedStatement extends AbstractJdbcPreparedStatement {
4752
} catch (SpannerException e) {
4853
throw JdbcSqlExceptionFactory.of(e);
4954
}
55+
this.generatedKeysColumns = Preconditions.checkNotNull(generatedKeysColumns);
5056
}
5157

5258
ParametersInfo getParametersInfo() {
@@ -76,19 +82,22 @@ ResultSet executeQueryWithOptions(QueryOption... options) throws SQLException {
7682

7783
@Override
7884
public int executeUpdate() throws SQLException {
79-
checkClosed();
80-
return executeUpdate(createStatement());
85+
long count = executeLargeUpdate(createStatement(), generatedKeysColumns);
86+
if (count > Integer.MAX_VALUE) {
87+
throw JdbcSqlExceptionFactory.of(
88+
"update count too large for executeUpdate: " + count, Code.OUT_OF_RANGE);
89+
}
90+
return (int) count;
8191
}
8292

93+
@Override
8394
public long executeLargeUpdate() throws SQLException {
84-
checkClosed();
85-
return executeLargeUpdate(createStatement());
95+
return executeLargeUpdate(createStatement(), generatedKeysColumns);
8696
}
8797

8898
@Override
8999
public boolean execute() throws SQLException {
90-
checkClosed();
91-
return executeStatement(createStatement());
100+
return executeStatement(createStatement(), generatedKeysColumns);
92101
}
93102

94103
@Override

0 commit comments

Comments
 (0)