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