Added support of generated keys

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

@@ -66,7 +66,8 @@ public final class YdbConst {
     public static final String INVALID_ROW = "Current row index is out of bounds: ";
     public static final String BULKS_UNSUPPORTED = "BULK mode is available only for prepared statement with one UPSERT";
     public static final String INVALID_BATCH_COLUMN = "Cannot prepared batch request: cannot find a column";
-    public static final String BULKS_DESCRIBE_ERROR = "Cannot parse BULK upsert: ";
+    public static final String BULK_DESCRIBE_ERROR = "Cannot parse BULK upsert: ";
+    public static final String BULK_NOT_SUPPORT_RETURNING = "BULK query doesn't support RETURNING";
     public static final String METADATA_RS_UNSUPPORTED_IN_PS = "ResultSet metadata is not supported " +
             "in prepared statements";
     public static final String CANNOT_UNWRAP_TO = "Cannot unwrap to ";

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

@@ -34,15 +34,24 @@ public class StaticQueryResult implements YdbQueryResult {
     private static class ExpressionResult {
         private final int updateCount;
         private final YdbResultSet resultSet;
+        private final YdbResultSet generatedKeys;
 
         ExpressionResult(int updateCount) {
             this.updateCount = updateCount;
             this.resultSet = null;
+            this.generatedKeys = null;
         }
 
-        ExpressionResult(YdbResultSet resultSet) {
+        ExpressionResult(int updateCount, YdbResultSet keys) {
+            this.updateCount = updateCount;
+            this.resultSet = null;
+            this.generatedKeys = keys;
+        }
+
+        ExpressionResult(YdbResultSet result) {
             this.updateCount = -1;
-            this.resultSet = resultSet;
+            this.resultSet = result;
+            this.generatedKeys = null;
         }
     }
 
@@ -59,6 +68,13 @@ public StaticQueryResult(YdbQuery query, List<YdbResultSet> list) {
                 results.add(NO_UPDATED);
                 continue;
             }
+
+            if (exp.hasUpdateWithGenerated() && idx < list.size()) {
+                results.add(new ExpressionResult(1, list.get(idx)));
+                idx++;
+                continue;
+            }
+
             if (exp.hasUpdateCount()) {
                 results.add(HAS_UPDATED);
                 continue;
@@ -124,6 +140,14 @@ public YdbResultSet getCurrentResultSet()  {
         return results.get(resultIndex).resultSet;
     }
 
+    @Override
+    public YdbResultSet getGeneratedKeys()  {
+        if (results == null || resultIndex >= results.size()) {
+            return null;
+        }
+        return results.get(resultIndex).generatedKeys;
+    }
+
     @Override
     public boolean getMoreResults(int current) throws SQLException {
         if (results == null || resultIndex >= results.size()) {

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

@@ -219,6 +219,12 @@ public YdbResultSet getCurrentResultSet() throws SQLException {
         }
     }
 
+    @Override
+    public YdbResultSet getGeneratedKeys() throws SQLException {
+        // TODO: Implement
+        return null;
+    }
+
     @Override
     public boolean hasResultSets() throws SQLException {
         if (resultIndex >= resultIndexes.length) {

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

@@ -29,6 +29,7 @@
 import tech.ydb.jdbc.exception.ExceptionFactory;
 import tech.ydb.jdbc.impl.YdbTracerImpl;
 import tech.ydb.jdbc.impl.YdbTracerNone;
+import tech.ydb.jdbc.query.QueryKey;
 import tech.ydb.jdbc.query.QueryType;
 import tech.ydb.jdbc.query.YdbPreparedQuery;
 import tech.ydb.jdbc.query.YdbQuery;
@@ -83,7 +84,7 @@ public class YdbContext implements AutoCloseable {
     private final String prefixPath;
     private final String prefixPragma;
 
-    private final Cache<String, YdbQuery> queriesCache;
+    private final Cache<QueryKey, YdbQuery> queriesCache;
     private final Cache<String, QueryStat> statsCache;
     private final Cache<String, Map<String, Type>> queryParamsCache;
     private final Cache<String, TableDescription> tableDescribeCache;
@@ -348,22 +349,21 @@ public <T extends BaseRequestSettings.BaseBuilder<T>> T withRequestTimeout(T bui
         return builder.withRequestTimeout(operation);
     }
 
-    public YdbQuery parseYdbQuery(String sql) throws SQLException {
-        return YdbQuery.parseQuery(sql, queryOptions, types);
+    public YdbQuery parseYdbQuery(String query) throws SQLException {
+        return YdbQuery.parseQuery(new QueryKey(query), queryOptions, types);
     }
 
-    public YdbQuery findOrParseYdbQuery(String sql) throws SQLException {
+    public YdbQuery findOrParseYdbQuery(QueryKey key) throws SQLException {
         if (queriesCache == null) {
-            return parseYdbQuery(sql);
+            return YdbQuery.parseQuery(key, queryOptions, types);
         }
 
-        YdbQuery cached = queriesCache.getIfPresent(sql);
+        YdbQuery cached = queriesCache.getIfPresent(key);
         if (cached == null) {
-            cached = parseYdbQuery(sql);
-            queriesCache.put(sql, cached);
+            cached = YdbQuery.parseQuery(key, queryOptions, types);
+            queriesCache.put(key, cached);
         }
 
-
         return cached;
     }
 
@@ -432,16 +432,19 @@ public YdbPreparedQuery findOrPrepareParams(YdbQuery query, YdbPrepareMode mode)
                     tableDescribeCache.put(tablePath, description);
                 } else {
                     if (type == QueryType.BULK_QUERY) {
-                        throw new SQLException(YdbConst.BULKS_DESCRIBE_ERROR + result.getStatus());
+                        throw new SQLException(YdbConst.BULK_DESCRIBE_ERROR + result.getStatus());
                     }
                 }
             }
             if (type == QueryType.BULK_QUERY) {
+                if (query.getReturning() != null) {
+                    throw new SQLException(YdbConst.BULK_NOT_SUPPORT_RETURNING);
+                }
                 return BulkUpsertQuery.build(types, tablePath, query.getYqlBatcher().getColumns(), description);
             }
 
             if (description != null) {
-                BatchedQuery params = BatchedQuery.createAutoBatched(types, query.getYqlBatcher(), description);
+                BatchedQuery params = BatchedQuery.createAutoBatched(types, query, description);
                 if (params != null) {
                     return params;
                 }

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

@@ -244,7 +244,7 @@ public void setCursorName(String name) throws SQLFeatureNotSupportedException {
 
     @Override
     public ResultSet getGeneratedKeys() throws SQLException {
-        return null; // --
+        return state.getGeneratedKeys();
     }
 
     @Override

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

@@ -29,6 +29,7 @@
 import tech.ydb.jdbc.context.YdbContext;
 import tech.ydb.jdbc.context.YdbExecutor;
 import tech.ydb.jdbc.context.YdbValidator;
+import tech.ydb.jdbc.query.QueryKey;
 import tech.ydb.jdbc.query.YdbPreparedQuery;
 import tech.ydb.jdbc.query.YdbQuery;
 
@@ -213,15 +214,15 @@ public YdbStatement createStatement(int resultSetType, int resultSetConcurrency,
     }
 
     @Override
-    public YdbPreparedStatement prepareStatement(String origSql, int resultSetType, int resultSetConcurrency,
+    public YdbPreparedStatement prepareStatement(String sql, int resultSetType, int resultSetConcurrency,
             int resultSetHoldability) throws SQLException {
         checkStatementParams(resultSetType, resultSetConcurrency, resultSetHoldability);
-        return prepareStatement(origSql, resultSetType, null, YdbPrepareMode.AUTO);
+        return prepareStatement(new QueryKey(sql), resultSetType, YdbPrepareMode.AUTO);
     }
 
     @Override
     public YdbPreparedStatement prepareStatement(String sql, YdbPrepareMode mode) throws SQLException {
-        return prepareStatement(sql, ResultSet.TYPE_SCROLL_INSENSITIVE, null, mode);
+        return prepareStatement(new QueryKey(sql), ResultSet.TYPE_SCROLL_INSENSITIVE, mode);
     }
 
     @Override
@@ -245,18 +246,18 @@ public YdbPreparedStatement prepareStatement(String sql, String[] columnNames) t
         if (columnNames == null || columnNames.length == 0) {
             return prepareStatement(sql);
         }
-        return prepareStatement(sql, ResultSet.TYPE_FORWARD_ONLY, columnNames, YdbPrepareMode.AUTO);
+        return prepareStatement(new QueryKey(sql, columnNames), ResultSet.TYPE_FORWARD_ONLY, YdbPrepareMode.AUTO);
     }
 
-    private YdbPreparedStatement prepareStatement(String sql, int type, String[] columnNames, YdbPrepareMode mode)
+    private YdbPreparedStatement prepareStatement(QueryKey key, int resultSetType, YdbPrepareMode mode)
             throws SQLException {
 
         validator.clearWarnings();
         ctx.getTracer().trace("prepare statement");
-        YdbQuery query = ctx.findOrParseYdbQuery(sql);
+        YdbQuery query = ctx.findOrParseYdbQuery(key);
         YdbPreparedQuery params = ctx.findOrPrepareParams(query, mode);
         ctx.getTracer().trace("create prepared statement");
-        return new YdbPreparedStatementImpl(this, query, params, type);
+        return new YdbPreparedStatementImpl(this, query, params, resultSetType);
     }
 
     @Override

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

@@ -20,6 +20,11 @@ public YdbResultSet getCurrentResultSet() throws SQLException {
             return null;
         }
 
+        @Override
+        public YdbResultSet getGeneratedKeys() throws SQLException {
+            return null;
+        }
+
         @Override
         public boolean hasResultSets() throws SQLException {
             return false;
@@ -36,6 +41,7 @@ public void close() throws SQLException { }
 
     int getUpdateCount() throws SQLException;
     YdbResultSet getCurrentResultSet() throws SQLException;
+    YdbResultSet getGeneratedKeys() throws SQLException;
 
     boolean hasResultSets() throws SQLException;
     boolean getMoreResults(int current) throws SQLException;

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

@@ -0,0 +1,72 @@
+package tech.ydb.jdbc.query;
+
+import java.util.Objects;
+
+/**
+ *
+ * @author Aleksandr Gorshenin
+ */
+public class QueryKey {
+    private final String query;
+    private final String returning;
+
+    public QueryKey(String query) {
+        this.query = query;
+        this.returning = null;
+    }
+
+    public QueryKey(String query, String[] columnNames) {
+        this.query = query;
+        this.returning = buildReturning(columnNames);
+    }
+
+    public String getQuery() {
+        return query;
+    }
+
+    public String getReturning() {
+        return returning;
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(query, returning);
+    }
+
+    @Override
+    public boolean equals(Object obj) {
+        if (this == obj) {
+            return true;
+        }
+        if (obj == null || getClass() != obj.getClass()) {
+            return false;
+        }
+
+        QueryKey other = (QueryKey) obj;
+        return Objects.equals(this.query, other.query) && Objects.equals(this.returning, other.returning);
+    }
+
+    private static String buildReturning(String[] columnNames) {
+        StringBuilder sb = new StringBuilder();
+        sb.append("RETURNING ");
+        if (columnNames.length == 1 && columnNames[0].charAt(0) == '*') {
+            sb.append('*');
+            return sb.toString();
+        }
+
+        for (int col = 0; col < columnNames.length; col++) {
+            String columnName = columnNames[col];
+            if (col > 0) {
+                sb.append(", ");
+            }
+
+            if (columnName.charAt(0) == '`') {
+                sb.append(columnName);
+            } else {
+                sb.append('`').append(columnName).append('`');
+            }
+        }
+
+        return sb.toString();
+    }
+}

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

@@ -16,6 +16,7 @@ public class QueryStatement {
     private final QueryCmd command;
     private final List<JdbcPrm.Factory> parameters = new ArrayList<>();
     private boolean hasReturinng = false;
+    private boolean hasGenerated = false;
 
     public QueryStatement(QueryType custom, QueryType baseType, QueryCmd command) {
         this.queryType = custom != null ? custom : baseType;
@@ -46,10 +47,18 @@ public void setHasReturning(boolean hasReturning) {
         this.hasReturinng = hasReturning;
     }
 
+    public void setHasGenerated(boolean hasGenerated) {
+        this.hasGenerated = hasGenerated;
+    }
+
     public boolean hasUpdateCount() {
         return (command == QueryCmd.INSERT_UPSERT || command == QueryCmd.UPDATE_REPLACE_DELETE) && !hasReturinng;
     }
 
+    public boolean hasUpdateWithGenerated() {
+        return hasGenerated;
+    }
+
     public boolean hasResults() {
         return command == QueryCmd.SELECT || hasReturinng;
     }

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

@@ -13,16 +13,16 @@
  * @author Aleksandr Gorshenin
  */
 public class YdbQuery {
-    private final String originQuery;
+    private final QueryKey key;
     private final String preparedYQL;
     private final List<QueryStatement> statements;
     private final YqlBatcher batcher;
 
     private final QueryType type;
     private final boolean isPlainYQL;
 
-    YdbQuery(String originQuery, String preparedYQL, List<QueryStatement> stats, YqlBatcher batcher, QueryType type) {
-        this.originQuery = originQuery;
+    YdbQuery(QueryKey key, String preparedYQL, List<QueryStatement> stats, YqlBatcher batcher, QueryType type) {
+        this.key = key;
         this.preparedYQL = preparedYQL;
         this.statements = stats;
         this.type = type;
@@ -48,7 +48,11 @@ public boolean isPlainYQL() {
     }
 
     public String getOriginQuery() {
-        return originQuery;
+        return key.getQuery();
+    }
+
+    public String getReturning() {
+        return key.getReturning();
     }
 
     public String getPreparedYql() {
@@ -59,7 +63,7 @@ public List<QueryStatement> getStatements() {
         return statements;
     }
 
-    public static YdbQuery parseQuery(String query, YdbQueryProperties opts, YdbTypes types) throws SQLException {
+    public static YdbQuery parseQuery(QueryKey query, YdbQueryProperties opts, YdbTypes types) throws SQLException {
         YdbQueryParser parser = new YdbQueryParser(types, query, opts);
         String preparedYQL = parser.parseSQL();