Added forced UInt64 type for OFFSET & LIMIT

Author: alex268

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

@@ -7,6 +7,8 @@
 import java.util.List;
 
 import tech.ydb.jdbc.YdbConst;
+import tech.ydb.jdbc.common.TypeDescription;
+import tech.ydb.table.values.PrimitiveType;
 
 
 /**
@@ -48,6 +50,7 @@ public QueryType detectQueryType() throws SQLException {
         return type != null ? type : QueryType.DATA_QUERY;
     }
 
+    @SuppressWarnings("MethodLength")
     public String parseSQL(String origin) throws SQLException {
         this.statements.clear();
 
@@ -59,6 +62,7 @@ public String parseSQL(String origin) throws SQLException {
 
         int parenLevel = 0;
         int keywordStart = -1;
+        boolean lastKeywordIsOffsetLimit = false;
 
         char[] chars = origin.toCharArray();
 
@@ -100,7 +104,11 @@ public String parseSQL(String origin) throws SQLException {
                             i++; // make sure the coming ? is not treated as a bind
                         } else {
                             String binded = argNameGenerator.createArgName(origin);
-                            currStatement.addParameter(binded, null);
+                            // force type UInt64 for OFFSET and LIMIT parameters
+                            TypeDescription type = lastKeywordIsOffsetLimit
+                                    ? TypeDescription.of(PrimitiveType.Uint64)
+                                    : null;
+                            currStatement.addParameter(binded, type);
                             parsed.append(binded);
                         }
                         fragmentStart = i + 1;
@@ -120,12 +128,17 @@ public String parseSQL(String origin) throws SQLException {
             }
 
             if (keywordStart >= 0 && (!isInsideKeyword || (i == chars.length - 1))) {
+                lastKeywordIsOffsetLimit = false;
 
                 if (currStatement != null) {
                     // Detect RETURNING keyword
                     if (parenLevel == 0 && parseReturningKeyword(chars, keywordStart)) {
                         currStatement.setHasReturning(true);
                     }
+
+                    if (parseOffsetKeyword(chars, keywordStart) || parseLimitKeyword(chars, keywordStart)) {
+                        lastKeywordIsOffsetLimit = true;
+                    }
                 } else {
                     // Detecting type of statement by the first keyword
                     currStatement = new QueryStatement(QueryType.UNKNOWN, QueryCmd.UNKNOWN);
@@ -179,7 +192,6 @@ public String parseSQL(String origin) throws SQLException {
                             || parseDropKeyword(chars, i)) {
                         currStatement = new QueryStatement(QueryType.SCHEME_QUERY, QueryCmd.CREATE_ALTER_DROP);
                         statements.add(currStatement);
-                        break;
                     }
                 }
 
@@ -441,4 +453,29 @@ private static boolean parseReturningKeyword(char[] query, int offset) {
                 && (query[offset + 7] | 32) == 'n'
                 && (query[offset + 8] | 32) == 'g';
     }
+
+    private static boolean parseOffsetKeyword(char[] query, int offset) {
+        if (query.length < (offset + 6)) {
+            return false;
+        }
+
+        return (query[offset] | 32) == 'o'
+                && (query[offset + 1] | 32) == 'f'
+                && (query[offset + 2] | 32) == 'f'
+                && (query[offset + 3] | 32) == 's'
+                && (query[offset + 4] | 32) == 'e'
+                && (query[offset + 5] | 32) == 't';
+    }
+
+    private static boolean parseLimitKeyword(char[] query, int offset) {
+        if (query.length < (offset + 5)) {
+            return false;
+        }
+
+        return (query[offset] | 32) == 'l'
+                && (query[offset + 1] | 32) == 'i'
+                && (query[offset + 2] | 32) == 'm'
+                && (query[offset + 3] | 32) == 'i'
+                && (query[offset + 4] | 32) == 't';
+    }
 }

jdbc/src/test/java/tech/ydb/jdbc/impl/YdbConnectionImplTest.java

@@ -779,6 +779,25 @@ public void testReturingStatements() throws SQLException {
         }
     }
 
+    @Test
+    public void testOffsetLimit() throws SQLException {
+        String query = QUERIES.withTableName("SELECT * FROM #tableName ORDER BY key LIMIT ? OFFSET ?");
+
+        try (PreparedStatement ps = jdbc.connection().prepareStatement(query)) {
+            ps.setInt(1, 0);
+            ps.setInt(2, 0);
+            try (ResultSet rs = ps.executeQuery()) {
+                Assertions.assertFalse(rs.next());
+            }
+
+            ps.setLong(1, 5);
+            ps.setLong(2, 5);
+            try (ResultSet rs = ps.executeQuery()) {
+                Assertions.assertFalse(rs.next());
+            }
+        }
+    }
+
     @Test
     public void testWarningsInQuery() throws SQLException {
         String createTempTable = QUERIES.withTableName(

jdbc/src/test/java/tech/ydb/jdbc/impl/YdbQueryConnectionImplTest.java

@@ -782,6 +782,25 @@ public void testReturingStatements() throws SQLException {
         }
     }
 
+    @Test
+    public void testOffsetLimit() throws SQLException {
+        String query = QUERIES.withTableName("SELECT * FROM #tableName ORDER BY key LIMIT ? OFFSET ?");
+
+        try (PreparedStatement ps = jdbc.connection().prepareStatement(query)) {
+            ps.setInt(1, 0);
+            ps.setInt(2, 0);
+            try (ResultSet rs = ps.executeQuery()) {
+                Assertions.assertFalse(rs.next());
+            }
+
+            ps.setLong(1, 5);
+            ps.setLong(2, 5);
+            try (ResultSet rs = ps.executeQuery()) {
+                Assertions.assertFalse(rs.next());
+            }
+        }
+    }
+
     @Test
     public void testWarningsInQuery() throws SQLException {
         String createTempTable = QUERIES.withTableName(

jdbc/src/test/java/tech/ydb/jdbc/query/YdbQueryParserTest.java

@@ -7,6 +7,8 @@
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.CsvSource;
 
+import tech.ydb.table.values.PrimitiveType;
+
 
 /**
  *
@@ -84,4 +86,29 @@ public void scanSelectWithKeyTest() throws SQLException {
 
         Assertions.assertEquals(QueryType.SCAN_QUERY, parser.detectQueryType());
     }
+
+    @Test
+    public void offsetParameterTest() throws SQLException {
+        String query = ""
+                + "select * from test_table where true=true -- test request\n"
+                + " offset /* comment */ ? limit 20";
+
+        YdbQueryParser parser = new YdbQueryParser(true, true);
+        String parsed = parser.parseSQL(query);
+        Assertions.assertEquals(""
+                + "select * from test_table where true=true -- test request\n"
+                + " offset /* comment */ $jp1 limit 20",
+                parsed);
+
+        Assertions.assertEquals(1, parser.getStatements().size());
+
+        QueryStatement statement = parser.getStatements().get(0);
+        Assertions.assertEquals(QueryType.DATA_QUERY, statement.getType());
+        Assertions.assertEquals(1, statement.getParams().size());
+
+        ParamDescription prm1 = statement.getParams().get(0);
+        Assertions.assertEquals("$jp1", prm1.name());
+        Assertions.assertNotNull(prm1.type());
+        Assertions.assertEquals(PrimitiveType.Uint64, prm1.type().ydbType());
+    }
 }