Skip to content

Commit

Permalink
Added support of convertion IN (?, ?, ?) to IN $list (#87)
Browse files Browse the repository at this point in the history
  • Loading branch information
alex268 authored Nov 13, 2024
2 parents 0e44158 + 1cc496a commit 0f626bb
Show file tree
Hide file tree
Showing 13 changed files with 798 additions and 200 deletions.
16 changes: 11 additions & 5 deletions jdbc/src/main/java/tech/ydb/jdbc/query/QueryStatement.java
Original file line number Diff line number Diff line change
@@ -1,9 +1,11 @@
package tech.ydb.jdbc.query;


import java.util.ArrayList;
import java.util.List;

import tech.ydb.jdbc.common.TypeDescription;
import tech.ydb.jdbc.query.params.JdbcPrm;


/**
*
Expand All @@ -12,7 +14,7 @@
public class QueryStatement {
private final QueryType queryType;
private final QueryCmd command;
private final List<ParamDescription> parameters = new ArrayList<>();
private final List<JdbcPrm.Factory> parameters = new ArrayList<>();
private boolean hasReturinng = false;

public QueryStatement(QueryType custom, QueryType baseType, QueryCmd command) {
Expand All @@ -28,12 +30,16 @@ public QueryCmd getCmd() {
return command;
}

public List<ParamDescription> getParams() {
public boolean hasJdbcParameters() {
return !parameters.isEmpty();
}

public List<JdbcPrm.Factory> getJdbcPrmFactories() {
return parameters;
}

public void addParameter(String name, TypeDescription type) {
this.parameters.add(new ParamDescription(name, type));
public void addJdbcPrmFactory(JdbcPrm.Factory prm) {
this.parameters.add(prm);
}

public void setHasReturning(boolean hasReturning) {
Expand Down
12 changes: 7 additions & 5 deletions jdbc/src/main/java/tech/ydb/jdbc/query/YdbQuery.java
Original file line number Diff line number Diff line change
Expand Up @@ -27,11 +27,11 @@ public class YdbQuery {
this.type = type;
this.batcher = batcher;

boolean hasJdbcParamters = false;
boolean hasJdbcParameters = false;
for (QueryStatement st: statements) {
hasJdbcParamters = hasJdbcParamters || !st.getParams().isEmpty();
hasJdbcParameters = hasJdbcParameters || st.hasJdbcParameters();
}
this.isPlainYQL = !hasJdbcParamters;
this.isPlainYQL = !hasJdbcParameters;
}

public QueryType getType() {
Expand Down Expand Up @@ -59,8 +59,10 @@ public List<QueryStatement> getStatements() {
}

public static YdbQuery parseQuery(String query, YdbQueryProperties opts) throws SQLException {
YdbQueryParser parser = new YdbQueryParser(opts.isDetectQueryType(), opts.isDetectJdbcParameters());
String preparedYQL = parser.parseSQL(query);
YdbQueryParser parser = new YdbQueryParser(
query, opts.isDetectQueryType(), opts.isDetectJdbcParameters(), opts.isReplaceJdbcInByYqlList()
);
String preparedYQL = parser.parseSQL();

QueryType type = null;
YqlBatcher batcher = parser.getYqlBatcher();
Expand Down
168 changes: 136 additions & 32 deletions jdbc/src/main/java/tech/ydb/jdbc/query/YdbQueryParser.java
Original file line number Diff line number Diff line change
Expand Up @@ -7,8 +7,7 @@
import java.util.List;

import tech.ydb.jdbc.YdbConst;
import tech.ydb.jdbc.common.TypeDescription;
import tech.ydb.table.values.PrimitiveType;
import tech.ydb.jdbc.query.params.JdbcPrm;


/**
Expand All @@ -18,13 +17,21 @@
public class YdbQueryParser {
private final boolean isDetectQueryType;
private final boolean isDetectJdbcParameters;
private final boolean isConvertJdbcInToList;

private final List<QueryStatement> statements = new ArrayList<>();
private final YqlBatcher batcher = new YqlBatcher();
private final String origin;
private final StringBuilder parsed;

public YdbQueryParser(boolean isDetectQueryType, boolean isDetectJdbcParameters) {
private int jdbcPrmIndex = 0;

public YdbQueryParser(String origin, boolean isDetectQueryType, boolean isDetectParameters, boolean isConvertIn) {
this.isDetectQueryType = isDetectQueryType;
this.isDetectJdbcParameters = isDetectJdbcParameters;
this.isDetectJdbcParameters = isDetectParameters;
this.isConvertJdbcInToList = isConvertIn;
this.origin = origin;
this.parsed = new StringBuilder(origin.length() + 10);
}

public List<QueryStatement> getStatements() {
Expand Down Expand Up @@ -56,7 +63,7 @@ public QueryType detectQueryType() throws SQLException {
}

@SuppressWarnings("MethodLength")
public String parseSQL(String origin) throws SQLException {
public String parseSQL() throws SQLException {
int fragmentStart = 0;
boolean detectJdbcArgs = false;

Expand All @@ -65,16 +72,13 @@ public String parseSQL(String origin) throws SQLException {

int parenLevel = 0;
int keywordStart = -1;
boolean lastKeywordIsOffsetLimit = false;

char[] chars = origin.toCharArray();

StringBuilder parsed = new StringBuilder(origin.length() + 10);
ArgNameGenerator argNameGenerator = new ArgNameGenerator();

for (int i = 0; i < chars.length; ++i) {
char ch = chars[i];
boolean isInsideKeyword = false;

int keywordEnd = i; // parseSingleQuotes, parseDoubleQuotes, etc move index so we keep old value
switch (ch) {
case '\'': // single-quotes
Expand Down Expand Up @@ -102,6 +106,7 @@ public String parseSQL(String origin) throws SQLException {
case '/': // possibly /* */ style comment
i = parseBlockComment(chars, i);
break;

case '?':
if (detectJdbcArgs && statement != null) {
parsed.append(chars, fragmentStart, i - fragmentStart);
Expand All @@ -110,22 +115,15 @@ public String parseSQL(String origin) throws SQLException {
batcher.readIdentifier(chars, i, 1);
i++; // make sure the coming ? is not treated as a bind
} else {
String binded = argNameGenerator.createArgName(origin);
// force type UInt64 for OFFSET and LIMIT parameters
TypeDescription forcedType = lastKeywordIsOffsetLimit
? TypeDescription.of(PrimitiveType.Uint64)
: null;
statement.addParameter(binded, forcedType);
parsed.append(binded);

String name = nextJdbcPrmName();
statement.addJdbcPrmFactory(JdbcPrm.simplePrm(name));
parsed.append(name);
batcher.readParameter();
}
fragmentStart = i + 1;
}
break;
default:
lastKeywordIsOffsetLimit = lastKeywordIsOffsetLimit && Character.isWhitespace(ch);

if (keywordStart >= 0) {
isInsideKeyword = Character.isJavaIdentifierPart(ch);
break;
Expand All @@ -140,7 +138,6 @@ public String parseSQL(String origin) throws SQLException {


if (keywordStart >= 0 && (!isInsideKeyword || (i == chars.length - 1))) {
lastKeywordIsOffsetLimit = false;
int keywordLength = (isInsideKeyword ? i + 1 : keywordEnd) - keywordStart;

if (statement != null) {
Expand All @@ -151,9 +148,23 @@ public String parseSQL(String origin) throws SQLException {
statement.setHasReturning(true);
}

if (parseOffsetKeyword(chars, keywordStart, keywordLength)
|| parseLimitKeyword(chars, keywordStart, keywordLength)) {
lastKeywordIsOffsetLimit = Character.isWhitespace(ch);
// Process ? after OFFSET and LIMIT
if (i < chars.length && detectJdbcArgs && Character.isWhitespace(ch)) {
if (parseOffsetKeyword(chars, keywordStart, keywordLength)
|| parseLimitKeyword(chars, keywordStart, keywordLength)) {
parsed.append(chars, fragmentStart, i - fragmentStart);
i = parseOffsetLimitParameter(chars, i, statement);
fragmentStart = i;
}
}

// Process IN (?, ?, ... )
if (i < chars.length && detectJdbcArgs && isConvertJdbcInToList) {
if (parseInKeyword(chars, keywordStart, keywordLength)) {
parsed.append(chars, fragmentStart, i - fragmentStart);
i = parseInListParameters(chars, i, statement);
fragmentStart = i;
}
}
} else {
boolean skipped = false;
Expand Down Expand Up @@ -272,18 +283,102 @@ public String parseSQL(String origin) throws SQLException {
return parsed.toString();
}

private static class ArgNameGenerator {
private int index = 0;
private String nextJdbcPrmName() {
while (true) {
jdbcPrmIndex += 1;
String name = YdbConst.AUTO_GENERATED_PARAMETER_PREFIX + jdbcPrmIndex;
if (!origin.contains(name)) {
return name;
}
}
}

public String createArgName(String origin) {
while (true) {
index += 1;
String name = YdbConst.AUTO_GENERATED_PARAMETER_PREFIX + index;
if (!origin.contains(name)) {
return name;
}
private int parseOffsetLimitParameter(char[] query, int offset, QueryStatement st) {
int start = offset;
while (++offset < query.length) {
char ch = query[offset];
switch (ch) {
case '?' :
if (offset + 1 < query.length && query[offset + 1] == '?') {
return start;
}
String name = nextJdbcPrmName();
parsed.append(query, start, offset - start);
parsed.append(name);
st.addJdbcPrmFactory(JdbcPrm.uint64Prm(name));
return offset + 1;
case '-': // possibly -- style comment
offset = parseLineComment(query, offset);
break;
case '/': // possibly /* */ style comment
offset = parseBlockComment(query, offset);
break;
default:
if (!Character.isWhitespace(query[offset])) {
return start;
}
break;
}
}

return start;
}

private int parseInListParameters(char[] query, int offset, QueryStatement st) {
int start = offset;
int listStartedAt = -1;
int listSize = 0;
boolean waitPrm = false;
while (offset < query.length) {
char ch = query[offset];
switch (ch) {
case '(': // start of list
if (listStartedAt >= 0) {
return start;
}
listStartedAt = offset;
waitPrm = true;
break;
case ',':
if (listStartedAt < 0 || waitPrm) {
return start;
}
waitPrm = true;
break;
case '?' :
if (!waitPrm || (offset + 1 < query.length && query[offset + 1] == '?')) {
return start;
}
listSize++;
waitPrm = false;
break;
case ')':
if (waitPrm || listSize == 0 || listStartedAt < 0) {
return start;
}

String name = nextJdbcPrmName();
parsed.append(query, start, listStartedAt - start);
parsed.append(' '); // add extra space to avoid IN$jpN
parsed.append(name);
st.addJdbcPrmFactory(JdbcPrm.inListOrm(name, listSize));
return offset + 1;
case '-': // possibly -- style comment
offset = parseLineComment(query, offset);
break;
case '/': // possibly /* */ style comment
offset = parseBlockComment(query, offset);
break;
default:
if (!Character.isWhitespace(query[offset])) {
return start;
}
break;
}
offset++;
}

return start;
}

private static int parseSingleQuotes(final char[] query, int offset) {
Expand Down Expand Up @@ -553,4 +648,13 @@ private static boolean parseLimitKeyword(char[] query, int offset, int length) {
&& (query[offset + 3] | 32) == 'i'
&& (query[offset + 4] | 32) == 't';
}

private static boolean parseInKeyword(char[] query, int offset, int length) {
if (length != 2) {
return false;
}

return (query[offset] | 32) == 'i'
&& (query[offset + 1] | 32) == 'n';
}
}
Loading

0 comments on commit 0f626bb

Please sign in to comment.