Skip to content

Commit cdc4a9b

Browse files
alex268alex268
authored
Added support of generated keys (#120)
2 parents 7d4cd74 + 2acfcf6 commit cdc4a9b

File tree

14 files changed

+289
-43
lines changed

14 files changed

+289
-43
lines changed

‎jdbc/src/main/java/tech/ydb/jdbc/YdbConst.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,8 @@ public final class YdbConst {
6666
public static final String INVALID_ROW = "Current row index is out of bounds: ";
6767
public static final String BULKS_UNSUPPORTED = "BULK mode is available only for prepared statement with one UPSERT";
6868
public static final String INVALID_BATCH_COLUMN = "Cannot prepared batch request: cannot find a column";
69-
public static final String BULKS_DESCRIBE_ERROR = "Cannot parse BULK upsert: ";
69+
public static final String BULK_DESCRIBE_ERROR = "Cannot parse BULK upsert: ";
70+
public static final String BULK_NOT_SUPPORT_RETURNING = "BULK query doesn't support RETURNING";
7071
public static final String METADATA_RS_UNSUPPORTED_IN_PS = "ResultSet metadata is not supported " +
7172
"in prepared statements";
7273
public static final String CANNOT_UNWRAP_TO = "Cannot unwrap to ";

‎jdbc/src/main/java/tech/ydb/jdbc/context/StaticQueryResult.java‎

Lines changed: 26 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -34,15 +34,24 @@ public class StaticQueryResult implements YdbQueryResult {
3434
private static class ExpressionResult {
3535
private final int updateCount;
3636
private final YdbResultSet resultSet;
37+
private final YdbResultSet generatedKeys;
3738

3839
ExpressionResult(int updateCount) {
3940
this.updateCount = updateCount;
4041
this.resultSet = null;
42+
this.generatedKeys = null;
4143
}
4244

43-
ExpressionResult(YdbResultSet resultSet) {
45+
ExpressionResult(int updateCount, YdbResultSet keys) {
46+
this.updateCount = updateCount;
47+
this.resultSet = null;
48+
this.generatedKeys = keys;
49+
}
50+
51+
ExpressionResult(YdbResultSet result) {
4452
this.updateCount = -1;
45-
this.resultSet = resultSet;
53+
this.resultSet = result;
54+
this.generatedKeys = null;
4655
}
4756
}
4857

@@ -59,6 +68,13 @@ public StaticQueryResult(YdbQuery query, List<YdbResultSet> list) {
5968
results.add(NO_UPDATED);
6069
continue;
6170
}
71+
72+
if (exp.hasUpdateWithGenerated() && idx < list.size()) {
73+
results.add(new ExpressionResult(1, list.get(idx)));
74+
idx++;
75+
continue;
76+
}
77+
6278
if (exp.hasUpdateCount()) {
6379
results.add(HAS_UPDATED);
6480
continue;
@@ -124,6 +140,14 @@ public YdbResultSet getCurrentResultSet() {
124140
return results.get(resultIndex).resultSet;
125141
}
126142

143+
@Override
144+
public YdbResultSet getGeneratedKeys() {
145+
if (results == null || resultIndex >= results.size()) {
146+
return null;
147+
}
148+
return results.get(resultIndex).generatedKeys;
149+
}
150+
127151
@Override
128152
public boolean getMoreResults(int current) throws SQLException {
129153
if (results == null || resultIndex >= results.size()) {

‎jdbc/src/main/java/tech/ydb/jdbc/context/StreamQueryResult.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -219,6 +219,12 @@ public YdbResultSet getCurrentResultSet() throws SQLException {
219219
}
220220
}
221221

222+
@Override
223+
public YdbResultSet getGeneratedKeys() throws SQLException {
224+
// TODO: Implement
225+
return null;
226+
}
227+
222228
@Override
223229
public boolean hasResultSets() throws SQLException {
224230
if (resultIndex >= resultIndexes.length) {

‎jdbc/src/main/java/tech/ydb/jdbc/context/YdbContext.java‎

Lines changed: 14 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
import tech.ydb.jdbc.exception.ExceptionFactory;
3030
import tech.ydb.jdbc.impl.YdbTracerImpl;
3131
import tech.ydb.jdbc.impl.YdbTracerNone;
32+
import tech.ydb.jdbc.query.QueryKey;
3233
import tech.ydb.jdbc.query.QueryType;
3334
import tech.ydb.jdbc.query.YdbPreparedQuery;
3435
import tech.ydb.jdbc.query.YdbQuery;
@@ -83,7 +84,7 @@ public class YdbContext implements AutoCloseable {
8384
private final String prefixPath;
8485
private final String prefixPragma;
8586

86-
private final Cache<String, YdbQuery> queriesCache;
87+
private final Cache<QueryKey, YdbQuery> queriesCache;
8788
private final Cache<String, QueryStat> statsCache;
8889
private final Cache<String, Map<String, Type>> queryParamsCache;
8990
private final Cache<String, TableDescription> tableDescribeCache;
@@ -348,22 +349,21 @@ public <T extends BaseRequestSettings.BaseBuilder<T>> T withRequestTimeout(T bui
348349
return builder.withRequestTimeout(operation);
349350
}
350351

351-
public YdbQuery parseYdbQuery(String sql) throws SQLException {
352-
return YdbQuery.parseQuery(sql, queryOptions, types);
352+
public YdbQuery parseYdbQuery(String query) throws SQLException {
353+
return YdbQuery.parseQuery(new QueryKey(query), queryOptions, types);
353354
}
354355

355-
public YdbQuery findOrParseYdbQuery(String sql) throws SQLException {
356+
public YdbQuery findOrParseYdbQuery(QueryKey key) throws SQLException {
356357
if (queriesCache == null) {
357-
return parseYdbQuery(sql);
358+
return YdbQuery.parseQuery(key, queryOptions, types);
358359
}
359360

360-
YdbQuery cached = queriesCache.getIfPresent(sql);
361+
YdbQuery cached = queriesCache.getIfPresent(key);
361362
if (cached == null) {
362-
cached = parseYdbQuery(sql);
363-
queriesCache.put(sql, cached);
363+
cached = YdbQuery.parseQuery(key, queryOptions, types);
364+
queriesCache.put(key, cached);
364365
}
365366

366-
367367
return cached;
368368
}
369369

@@ -432,16 +432,19 @@ public YdbPreparedQuery findOrPrepareParams(YdbQuery query, YdbPrepareMode mode)
432432
tableDescribeCache.put(tablePath, description);
433433
} else {
434434
if (type == QueryType.BULK_QUERY) {
435-
throw new SQLException(YdbConst.BULKS_DESCRIBE_ERROR + result.getStatus());
435+
throw new SQLException(YdbConst.BULK_DESCRIBE_ERROR + result.getStatus());
436436
}
437437
}
438438
}
439439
if (type == QueryType.BULK_QUERY) {
440+
if (query.getReturning() != null) {
441+
throw new SQLException(YdbConst.BULK_NOT_SUPPORT_RETURNING);
442+
}
440443
return BulkUpsertQuery.build(types, tablePath, query.getYqlBatcher().getColumns(), description);
441444
}
442445

443446
if (description != null) {
444-
BatchedQuery params = BatchedQuery.createAutoBatched(types, query.getYqlBatcher(), description);
447+
BatchedQuery params = BatchedQuery.createAutoBatched(types, query, description);
445448
if (params != null) {
446449
return params;
447450
}

‎jdbc/src/main/java/tech/ydb/jdbc/impl/BaseYdbStatement.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -244,7 +244,7 @@ public void setCursorName(String name) throws SQLFeatureNotSupportedException {
244244

245245
@Override
246246
public ResultSet getGeneratedKeys() throws SQLException {
247-
return null; // --
247+
return state.getGeneratedKeys();
248248
}
249249

250250
@Override

‎jdbc/src/main/java/tech/ydb/jdbc/impl/YdbConnectionImpl.java‎

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
import tech.ydb.jdbc.context.YdbContext;
3030
import tech.ydb.jdbc.context.YdbExecutor;
3131
import tech.ydb.jdbc.context.YdbValidator;
32+
import tech.ydb.jdbc.query.QueryKey;
3233
import tech.ydb.jdbc.query.YdbPreparedQuery;
3334
import tech.ydb.jdbc.query.YdbQuery;
3435

@@ -213,15 +214,15 @@ public YdbStatement createStatement(int resultSetType, int resultSetConcurrency,
213214
}
214215

215216
@Override
216-
public YdbPreparedStatement prepareStatement(String origSql, int resultSetType, int resultSetConcurrency,
217+
public YdbPreparedStatement prepareStatement(String sql, int resultSetType, int resultSetConcurrency,
217218
int resultSetHoldability) throws SQLException {
218219
checkStatementParams(resultSetType, resultSetConcurrency, resultSetHoldability);
219-
return prepareStatement(origSql, resultSetType, null, YdbPrepareMode.AUTO);
220+
return prepareStatement(new QueryKey(sql), resultSetType, YdbPrepareMode.AUTO);
220221
}
221222

222223
@Override
223224
public YdbPreparedStatement prepareStatement(String sql, YdbPrepareMode mode) throws SQLException {
224-
return prepareStatement(sql, ResultSet.TYPE_SCROLL_INSENSITIVE, null, mode);
225+
return prepareStatement(new QueryKey(sql), ResultSet.TYPE_SCROLL_INSENSITIVE, mode);
225226
}
226227

227228
@Override
@@ -245,18 +246,18 @@ public YdbPreparedStatement prepareStatement(String sql, String[] columnNames) t
245246
if (columnNames == null || columnNames.length == 0) {
246247
return prepareStatement(sql);
247248
}
248-
return prepareStatement(sql, ResultSet.TYPE_FORWARD_ONLY, columnNames, YdbPrepareMode.AUTO);
249+
return prepareStatement(new QueryKey(sql, columnNames), ResultSet.TYPE_FORWARD_ONLY, YdbPrepareMode.AUTO);
249250
}
250251

251-
private YdbPreparedStatement prepareStatement(String sql, int type, String[] columnNames, YdbPrepareMode mode)
252+
private YdbPreparedStatement prepareStatement(QueryKey key, int resultSetType, YdbPrepareMode mode)
252253
throws SQLException {
253254

254255
validator.clearWarnings();
255256
ctx.getTracer().trace("prepare statement");
256-
YdbQuery query = ctx.findOrParseYdbQuery(sql);
257+
YdbQuery query = ctx.findOrParseYdbQuery(key);
257258
YdbPreparedQuery params = ctx.findOrPrepareParams(query, mode);
258259
ctx.getTracer().trace("create prepared statement");
259-
return new YdbPreparedStatementImpl(this, query, params, type);
260+
return new YdbPreparedStatementImpl(this, query, params, resultSetType);
260261
}
261262

262263
@Override

‎jdbc/src/main/java/tech/ydb/jdbc/impl/YdbQueryResult.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,11 @@ public YdbResultSet getCurrentResultSet() throws SQLException {
2020
return null;
2121
}
2222

23+
@Override
24+
public YdbResultSet getGeneratedKeys() throws SQLException {
25+
return null;
26+
}
27+
2328
@Override
2429
public boolean hasResultSets() throws SQLException {
2530
return false;
@@ -36,6 +41,7 @@ public void close() throws SQLException { }
3641

3742
int getUpdateCount() throws SQLException;
3843
YdbResultSet getCurrentResultSet() throws SQLException;
44+
YdbResultSet getGeneratedKeys() throws SQLException;
3945

4046
boolean hasResultSets() throws SQLException;
4147
boolean getMoreResults(int current) throws SQLException;

‎jdbc/src/main/java/tech/ydb/jdbc/query/QueryKey.java‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
package tech.ydb.jdbc.query;
2+
3+
import java.util.Objects;
4+
5+
/**
6+
*
7+
* @author Aleksandr Gorshenin
8+
*/
9+
public class QueryKey {
10+
private final String query;
11+
private final String returning;
12+
13+
public QueryKey(String query) {
14+
this.query = query;
15+
this.returning = null;
16+
}
17+
18+
public QueryKey(String query, String[] columnNames) {
19+
this.query = query;
20+
this.returning = buildReturning(columnNames);
21+
}
22+
23+
public String getQuery() {
24+
return query;
25+
}
26+
27+
public String getReturning() {
28+
return returning;
29+
}
30+
31+
@Override
32+
public int hashCode() {
33+
return Objects.hash(query, returning);
34+
}
35+
36+
@Override
37+
public boolean equals(Object obj) {
38+
if (this == obj) {
39+
return true;
40+
}
41+
if (obj == null || getClass() != obj.getClass()) {
42+
return false;
43+
}
44+
45+
QueryKey other = (QueryKey) obj;
46+
return Objects.equals(this.query, other.query) && Objects.equals(this.returning, other.returning);
47+
}
48+
49+
private static String buildReturning(String[] columnNames) {
50+
StringBuilder sb = new StringBuilder();
51+
sb.append("RETURNING ");
52+
if (columnNames.length == 1 && columnNames[0].charAt(0) == '*') {
53+
sb.append('*');
54+
return sb.toString();
55+
}
56+
57+
for (int col = 0; col < columnNames.length; col++) {
58+
String columnName = columnNames[col];
59+
if (col > 0) {
60+
sb.append(", ");
61+
}
62+
63+
if (columnName.charAt(0) == '`') {
64+
sb.append(columnName);
65+
} else {
66+
sb.append('`').append(columnName).append('`');
67+
}
68+
}
69+
70+
return sb.toString();
71+
}
72+
}

‎jdbc/src/main/java/tech/ydb/jdbc/query/QueryStatement.java‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ public class QueryStatement {
1616
private final QueryCmd command;
1717
private final List<JdbcPrm.Factory> parameters = new ArrayList<>();
1818
private boolean hasReturinng = false;
19+
private boolean hasGenerated = false;
1920

2021
public QueryStatement(QueryType custom, QueryType baseType, QueryCmd command) {
2122
this.queryType = custom != null ? custom : baseType;
@@ -46,10 +47,18 @@ public void setHasReturning(boolean hasReturning) {
4647
this.hasReturinng = hasReturning;
4748
}
4849

50+
public void setHasGenerated(boolean hasGenerated) {
51+
this.hasGenerated = hasGenerated;
52+
}
53+
4954
public boolean hasUpdateCount() {
5055
return (command == QueryCmd.INSERT_UPSERT || command == QueryCmd.UPDATE_REPLACE_DELETE) && !hasReturinng;
5156
}
5257

58+
public boolean hasUpdateWithGenerated() {
59+
return hasGenerated;
60+
}
61+
5362
public boolean hasResults() {
5463
return command == QueryCmd.SELECT || hasReturinng;
5564
}

‎jdbc/src/main/java/tech/ydb/jdbc/query/YdbQuery.java‎

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -13,16 +13,16 @@
1313
* @author Aleksandr Gorshenin
1414
*/
1515
public class YdbQuery {
16-
private final String originQuery;
16+
private final QueryKey key;
1717
private final String preparedYQL;
1818
private final List<QueryStatement> statements;
1919
private final YqlBatcher batcher;
2020

2121
private final QueryType type;
2222
private final boolean isPlainYQL;
2323

24-
YdbQuery(String originQuery, String preparedYQL, List<QueryStatement> stats, YqlBatcher batcher, QueryType type) {
25-
this.originQuery = originQuery;
24+
YdbQuery(QueryKey key, String preparedYQL, List<QueryStatement> stats, YqlBatcher batcher, QueryType type) {
25+
this.key = key;
2626
this.preparedYQL = preparedYQL;
2727
this.statements = stats;
2828
this.type = type;
@@ -48,7 +48,11 @@ public boolean isPlainYQL() {
4848
}
4949

5050
public String getOriginQuery() {
51-
return originQuery;
51+
return key.getQuery();
52+
}
53+
54+
public String getReturning() {
55+
return key.getReturning();
5256
}
5357

5458
public String getPreparedYql() {
@@ -59,7 +63,7 @@ public List<QueryStatement> getStatements() {
5963
return statements;
6064
}
6165

62-
public static YdbQuery parseQuery(String query, YdbQueryProperties opts, YdbTypes types) throws SQLException {
66+
public static YdbQuery parseQuery(QueryKey query, YdbQueryProperties opts, YdbTypes types) throws SQLException {
6367
YdbQueryParser parser = new YdbQueryParser(types, query, opts);
6468
String preparedYQL = parser.parseSQL();
6569

0 commit comments

Comments
 (0)