feat: handle another corner case

Author: rajatbhatta

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AbstractStatementParser.java

@@ -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.
@@ -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) {

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/PostgreSQLStatementParser.java

@@ -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";
 
@@ -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);
@@ -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

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/SpannerStatementParser.java

@@ -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";
 

google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/StatementParserTest.java

@@ -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
@@ -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());
+ 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) {