Skip to content

Commit

Permalink
feat: handle another corner case
Browse files Browse the repository at this point in the history
  • Loading branch information
rajatbhatta committed Aug 30, 2022
1 parent ea4152f commit e2a463e
Show file tree
Hide file tree
Showing 4 changed files with 69 additions and 9 deletions.
Expand Up @@ -480,6 +480,10 @@ private boolean statementStartsWith(String sql, Iterable<String> checkStatements
static final char SLASH = '/';
static final char ASTERISK = '*';
static final char DOLLAR = '$';
static final char SPACE = ' ';
static final char CLOSE_PARENTHESIS = ')';
static final char COMMA = ',';
static final char UNDERSCORE = '_';

/**
* Removes comments from and trims the given sql statement using the dialect of this parser.
Expand Down Expand Up @@ -551,7 +555,7 @@ static int countOccurrencesOf(char c, String string) {
* @return A boolean indicating whether the sql string has a Returning clause or not.
*/
@InternalApi
abstract protected boolean checkReturningClauseInternal(String sql);
protected abstract boolean checkReturningClauseInternal(String sql);

@InternalApi
public boolean checkReturningClause(String sql) {
Expand Down
Expand Up @@ -30,7 +30,7 @@

@InternalApi
public class PostgreSQLStatementParser extends AbstractStatementParser {
private static final Pattern RETURNING_PATTERN = Pattern.compile("[ ')\"]returning[ '(\"]");
private static final Pattern RETURNING_PATTERN = Pattern.compile("returning[ '(\"*]");
private static final Pattern AS_RETURNING_PATTERN = Pattern.compile("[ ')\"]as returning[ '(\"]");
private static final String RETURNING_STRING = "returning";

Expand Down Expand Up @@ -144,7 +144,7 @@ String parseDollarQuotedString(String sql, int index) {
if (c == DOLLAR) {
return tag.toString();
}
if (!Character.isJavaIdentifierPart(c)) {
if (!isValidIdentifierChar(c)) {
break;
}
tag.append(c);
Expand Down Expand Up @@ -292,16 +292,55 @@ private void appendIfNotNull(
}
}

private boolean isValidIdentifierFirstChar(char c) {
return Character.isLetter(c) || c == UNDERSCORE;
}

private boolean isValidIdentifierChar(char c) {
return isValidIdentifierFirstChar(c) || Character.isDigit(c) || c == DOLLAR;
}

private boolean checkCharPrecedingReturning(char ch) {
return (ch == SPACE)
|| (ch == SINGLE_QUOTE)
|| (ch == CLOSE_PARENTHESIS)
|| (ch == DOUBLE_QUOTE)
|| (ch == DOLLAR);
}

private boolean checkCharPrecedingSubstrWithReturning(char ch) {
return (ch == SPACE)
|| (ch == SINGLE_QUOTE)
|| (ch == CLOSE_PARENTHESIS)
|| (ch == DOUBLE_QUOTE)
|| (ch == COMMA);
}

private boolean isReturning(String sql, int index) {
// RETURNING is a reserved keyword in PG, but requires a
// leading AS to be used as column label, to avoid ambiguity.
// We thus check for cases which do not have a leading AS.
// (https://www.postgresql.org/docs/current/sql-keywords-appendix.html)
return (index >= 1)
&& (index + 10 <= sql.length())
&& RETURNING_PATTERN.matcher(sql.substring(index - 1, index + 10)).matches()
&& !((index >= 4)
&& AS_RETURNING_PATTERN.matcher(sql.substring(index - 4, index + 10)).matches());
if (index >= 1) {
if (((index + 10 <= sql.length())
&& RETURNING_PATTERN.matcher(sql.substring(index, index + 10)).matches()
&& !((index >= 4)
&& AS_RETURNING_PATTERN.matcher(sql.substring(index - 4, index + 10)).matches()))) {
if (checkCharPrecedingReturning(sql.charAt(index - 1))) {
return true;
}
// Check for cases where returning clause is part of a substring which starts with an
// invalid first character of an identifier.
// For example,
// insert into t select 2returning *;
int ind = index - 1;
while ((ind >= 0) && !checkCharPrecedingSubstrWithReturning(sql.charAt(ind))) {
ind--;
}
return !isValidIdentifierFirstChar(sql.charAt(ind + 1));
}
}
return false;
}

@InternalApi
Expand Down
Expand Up @@ -30,7 +30,8 @@
@InternalApi
public class SpannerStatementParser extends AbstractStatementParser {

private static final Pattern THEN_RETURN_PATTERN = Pattern.compile("[ `')\"]then return[ `'(\"]");
private static final Pattern THEN_RETURN_PATTERN =
Pattern.compile("[ `')\"]then return[ *`'(\"]");
private static final String THEN_STRING = "then";
private static final String RETURN_STRING = "return";

Expand Down
Expand Up @@ -1360,6 +1360,10 @@ public void testGoogleSQLReturningClause() {
parser
.parse(Statement.of("insert into x (a,b) values (1,2)thenreturn*"))
.hasReturningClause());
assertTrue(
parser
.parse(Statement.of("insert into t(a) select \"x\"then return*"))
.hasReturningClause());
}

@Test
Expand Down Expand Up @@ -1455,6 +1459,18 @@ public void testPostgreSQLReturningClause() {
Statement.of(
"UPDATE x SET y = $returning$returning$returning$ WHERE z = 123 ReTuRnInG *"))
.hasReturningClause());
assertTrue(
parser.parse(Statement.of("insert into t1 select 1 returning*")).hasReturningClause());
assertTrue(
parser.parse(Statement.of("insert into t1 select 2returning*")).hasReturningClause());
assertTrue(
parser.parse(Statement.of("insert into t1 select 10e2returning*")).hasReturningClause());

This comment has been minimized.

Copy link
@olavloite

olavloite Aug 30, 2022

Collaborator

Can we also add a test for insert into t1 select 10.returning* (this is equivalent with insert into t1 select 10.0 returning *)

assertFalse(
parser
.parse(Statement.of("insert into t1 select 'test''returning *'"))
.hasReturningClause());
assertTrue(
parser.parse(Statement.of("insert into t select 2,3returning*")).hasReturningClause());
}

private void assertUnclosedLiteral(String sql) {
Expand Down

0 comments on commit e2a463e

Please sign in to comment.