Rewrite query parser to detect every keyword

Author: alex268

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

@@ -18,7 +18,7 @@
 import tech.ydb.jdbc.common.FixedResultSetFactory;
 import tech.ydb.jdbc.context.YdbValidator;
 import tech.ydb.jdbc.query.ExplainedQuery;
-import tech.ydb.jdbc.query.QueryExpression;
+import tech.ydb.jdbc.query.QueryStatement;
 import tech.ydb.jdbc.query.YdbQuery;
 import tech.ydb.jdbc.settings.YdbOperationProperties;
 import tech.ydb.table.query.Params;
@@ -168,7 +168,7 @@ protected boolean updateState(List<YdbResult> results) {
     protected List<YdbResult> executeSchemeQuery(YdbQuery query) throws SQLException {
         connection.executeSchemeQuery(query, validator);
 
-        int expressionsCount = query.getExpressions().isEmpty() ? 1 : query.getExpressions().size();
+        int expressionsCount = query.getStatements().isEmpty() ? 1 : query.getStatements().size();
         List<YdbResult> results = new ArrayList<>();
         for (int i = 0; i < expressionsCount; i++) {
             results.add(NO_UPDATED);
@@ -200,17 +200,17 @@ protected List<YdbResult> executeDataQuery(YdbQuery query, Params params) throws
 
         List<YdbResult> results = new ArrayList<>();
         int idx = 0;
-        for (QueryExpression exp: query.getExpressions()) {
+        for (QueryStatement exp: query.getStatements()) {
             if (exp.isDDL()) {
                 results.add(NO_UPDATED);
                 continue;
             }
-            if (!exp.hasResults()) {
+            if (exp.hasUpdateCount()) {
                 results.add(HAS_UPDATED);
                 continue;
             }
 
-            if (idx < resultSets.size())  {
+            if (exp.hasResults() && idx < resultSets.size()) {
                 ResultSetReader rs = resultSets.get(idx);
                 if (failOnTruncatedResult && rs.isTruncated()) {
                     String msg = String.format(YdbConst.RESULT_IS_TRUNCATED, idx, rs.getRowCount());

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

@@ -5,8 +5,9 @@
  * @author Aleksandr Gorshenin
  */
 public enum QueryCmd {
+    UNKNOWN,
     SELECT,
     CREATE_ALTER_DROP,
     INSERT_UPSERT,
-    UPDATE_REPLACE_REMOVE
+    UPDATE_REPLACE_DELETE
 }

jdbc/src/main/java/tech/ydb/jdbc/query/QueryExpression.java renamed to jdbc/src/main/java/tech/ydb/jdbc/query/QueryStatement.java

@@ -7,12 +7,13 @@
  *
  * @author Aleksandr Gorshenin
  */
-public class QueryExpression {
+public class QueryStatement {
     private final QueryType queryType;
     private final QueryCmd command;
     private final List<String> paramNames = new ArrayList<>();
+    private boolean hasReturinng = false;
 
-    public QueryExpression(QueryType type, QueryCmd command) {
+    public QueryStatement(QueryType type, QueryCmd command) {
         this.queryType = type;
         this.command = command;
     }
@@ -29,8 +30,16 @@ public void addParamName(String name) {
         this.paramNames.add(name);
     }
 
+    public void setHasReturning(boolean hasReturning) {
+        this.hasReturinng = hasReturning;
+    }
+
+    public boolean hasUpdateCount() {
+        return (command == QueryCmd.INSERT_UPSERT || command == QueryCmd.UPDATE_REPLACE_DELETE); // && !hasReturinng;
+    }
+
     public boolean hasResults() {
-        return command == QueryCmd.SELECT;
+        return command == QueryCmd.SELECT; // || hasReturinng;
     }
 
     public boolean isDDL() {

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

@@ -1,6 +1,8 @@
 package tech.ydb.jdbc.query;
 
 public enum QueryType {
+    UNKNOWN,
+
     // DDL
     SCHEME_QUERY,
 

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

@@ -4,7 +4,6 @@
 
 import java.sql.SQLDataException;
 import java.sql.SQLException;
-import java.sql.SQLFeatureNotSupportedException;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
@@ -19,16 +18,22 @@
  * @author Aleksandr Gorshenin
  */
 public class YdbQuery {
-    private final ParsedQuery query;
+    private final String originSQL;
+    private final String preparedYQL;
+    private final List<QueryStatement> statements;
+
     private final QueryType type;
     private final boolean isAutoDeclare;
     private final List<String> paramNames = new ArrayList<>();
 
-    YdbQuery(ParsedQuery query, QueryType type, boolean isAutoDeclare) {
-        this.query = query;
+    YdbQuery(String originSQL, String preparedYQL, List<QueryStatement> stats, QueryType type, boolean isAutoDeclare) {
+        this.originSQL = originSQL;
+        this.preparedYQL = preparedYQL;
+        this.statements = stats;
         this.type = type;
         this.isAutoDeclare = isAutoDeclare;
-        for (QueryExpression expression : query.getExpressions()) {
+
+        for (QueryStatement expression : stats) {
             paramNames.addAll(expression.getParamNames());
         }
     }
@@ -38,11 +43,11 @@ public QueryType getType() {
     }
 
     public String getOriginSQL() {
-        return query.getOriginSQL();
+        return originSQL;
     }
 
-    public List<QueryExpression> getExpressions() {
-        return query.getExpressions();
+    public List<QueryStatement> getStatements() {
+        return statements;
     }
 
     public boolean hasFreeParams() {
@@ -55,15 +60,15 @@ public List<String> getFreeParams() {
 
     public String withParams(Params params) throws SQLException {
         if (paramNames.isEmpty()) {
-            return query.getPreparedYQL();
+            return preparedYQL;
         }
 
         if (params == null) {
             if (!paramNames.isEmpty() && isAutoDeclare) {
                 // Comment in place where must be declare section
-                return "-- DECLARE " + paramNames.size() + " PARAMETERS\n" + query.getPreparedYQL();
+                return "-- DECLARE " + paramNames.size() + " PARAMETERS\n" + preparedYQL;
             }
-            return query.getPreparedYQL();
+            return preparedYQL;
         }
 
         StringBuilder yql = new StringBuilder();
@@ -84,30 +89,19 @@ public String withParams(Params params) throws SQLException {
             }
         }
 
-        yql.append(query.getPreparedYQL());
+        yql.append(preparedYQL);
         return yql.toString();
     }
 
-    public static YdbQuery parseQuery(String queryText, YdbQueryProperties opts) throws SQLException {
-        ParsedQuery query = ParsedQuery.parse(queryText, opts);
+    public static YdbQuery parseQuery(String query, YdbQueryProperties opts) throws SQLException {
+        YdbQueryParser parser = new YdbQueryParser(opts.isDetectQueryType(), opts.isDetectJdbcParameters());
+        String preparedYQL = parser.parseSQL(query);
+
         QueryType type = opts.getForcedQueryType();
-        if (type != null) {
-            return new YdbQuery(query, type, opts.isDeclareJdbcParameters());
-        }
-        for (QueryExpression exp: query.getExpressions()) {
-            if (type == null) {
-                type = exp.getType();
-            } else {
-                if (type != exp.getType()) {
-                    throw new SQLFeatureNotSupportedException(
-                            YdbConst.MULTI_TYPES_IN_ONE_QUERY + type + ", " + exp.getType()
-                    );
-                }
-            }
-        }
         if (type == null) {
-            type = QueryType.DATA_QUERY;
+            type = parser.detectQueryType();
         }
-        return new YdbQuery(query, type, opts.isDeclareJdbcParameters());
+
+        return new YdbQuery(query, preparedYQL, parser.getStatements(), type, opts.isDeclareJdbcParameters());
     }
 }