Skip to content

Commit e89eb07

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 455a45f commit e89eb07

15 files changed

+257
-69
lines changed

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

Lines changed: 2 additions & 1 deletion
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

Lines changed: 11 additions & 1 deletion
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

Lines changed: 1 addition & 1 deletion
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

Lines changed: 1 addition & 1 deletion
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

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,7 @@
4949
* <p>
5050
* Supports creating {@link Statement} and {@link PreparedStatement} instances
5151
*/
52-
public class SQLConnection implements Connection {
52+
public class SQLConnection implements TarantoolConnection {
5353

5454
private static final int UNSET_HOLDABILITY = 0;
5555
private static final String PING_QUERY = "SELECT 1";
@@ -148,10 +148,7 @@ public PreparedStatement prepareStatement(String sql,
148148
public PreparedStatement prepareStatement(String sql, int autoGeneratedKeys) throws SQLException {
149149
checkNotClosed();
150150
JdbcConstants.checkGeneratedKeysConstant(autoGeneratedKeys);
151-
if (autoGeneratedKeys != Statement.NO_GENERATED_KEYS) {
152-
throw new SQLFeatureNotSupportedException();
153-
}
154-
return prepareStatement(sql);
151+
return new SQLPreparedStatement(this, sql, autoGeneratedKeys);
155152
}
156153

157154
@Override
@@ -527,14 +524,17 @@ public int getNetworkTimeout() throws SQLException {
527524
return (int) client.getOperationTimeout();
528525
}
529526

530-
protected SQLResultHolder execute(long timeout, SQLQueryHolder query) throws SQLException {
527+
@Override
528+
public SQLResultHolder execute(long timeout, SQLQueryHolder query) throws SQLException {
531529
checkNotClosed();
532530
return (useNetworkTimeout(timeout))
533531
? executeWithNetworkTimeout(query)
534532
: executeWithQueryTimeout(timeout, query);
535533
}
536534

537-
protected SQLBatchResultHolder executeBatch(long timeout, List<SQLQueryHolder> queries) throws SQLException {
535+
@Override
536+
public SQLBatchResultHolder executeBatch(long timeout, List<SQLQueryHolder> queries)
537+
throws SQLException {
538538
checkNotClosed();
539539
SQLTarantoolClientImpl.SQLRawOps sqlOps = client.sqlRawOps();
540540
SQLBatchResultHolder batchResult = useNetworkTimeout(timeout)
@@ -810,10 +810,10 @@ SQLRawOps sqlRawOps() {
810810

811811
@Override
812812
protected void completeSql(TarantoolOp<?> future, TarantoolPacket pack) {
813-
Long rowCount = SqlProtoUtils.getSqlRowCount(pack);
813+
Long rowCount = SqlProtoUtils.getSQLRowCount(pack);
814814
SQLResultHolder result = (rowCount == null)
815815
? SQLResultHolder.ofQuery(SqlProtoUtils.getSQLMetadata(pack), SqlProtoUtils.getSQLData(pack))
816-
: SQLResultHolder.ofUpdate(rowCount.intValue());
816+
: SQLResultHolder.ofUpdate(rowCount.intValue(), SqlProtoUtils.getSQLAutoIncrementIds(pack));
817817
((TarantoolOp) future).complete(result);
818818
}
819819

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

Lines changed: 2 additions & 2 deletions
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

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -34,13 +34,14 @@ public class SQLPreparedStatement extends SQLStatement implements PreparedStatem
3434

3535
private final String sql;
3636
private final Map<Integer, Object> parameters;
37-
37+
private final int autoGeneratedKeys;
3838
private List<Map<Integer, Object>> batchParameters = new ArrayList<>();
3939

40-
public SQLPreparedStatement(SQLConnection connection, String sql) throws SQLException {
40+
public SQLPreparedStatement(SQLConnection connection, String sql, int autoGeneratedKeys) throws SQLException {
4141
super(connection);
4242
this.sql = sql;
4343
this.parameters = new HashMap<>();
44+
this.autoGeneratedKeys = autoGeneratedKeys;
4445
setPoolable(true);
4546
}
4647

@@ -52,13 +53,14 @@ public SQLPreparedStatement(SQLConnection connection,
5253
super(connection, resultSetType, resultSetConcurrency, resultSetHoldability);
5354
this.sql = sql;
5455
this.parameters = new HashMap<>();
56+
this.autoGeneratedKeys = NO_GENERATED_KEYS;
5557
setPoolable(true);
5658
}
5759

5860
@Override
5961
public ResultSet executeQuery() throws SQLException {
6062
checkNotClosed();
61-
if (!executeInternal(sql, toParametersList(parameters))) {
63+
if (!executeInternal(autoGeneratedKeys, sql, toParametersList(parameters))) {
6264
throw new SQLException("No results were returned", SQLStates.NO_DATA.getSqlState());
6365
}
6466
return resultSet;
@@ -73,7 +75,7 @@ public ResultSet executeQuery(String sql) throws SQLException {
7375
@Override
7476
public int executeUpdate() throws SQLException {
7577
checkNotClosed();
76-
if (executeInternal(sql, toParametersList(parameters))) {
78+
if (executeInternal(autoGeneratedKeys, sql, toParametersList(parameters))) {
7779
throw new SQLException(
7880
"Result was returned but nothing was expected",
7981
SQLStates.TOO_MANY_RESULTS.getSqlState()
@@ -244,7 +246,7 @@ private void setParameter(int parameterIndex, Object value) throws SQLException
244246
@Override
245247
public boolean execute() throws SQLException {
246248
checkNotClosed();
247-
return executeInternal(sql, toParametersList(parameters));
249+
return executeInternal(autoGeneratedKeys, sql, toParametersList(parameters));
248250
}
249251

250252
@Override

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

Lines changed: 13 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -16,24 +16,29 @@ public class SQLResultHolder {
1616
private final List<SqlProtoUtils.SQLMetaData> sqlMetadata;
1717
private final List<List<Object>> rows;
1818
private final int updateCount;
19+
private final List<Integer> generatedIds;
1920

20-
public SQLResultHolder(List<SqlProtoUtils.SQLMetaData> sqlMetadata, List<List<Object>> rows, int updateCount) {
21+
public SQLResultHolder(List<SqlProtoUtils.SQLMetaData> sqlMetadata,
22+
List<List<Object>> rows,
23+
int updateCount,
24+
List<Integer> generatedIds) {
2125
this.sqlMetadata = sqlMetadata;
2226
this.rows = rows;
2327
this.updateCount = updateCount;
28+
this.generatedIds = generatedIds;
2429
}
2530

2631
public static SQLResultHolder ofQuery(final List<SqlProtoUtils.SQLMetaData> sqlMetadata,
2732
final List<List<Object>> rows) {
28-
return new SQLResultHolder(sqlMetadata, rows, NO_UPDATE_COUNT);
33+
return new SQLResultHolder(sqlMetadata, rows, NO_UPDATE_COUNT, Collections.emptyList());
2934
}
3035

3136
public static SQLResultHolder ofEmptyQuery() {
3237
return ofQuery(Collections.emptyList(), Collections.emptyList());
3338
}
3439

35-
public static SQLResultHolder ofUpdate(int updateCount) {
36-
return new SQLResultHolder(null, null, updateCount);
40+
public static SQLResultHolder ofUpdate(int updateCount, List<Integer> generatedIds) {
41+
return new SQLResultHolder(null, null, updateCount, generatedIds);
3742
}
3843

3944
public List<SqlProtoUtils.SQLMetaData> getSqlMetadata() {
@@ -48,6 +53,10 @@ public int getUpdateCount() {
4853
return updateCount;
4954
}
5055

56+
public List<Integer> getGeneratedIds() {
57+
return generatedIds;
58+
}
59+
5160
public boolean isQueryResult() {
5261
return sqlMetadata != null && rows != null;
5362
}

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

Lines changed: 40 additions & 24 deletions
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

@@ -13,6 +15,7 @@
1315
import java.sql.SQLWarning;
1416
import java.sql.Statement;
1517
import java.util.ArrayList;
18+
import java.util.Collections;
1619
import java.util.List;
1720
import java.util.concurrent.TimeUnit;
1821
import java.util.concurrent.atomic.AtomicBoolean;
@@ -27,13 +30,17 @@
2730
*/
2831
public class SQLStatement implements TarantoolStatement {
2932

30-
protected final SQLConnection connection;
33+
private static final String GENERATED_KEY_COLUMN_NAME = "GENERATED_KEY";
34+
35+
protected final TarantoolConnection connection;
36+
private final SQLResultSet emptyGeneratedKeys;
3137

3238
/**
3339
* Current result set / update count associated to this statement.
3440
*/
3541
protected SQLResultSet resultSet;
3642
protected int updateCount;
43+
protected SQLResultSet generatedKeys;
3744

3845
private List<String> batchQueries = new ArrayList<>();
3946

@@ -61,10 +68,12 @@ public class SQLStatement implements TarantoolStatement {
6168
private final AtomicBoolean isClosed = new AtomicBoolean(false);
6269

6370
protected SQLStatement(SQLConnection sqlConnection) throws SQLException {
64-
this.connection = sqlConnection;
65-
this.resultSetType = ResultSet.TYPE_FORWARD_ONLY;
66-
this.resultSetConcurrency = ResultSet.CONCUR_READ_ONLY;
67-
this.resultSetHoldability = sqlConnection.getHoldability();
71+
this(
72+
sqlConnection,
73+
ResultSet.TYPE_FORWARD_ONLY,
74+
ResultSet.CONCUR_READ_ONLY,
75+
sqlConnection.getHoldability()
76+
);
6877
}
6978

7079
protected SQLStatement(SQLConnection sqlConnection,
@@ -75,37 +84,34 @@ protected SQLStatement(SQLConnection sqlConnection,
7584
this.resultSetType = resultSetType;
7685
this.resultSetConcurrency = resultSetConcurrency;
7786
this.resultSetHoldability = resultSetHoldability;
87+
this.emptyGeneratedKeys = this.generatedKeys = executeGeneratedKeys(Collections.emptyList());
7888
}
7989

8090
@Override
8191
public ResultSet executeQuery(String sql) throws SQLException {
8292
checkNotClosed();
83-
if (!executeInternal(sql)) {
93+
if (!executeInternal(NO_GENERATED_KEYS, sql)) {
8494
throw new SQLException("No results were returned", SQLStates.NO_DATA.getSqlState());
8595
}
8696
return resultSet;
8797
}
8898

8999
@Override
90100
public int executeUpdate(String sql) throws SQLException {
91-
checkNotClosed();
92-
if (executeInternal(sql)) {
93-
throw new SQLException(
94-
"Result was returned but nothing was expected",
95-
SQLStates.TOO_MANY_RESULTS.getSqlState()
96-
);
97-
}
98-
return updateCount;
101+
return executeUpdate(sql, NO_GENERATED_KEYS);
99102
}
100103

101104
@Override
102105
public int executeUpdate(String sql, int autoGeneratedKeys) throws SQLException {
103106
checkNotClosed();
104107
JdbcConstants.checkGeneratedKeysConstant(autoGeneratedKeys);
105-
if (autoGeneratedKeys != Statement.NO_GENERATED_KEYS) {
106-
throw new SQLFeatureNotSupportedException();
108+
if (executeInternal(autoGeneratedKeys, sql)) {
109+
throw new SQLException(
110+
"Result was returned but nothing was expected",
111+
SQLStates.TOO_MANY_RESULTS.getSqlState()
112+
);
107113
}
108-
return executeUpdate(sql);
114+
return updateCount;
109115
}
110116

111117
@Override
@@ -195,17 +201,14 @@ public void setCursorName(String name) throws SQLException {
195201
@Override
196202
public boolean execute(String sql) throws SQLException {
197203
checkNotClosed();
198-
return executeInternal(sql);
204+
return executeInternal(NO_GENERATED_KEYS, sql);
199205
}
200206

201207
@Override
202208
public boolean execute(String sql, int autoGeneratedKeys) throws SQLException {
203209
checkNotClosed();
204210
JdbcConstants.checkGeneratedKeysConstant(autoGeneratedKeys);
205-
if (autoGeneratedKeys != Statement.NO_GENERATED_KEYS) {
206-
throw new SQLFeatureNotSupportedException();
207-
}
208-
return execute(sql);
211+
return executeInternal(autoGeneratedKeys, sql);
209212
}
210213

211214
@Override
@@ -321,7 +324,7 @@ public Connection getConnection() throws SQLException {
321324
@Override
322325
public ResultSet getGeneratedKeys() throws SQLException {
323326
checkNotClosed();
324-
return new SQLResultSet(SQLResultHolder.ofEmptyQuery(), this);
327+
return generatedKeys;
325328
}
326329

327330
@Override
@@ -401,6 +404,7 @@ protected void discardLastResults() throws SQLException {
401404
clearWarnings();
402405
updateCount = -1;
403406
resultSet = null;
407+
generatedKeys = emptyGeneratedKeys;
404408

405409
if (lastResultSet != null) {
406410
try {
@@ -419,7 +423,7 @@ protected void discardLastResults() throws SQLException {
419423
*
420424
* @return {@code true}, if the result is a ResultSet object;
421425
*/
422-
protected boolean executeInternal(String sql, Object... params) throws SQLException {
426+
protected boolean executeInternal(int autoGeneratedKeys, String sql, Object... params) throws SQLException {
423427
discardLastResults();
424428
SQLResultHolder holder;
425429
try {
@@ -433,6 +437,9 @@ protected boolean executeInternal(String sql, Object... params) throws SQLExcept
433437
resultSet = new SQLResultSet(holder, this);
434438
}
435439
updateCount = holder.getUpdateCount();
440+
if (autoGeneratedKeys == Statement.RETURN_GENERATED_KEYS) {
441+
generatedKeys = executeGeneratedKeys(holder.getGeneratedIds());
442+
}
436443
return holder.isQueryResult();
437444
}
438445

@@ -474,4 +481,13 @@ protected void checkNotClosed() throws SQLException {
474481
}
475482
}
476483

484+
protected SQLResultSet executeGeneratedKeys(List<Integer> generatedKeys) throws SQLException {
485+
SqlProtoUtils.SQLMetaData sqlMetaData =
486+
new SqlProtoUtils.SQLMetaData(GENERATED_KEY_COLUMN_NAME, TarantoolSqlType.INTEGER);
487+
List<List<Object>> rows = generatedKeys.stream()
488+
.map(Collections::<Object>singletonList)
489+
.collect(Collectors.toList());
490+
return createResultSet(SQLResultHolder.ofQuery(Collections.singletonList(sqlMetaData), rows));
491+
}
492+
477493
}

0 commit comments

Comments
 (0)