Skip to content

Commit 13b829c

Browse files
nicktorwaldnicktorwald
committed
Support generated keys from INSERT query
Parse the 'generated_ids' array which is returned after successful INSERT command has applied. This makes sense when a table primary key has an autoincrement property. The driver always returns a predefined result set with a single-column table (column name is 'GENERATED_KEYS') where each row is one generated value. Closes: #77
1 parent 70e65fa commit 13b829c

15 files changed

+242
-71
lines changed

‎src/main/java/org/tarantool/Key.java

+2-1
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,8 @@ public enum Key implements Callable<Integer> {
3333
SQL_BIND(0x41),
3434
SQL_OPTIONS(0x42),
3535
SQL_INFO(0x42),
36-
SQL_ROW_COUNT(0);
36+
SQL_ROW_COUNT(0x00),
37+
SQL_INFO_AUTOINCREMENT_IDS(0x01);
3738

3839
int id;
3940

‎src/main/java/org/tarantool/SqlProtoUtils.java

+11-1
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
import org.tarantool.protocol.TarantoolPacket;
66

77
import java.util.ArrayList;
8+
import java.util.Collections;
89
import java.util.LinkedHashMap;
910
import java.util.List;
1011
import java.util.Map;
@@ -41,7 +42,7 @@ public static List<SQLMetaData> getSQLMetadata(TarantoolPacket pack) {
4142
return values;
4243
}
4344

44-
public static Long getSqlRowCount(TarantoolPacket pack) {
45+
public static Long getSQLRowCount(TarantoolPacket pack) {
4546
Map<Key, Object> info = (Map<Key, Object>) pack.getBody().get(Key.SQL_INFO.getId());
4647
Number rowCount;
4748
if (info != null && (rowCount = ((Number) info.get(Key.SQL_ROW_COUNT.getId()))) != null) {
@@ -50,6 +51,15 @@ public static Long getSqlRowCount(TarantoolPacket pack) {
5051
return null;
5152
}
5253

54+
public static List<Integer> getSQLAutoIncrementIds(TarantoolPacket pack) {
55+
Map<Key, Object> info = (Map<Key, Object>) pack.getBody().get(Key.SQL_INFO.getId());
56+
if (info != null) {
57+
List<Integer> generatedIds = (List<Integer>) info.get(Key.SQL_INFO_AUTOINCREMENT_IDS.getId());
58+
return generatedIds == null ? Collections.emptyList() : generatedIds;
59+
}
60+
return Collections.emptyList();
61+
}
62+
5363
public static class SQLMetaData {
5464
private String name;
5565
private TarantoolSqlType type;

‎src/main/java/org/tarantool/TarantoolClientImpl.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -482,7 +482,7 @@ protected void complete(TarantoolPacket packet, TarantoolOp<?> future) {
482482
}
483483

484484
protected void completeSql(TarantoolOp<?> future, TarantoolPacket pack) {
485-
Long rowCount = SqlProtoUtils.getSqlRowCount(pack);
485+
Long rowCount = SqlProtoUtils.getSQLRowCount(pack);
486486
if (rowCount != null) {
487487
((TarantoolOp) future).complete(rowCount);
488488
} else {

‎src/main/java/org/tarantool/TarantoolConnection.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -80,7 +80,7 @@ public void close() {
8080
@Override
8181
public Long update(String sql, Object... bind) {
8282
TarantoolPacket pack = sql(sql, bind);
83-
return SqlProtoUtils.getSqlRowCount(pack);
83+
return SqlProtoUtils.getSQLRowCount(pack);
8484
}
8585

8686
@Override

‎src/main/java/org/tarantool/jdbc/SQLConnection.java

+7-8
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@
4747
* <p>
4848
* Supports creating {@link Statement} and {@link PreparedStatement} instances
4949
*/
50-
public class SQLConnection implements Connection {
50+
public class SQLConnection implements TarantoolConnection {
5151

5252
private static final int UNSET_HOLDABILITY = 0;
5353
private static final String PING_QUERY = "SELECT 1";
@@ -146,10 +146,7 @@ public PreparedStatement prepareStatement(String sql,
146146
public PreparedStatement prepareStatement(String sql, int autoGeneratedKeys) throws SQLException {
147147
checkNotClosed();
148148
JdbcConstants.checkGeneratedKeysConstant(autoGeneratedKeys);
149-
if (autoGeneratedKeys != Statement.NO_GENERATED_KEYS) {
150-
throw new SQLFeatureNotSupportedException();
151-
}
152-
return prepareStatement(sql);
149+
return new SQLPreparedStatement(this, sql, autoGeneratedKeys);
153150
}
154151

155152
@Override
@@ -525,7 +522,9 @@ public int getNetworkTimeout() throws SQLException {
525522
return (int) client.getOperationTimeout();
526523
}
527524

528-
protected SQLResultHolder execute(long timeout, String sql, Object... args) throws SQLException {
525+
@Override
526+
public SQLResultHolder execute(long timeout, String sql, Object... args)
527+
throws SQLException {
529528
checkNotClosed();
530529
int networkTimeout = getNetworkTimeout();
531530
return (timeout == 0 || (networkTimeout > 0 && networkTimeout < timeout))
@@ -749,10 +748,10 @@ SQLRawOps sqlRawOps() {
749748

750749
@Override
751750
protected void completeSql(TarantoolOp<?> future, TarantoolPacket pack) {
752-
Long rowCount = SqlProtoUtils.getSqlRowCount(pack);
751+
Long rowCount = SqlProtoUtils.getSQLRowCount(pack);
753752
SQLResultHolder result = (rowCount == null)
754753
? SQLResultHolder.ofQuery(SqlProtoUtils.getSQLMetadata(pack), SqlProtoUtils.getSQLData(pack))
755-
: SQLResultHolder.ofUpdate(rowCount.intValue());
754+
: SQLResultHolder.ofUpdate(rowCount.intValue(), SqlProtoUtils.getSQLAutoIncrementIds(pack));
756755
((TarantoolOp) future).complete(result);
757756
}
758757

‎src/main/java/org/tarantool/jdbc/SQLDatabaseMetadata.java

+2-2
Original file line numberDiff line numberDiff line change
@@ -978,7 +978,7 @@ public boolean supportsMultipleOpenResults() throws SQLException {
978978

979979
@Override
980980
public boolean supportsGetGeneratedKeys() throws SQLException {
981-
return false;
981+
return true;
982982
}
983983

984984
@Override
@@ -1104,7 +1104,7 @@ private ResultSet asEmptyMetadataResultSet(List<TupleTwo<String, TarantoolSqlTyp
11041104

11051105
@Override
11061106
public boolean generatedKeyAlwaysReturned() throws SQLException {
1107-
return false;
1107+
return true;
11081108
}
11091109

11101110
@Override

‎src/main/java/org/tarantool/jdbc/SQLPreparedStatement.java

+7-4
Original file line numberDiff line numberDiff line change
@@ -31,12 +31,14 @@ public class SQLPreparedStatement extends SQLStatement implements PreparedStatem
3131
static final String INVALID_CALL_MSG = "The method cannot be called on a PreparedStatement.";
3232
final String sql;
3333
final Map<Integer, Object> params;
34+
final int autoGeneratedKeys;
3435

3536

36-
public SQLPreparedStatement(SQLConnection connection, String sql) throws SQLException {
37+
public SQLPreparedStatement(SQLConnection connection, String sql, int autoGeneratedKeys) throws SQLException {
3738
super(connection);
3839
this.sql = sql;
3940
this.params = new HashMap<>();
41+
this.autoGeneratedKeys = autoGeneratedKeys;
4042
}
4143

4244
public SQLPreparedStatement(SQLConnection connection,
@@ -47,12 +49,13 @@ public SQLPreparedStatement(SQLConnection connection,
4749
super(connection, resultSetType, resultSetConcurrency, resultSetHoldability);
4850
this.sql = sql;
4951
this.params = new HashMap<>();
52+
this.autoGeneratedKeys = NO_GENERATED_KEYS;
5053
}
5154

5255
@Override
5356
public ResultSet executeQuery() throws SQLException {
5457
checkNotClosed();
55-
if (!executeInternal(sql, getParams())) {
58+
if (!executeInternal(autoGeneratedKeys, sql, getParams())) {
5659
throw new SQLException("No results were returned", SQLStates.NO_DATA.getSqlState());
5760
}
5861
return resultSet;
@@ -78,7 +81,7 @@ protected Object[] getParams() throws SQLException {
7881
@Override
7982
public int executeUpdate() throws SQLException {
8083
checkNotClosed();
81-
if (executeInternal(sql, getParams())) {
84+
if (executeInternal(autoGeneratedKeys, sql, getParams())) {
8285
throw new SQLException(
8386
"Result was returned but nothing was expected",
8487
SQLStates.TOO_MANY_RESULTS.getSqlState()
@@ -248,7 +251,7 @@ private void setParameter(int parameterIndex, Object value) throws SQLException
248251
@Override
249252
public boolean execute() throws SQLException {
250253
checkNotClosed();
251-
return executeInternal(sql, getParams());
254+
return executeInternal(autoGeneratedKeys, sql, getParams());
252255
}
253256

254257
@Override

‎src/main/java/org/tarantool/jdbc/SQLResultHolder.java

+13-8
Original file line numberDiff line numberDiff line change
@@ -14,24 +14,25 @@ public class SQLResultHolder {
1414
final List<SqlProtoUtils.SQLMetaData> sqlMetadata;
1515
final List<List<Object>> rows;
1616
final int updateCount;
17+
final List<Integer> generatedIds;
1718

18-
public SQLResultHolder(List<SqlProtoUtils.SQLMetaData> sqlMetadata, List<List<Object>> rows, int updateCount) {
19+
public SQLResultHolder(List<SqlProtoUtils.SQLMetaData> sqlMetadata,
20+
List<List<Object>> rows,
21+
int updateCount,
22+
List<Integer> generatedIds) {
1923
this.sqlMetadata = sqlMetadata;
2024
this.rows = rows;
2125
this.updateCount = updateCount;
26+
this.generatedIds = generatedIds;
2227
}
2328

2429
public static SQLResultHolder ofQuery(final List<SqlProtoUtils.SQLMetaData> sqlMetadata,
2530
final List<List<Object>> rows) {
26-
return new SQLResultHolder(sqlMetadata, rows, -1);
31+
return new SQLResultHolder(sqlMetadata, rows, -1, Collections.emptyList());
2732
}
2833

29-
public static SQLResultHolder ofEmptyQuery() {
30-
return ofQuery(Collections.emptyList(), Collections.emptyList());
31-
}
32-
33-
public static SQLResultHolder ofUpdate(int updateCount) {
34-
return new SQLResultHolder(null, null, updateCount);
34+
public static SQLResultHolder ofUpdate(int updateCount, List<Integer> generatedIds) {
35+
return new SQLResultHolder(null, null, updateCount, generatedIds);
3536
}
3637

3738
public List<SqlProtoUtils.SQLMetaData> getSqlMetadata() {
@@ -46,6 +47,10 @@ public int getUpdateCount() {
4647
return updateCount;
4748
}
4849

50+
public List<Integer> getGeneratedIds() {
51+
return generatedIds;
52+
}
53+
4954
public boolean isQueryResult() {
5055
return sqlMetadata != null && rows != null;
5156
}

‎src/main/java/org/tarantool/jdbc/SQLStatement.java

+42-24
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
package org.tarantool.jdbc;
22

3+
import org.tarantool.SqlProtoUtils;
4+
import org.tarantool.jdbc.type.TarantoolSqlType;
35
import org.tarantool.util.JdbcConstants;
46
import org.tarantool.util.SQLStates;
57

@@ -11,8 +13,11 @@
1113
import java.sql.SQLTimeoutException;
1214
import java.sql.SQLWarning;
1315
import java.sql.Statement;
16+
import java.util.Collections;
17+
import java.util.List;
1418
import java.util.concurrent.TimeUnit;
1519
import java.util.concurrent.atomic.AtomicBoolean;
20+
import java.util.stream.Collectors;
1621

1722
/**
1823
* Tarantool {@link Statement} implementation.
@@ -23,13 +28,17 @@
2328
*/
2429
public class SQLStatement implements TarantoolStatement {
2530

26-
protected final SQLConnection connection;
31+
private static final String GENERATED_KEY_COLUMN_NAME = "GENERATED_KEY";
32+
33+
protected final TarantoolConnection connection;
34+
private final SQLResultSet emptyGeneratedKeys;
2735

2836
/**
2937
* Current result set / update count associated to this statement.
3038
*/
3139
protected SQLResultSet resultSet;
3240
protected int updateCount;
41+
protected SQLResultSet generatedKeys;
3342

3443
private boolean isCloseOnCompletion;
3544

@@ -47,10 +56,12 @@ public class SQLStatement implements TarantoolStatement {
4756
private final AtomicBoolean isClosed = new AtomicBoolean(false);
4857

4958
protected SQLStatement(SQLConnection sqlConnection) throws SQLException {
50-
this.connection = sqlConnection;
51-
this.resultSetType = ResultSet.TYPE_FORWARD_ONLY;
52-
this.resultSetConcurrency = ResultSet.CONCUR_READ_ONLY;
53-
this.resultSetHoldability = sqlConnection.getHoldability();
59+
this(
60+
sqlConnection,
61+
ResultSet.TYPE_FORWARD_ONLY,
62+
ResultSet.CONCUR_READ_ONLY,
63+
sqlConnection.getHoldability()
64+
);
5465
}
5566

5667
protected SQLStatement(SQLConnection sqlConnection,
@@ -61,37 +72,34 @@ protected SQLStatement(SQLConnection sqlConnection,
6172
this.resultSetType = resultSetType;
6273
this.resultSetConcurrency = resultSetConcurrency;
6374
this.resultSetHoldability = resultSetHoldability;
75+
this.emptyGeneratedKeys = this.generatedKeys = executeGeneratedKeys(Collections.emptyList());
6476
}
6577

6678
@Override
6779
public ResultSet executeQuery(String sql) throws SQLException {
6880
checkNotClosed();
69-
if (!executeInternal(sql)) {
81+
if (!executeInternal(NO_GENERATED_KEYS, sql)) {
7082
throw new SQLException("No results were returned", SQLStates.NO_DATA.getSqlState());
7183
}
7284
return resultSet;
7385
}
7486

7587
@Override
7688
public int executeUpdate(String sql) throws SQLException {
77-
checkNotClosed();
78-
if (executeInternal(sql)) {
79-
throw new SQLException(
80-
"Result was returned but nothing was expected",
81-
SQLStates.TOO_MANY_RESULTS.getSqlState()
82-
);
83-
}
84-
return updateCount;
89+
return executeUpdate(sql, NO_GENERATED_KEYS);
8590
}
8691

8792
@Override
8893
public int executeUpdate(String sql, int autoGeneratedKeys) throws SQLException {
8994
checkNotClosed();
9095
JdbcConstants.checkGeneratedKeysConstant(autoGeneratedKeys);
91-
if (autoGeneratedKeys != Statement.NO_GENERATED_KEYS) {
92-
throw new SQLFeatureNotSupportedException();
96+
if (executeInternal(autoGeneratedKeys, sql)) {
97+
throw new SQLException(
98+
"Result was returned but nothing was expected",
99+
SQLStates.TOO_MANY_RESULTS.getSqlState()
100+
);
93101
}
94-
return executeUpdate(sql);
102+
return updateCount;
95103
}
96104

97105
@Override
@@ -181,17 +189,14 @@ public void setCursorName(String name) throws SQLException {
181189
@Override
182190
public boolean execute(String sql) throws SQLException {
183191
checkNotClosed();
184-
return executeInternal(sql);
192+
return executeInternal(NO_GENERATED_KEYS, sql);
185193
}
186194

187195
@Override
188196
public boolean execute(String sql, int autoGeneratedKeys) throws SQLException {
189197
checkNotClosed();
190198
JdbcConstants.checkGeneratedKeysConstant(autoGeneratedKeys);
191-
if (autoGeneratedKeys != Statement.NO_GENERATED_KEYS) {
192-
throw new SQLFeatureNotSupportedException();
193-
}
194-
return execute(sql);
199+
return executeInternal(autoGeneratedKeys, sql);
195200
}
196201

197202
@Override
@@ -296,7 +301,7 @@ public Connection getConnection() throws SQLException {
296301
@Override
297302
public ResultSet getGeneratedKeys() throws SQLException {
298303
checkNotClosed();
299-
return new SQLResultSet(SQLResultHolder.ofEmptyQuery(), this);
304+
return generatedKeys;
300305
}
301306

302307
@Override
@@ -374,6 +379,7 @@ protected void discardLastResults() throws SQLException {
374379
clearWarnings();
375380
updateCount = -1;
376381
resultSet = null;
382+
generatedKeys = emptyGeneratedKeys;
377383

378384
if (lastResultSet != null) {
379385
try {
@@ -392,7 +398,7 @@ protected void discardLastResults() throws SQLException {
392398
*
393399
* @return {@code true}, if the result is a ResultSet object;
394400
*/
395-
protected boolean executeInternal(String sql, Object... params) throws SQLException {
401+
protected boolean executeInternal(int autoGeneratedKeys, String sql, Object... params) throws SQLException {
396402
discardLastResults();
397403
SQLResultHolder holder;
398404
try {
@@ -406,6 +412,9 @@ protected boolean executeInternal(String sql, Object... params) throws SQLExcept
406412
resultSet = new SQLResultSet(holder, this);
407413
}
408414
updateCount = holder.getUpdateCount();
415+
if (autoGeneratedKeys == Statement.RETURN_GENERATED_KEYS) {
416+
generatedKeys = executeGeneratedKeys(holder.getGeneratedIds());
417+
}
409418
return holder.isQueryResult();
410419
}
411420

@@ -425,4 +434,13 @@ protected void checkNotClosed() throws SQLException {
425434
}
426435
}
427436

437+
protected SQLResultSet executeGeneratedKeys(List<Integer> generatedKeys) throws SQLException {
438+
SqlProtoUtils.SQLMetaData sqlMetaData =
439+
new SqlProtoUtils.SQLMetaData(GENERATED_KEY_COLUMN_NAME, TarantoolSqlType.INTEGER);
440+
List<List<Object>> rows = generatedKeys.stream()
441+
.map(Collections::<Object>singletonList)
442+
.collect(Collectors.toList());
443+
return createResultSet(SQLResultHolder.ofQuery(Collections.singletonList(sqlMetaData), rows));
444+
}
445+
428446
}

0 commit comments

Comments
 (0)