Repository: GoogleCloudPlatform/pgadapter
Pull request #330 (2022-08-21 09:52:15)

refactor: combine StatementParser and SimpleParser
Combines the StatementParser and SimpleParser into one parser to reduce code duplication
and to make the SimpleParser more complete, as it now also supports skipping comments
and dollar-quoted strings.

codecov[bot]: # [Codecov](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/330?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) Report
> Merging [#330](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/330?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) (713eaa2) into [postgresql-dialect](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/commit/a5e563419cd0cc83557719991469ff606429a976?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) (a5e5634) will **decrease** coverage by `0.15%`.
> The diff coverage is `83.58%`.

```diff
@@                   Coverage Diff                    @@
##             postgresql-dialect     #330      +/-   ##
========================================================
- Coverage                 83.54%   83.39%   -0.16%     
+ Complexity                 1505     1501       -4     
========================================================
  Files                       113      112       -1     
  Lines                      5203     5214      +11     
  Branches                    674      674              
========================================================
+ Hits                       4347     4348       +1     
- Misses                      639      643       +4     
- Partials                    217      223       +6     
```

| Flag | Coverage Δ | |
|---|---|---|
| all_tests | `83.39% <83.58%> (-0.16%)` | :arrow_down: |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#carryforward-flags-in-the-pull-request-comment) to find out more.

| [Impacted Files](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/330?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) | Coverage Δ | |
|---|---|---|
| [...oud/spanner/pgadapter/statements/SimpleParser.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/330/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3N0YXRlbWVudHMvU2ltcGxlUGFyc2VyLmphdmE=) | `88.16% <82.53%> (-6.91%)` | :arrow_down: |
| [...oud/spanner/pgadapter/commands/DynamicCommand.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/330/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL2NvbW1hbmRzL0R5bmFtaWNDb21tYW5kLmphdmE=) | `92.30% <100.00%> (+0.64%)` | :arrow_up: |
| [...panner/pgadapter/statements/BackendConnection.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/330/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3N0YXRlbWVudHMvQmFja2VuZENvbm5lY3Rpb24uamF2YQ==) | `90.63% <100.00%> (-0.03%)` | :arrow_down: |
| [...er/pgadapter/statements/IntermediateStatement.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/330/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3N0YXRlbWVudHMvSW50ZXJtZWRpYXRlU3RhdGVtZW50LmphdmE=) | `95.34% <100.00%> (ø)` | |
| [...ner/pgadapter/statements/SimpleQueryStatement.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/330/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3N0YXRlbWVudHMvU2ltcGxlUXVlcnlTdGF0ZW1lbnQuamF2YQ==) | `93.54% <100.00%> (+0.10%)` | :arrow_up: |
| [...d/spanner/pgadapter/wireprotocol/ParseMessage.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/330/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3dpcmVwcm90b2NvbC9QYXJzZU1lc3NhZ2UuamF2YQ==) | `81.13% <100.00%> (ø)` | |

:mega: We’re building smart automated test selection to slash your CI/CD build times. [Learn more](https://about.codecov.io/iterative-testing/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform)


Patch:
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/commands/DynamicCommand.java b/src/main/java/com/google/cloud/spanner/pgadapter/commands/DynamicCommand.java
index 8800e262..b7710859 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/commands/DynamicCommand.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/commands/DynamicCommand.java
@@ -16,7 +16,6 @@ package com.google.cloud.spanner.pgadapter.commands;
 
 import com.google.api.core.InternalApi;
 import com.google.cloud.spanner.pgadapter.metadata.DynamicCommandMetadata;
-import com.google.cloud.spanner.pgadapter.utils.StatementParser;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.regex.Pattern;
@@ -51,9 +50,18 @@ public class DynamicCommand extends Command {
     } else {
       List<String> matcherList = new ArrayList<>();
       for (String argumentPosition : this.metadata.getMatcherOrder()) {
-        matcherList.add(StatementParser.singleQuoteEscape(this.matcher.group(argumentPosition)));
+        matcherList.add(singleQuoteEscape(this.matcher.group(argumentPosition)));
       }
       return String.format(this.metadata.getOutputPattern(), matcherList.toArray());
     }
   }
+
+  /**
+   * Simple method to escape SQL. Ultimately it is preferred that a user uses PreparedStatements but
+   * for the case of psql emulation, we apply this to provide a simple layer of protection to the
+   * user. This method simply duplicates all single quotes (i.e. ' becomes '').
+   */
+  public static String singleQuoteEscape(String sql) {
+    return sql.replace("'", "''");
+  }
 }
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/statements/BackendConnection.java b/src/main/java/com/google/cloud/spanner/pgadapter/statements/BackendConnection.java
index aa3bc91c..ee1e2484 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/statements/BackendConnection.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/statements/BackendConnection.java
@@ -14,6 +14,7 @@
 
 package com.google.cloud.spanner.pgadapter.statements;
 
+import static com.google.cloud.spanner.pgadapter.statements.SimpleParser.isCommand;
 import static com.google.cloud.spanner.pgadapter.wireprotocol.QueryMessage.COPY;
 
 import com.google.api.core.InternalApi;
@@ -52,7 +53,6 @@ import com.google.cloud.spanner.pgadapter.statements.SimpleParser.TableOrIndexNa
 import com.google.cloud.spanner.pgadapter.statements.local.LocalStatement;
 import com.google.cloud.spanner.pgadapter.utils.CopyDataReceiver;
 import com.google.cloud.spanner.pgadapter.utils.MutationWriter;
-import com.google.cloud.spanner.pgadapter.utils.StatementParser;
 import com.google.cloud.spanner.pgadapter.wireoutput.ReadyResponse;
 import com.google.cloud.spanner.pgadapter.wireoutput.ReadyResponse.Status;
 import com.google.common.annotations.VisibleForTesting;
@@ -771,8 +771,7 @@ public class BackendConnection {
             statement ->
                 statement.parsedStatement.getType() == StatementType.UPDATE
                     || statement.parsedStatement.getType() == StatementType.UNKNOWN
-                        && StatementParser.isCommand(
-                            COPY, statement.parsedStatement.getSqlWithoutComments()));
+                        && isCommand(COPY, statement.parsedStatement.getSqlWithoutComments()));
   }
 
   private int getStatementCount() {
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/statements/IntermediatePreparedStatement.java b/src/main/java/com/google/cloud/spanner/pgadapter/statements/IntermediatePreparedStatement.java
index f4ac9ac5..60165e46 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/statements/IntermediatePreparedStatement.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/statements/IntermediatePreparedStatement.java
@@ -401,7 +401,7 @@ public class IntermediatePreparedStatement extends IntermediateStatement {
     if (!parser.eatKeyword("set")) {
       return null;
     }
-    List<String> assignmentsList = parser.parseExpressionListUntilKeyword("where");
+    List<String> assignmentsList = parser.parseExpressionListUntilKeyword("where", true);
     if (assignmentsList == null || assignmentsList.isEmpty()) {
       return null;
     }
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/statements/SimpleParser.java b/src/main/java/com/google/cloud/spanner/pgadapter/statements/SimpleParser.java
index eaf03a8f..0bab4e1d 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/statements/SimpleParser.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/statements/SimpleParser.java
@@ -14,14 +14,27 @@
 
 package com.google.cloud.spanner.pgadapter.statements;
 
+import com.google.api.client.util.Strings;
+import com.google.api.core.InternalApi;
+import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Objects;
+import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 
 /** A very simple parser that can interpret SQL statements to find specific parts in the string. */
-class SimpleParser {
+@InternalApi
+public class SimpleParser {
+  private static final char STATEMENT_DELIMITER = ';';
+  private static final char SINGLE_QUOTE = '\'';
+  private static final char DOUBLE_QUOTE = '"';
+  private static final char HYPHEN = '-';
+  private static final char SLASH = '/';
+  private static final char ASTERISK = '*';
+  private static final char DOLLAR = '$';
+
   /** Name of table or index. */
   static class TableOrIndexName {
     /** Schema is an optional schema name prefix. */
@@ -81,18 +94,107 @@ class SimpleParser {
     this.pos = pos;
   }
 
+  /** Returns the command tag of the given SQL string */
+  public static String parseCommand(String sql) {
+    SimpleParser parser = new SimpleParser(sql);
+    if (parser.eatKeyword("with")) {
+      do {
+        if (!parser.skipCommonTableExpression()) {
+          // Return WITH as the command tag if we encounter an invalid CTE. This is for safety, as
+          // the chances that we encounter an invalid CTE is a lot larger than encountering any
+          // other statement without a logical first keyword.
+          return "WITH";
+        }
+      } while (parser.eatToken(","));
+    }
+    String keyword = parser.readKeyword();
+    if (Strings.isNullOrEmpty(keyword)) {
+      keyword = new SimpleParser(sql).readKeyword();
+    }
+    return keyword.toUpperCase();
+  }
+
+  /** Returns true if the given sql string is the given command. */
+  public static boolean isCommand(String command, String query) {
+    Preconditions.checkNotNull(command);
+    Preconditions.checkNotNull(query);
+    return new SimpleParser(query).peekKeyword(command);
+  }
+
+  /**
+   * Splits the given sql string into multiple sql statements. A semi-colon (;) indicates the end of
+   * a statement.
+   */
+  ImmutableList<String> splitStatements() {
+    // First check trivial cases with only one statement.
+    int firstIndexOfDelimiter = sql.indexOf(STATEMENT_DELIMITER);
+    if (firstIndexOfDelimiter == -1) {
+      return ImmutableList.of(sql);
+    }
+    if (firstIndexOfDelimiter == sql.length() - 1) {
+      return ImmutableList.of(sql.substring(0, sql.length() - 1));
+    }
+
+    ImmutableList.Builder<String> builder = ImmutableList.builder();
+    int lastFoundSeparatorPos = 0;
+    while (skipCommentsAndLiterals() && pos < sql.length()) {
+      if (sql.charAt(pos) == STATEMENT_DELIMITER) {
+        String stmt = sql.substring(lastFoundSeparatorPos, pos).trim();
+        builder.add(stmt);
+        lastFoundSeparatorPos = pos + 1;
+      }
+      pos++;
+    }
+
+    if (lastFoundSeparatorPos < sql.length()) {
+      String trimmed = sql.substring(lastFoundSeparatorPos).trim();
+      if (trimmed.length() > 0) {
+        builder.add(trimmed);
+      }
+    }
+    return builder.build();
+  }
+
+  boolean skipCommonTableExpression() {
+    if (readIdentifierPart() == null) {
+      return false;
+    }
+    if (eatToken("(")) {
+      List<String> columnNames = parseExpressionList();
+      if (columnNames == null || columnNames.isEmpty()) {
+        return false;
+      }
+      if (!eatToken(")")) {
+        return false;
+      }
+    }
+    if (!eatKeyword("as")) {
+      return false;
+    }
+    if (!eatToken("(")) {
+      return false;
+    }
+    parseExpressionUntilKeyword(ImmutableList.of());
+    if (!eatToken(")")) {
+      return false;
+    }
+    return true;
+  }
+
   List<String> parseExpressionList() {
-    return parseExpressionListUntilKeyword(null);
+    return parseExpressionListUntilKeyword(null, false);
   }
 
-  List<String> parseExpressionListUntilKeyword(@Nullable String keyword) {
+  List<String> parseExpressionListUntilKeyword(
+      @Nullable String keyword, boolean sameParensLevelAsStart) {
     skipWhitespaces();
     List<String> result = new ArrayList<>();
     int start = pos;
     while (pos < sql.length()) {
       String expression =
           parseExpressionUntilKeyword(
-              keyword == null ? ImmutableList.of() : ImmutableList.of(keyword));
+              keyword == null ? ImmutableList.of() : ImmutableList.of(keyword),
+              sameParensLevelAsStart);
       if (expression == null) {
         return null;
       }
@@ -112,46 +214,48 @@ class SimpleParser {
   }
 
   String parseExpressionUntilKeyword(ImmutableList<String> keywords) {
+    return parseExpressionUntilKeyword(keywords, false);
+  }
+
+  String parseExpressionUntilKeyword(
+      ImmutableList<String> keywords, boolean sameParensLevelAsStart) {
     skipWhitespaces();
     int start = pos;
-    boolean quoted = false;
-    char startQuote = 0;
+    boolean valid;
     int parens = 0;
-    while (pos < sql.length()) {
-      if (quoted) {
-        if (sql.charAt(pos) == startQuote) {
-          if (pos == (sql.length() - 1) || sql.charAt(pos + 1) != startQuote) {
-            quoted = false;
-          } else {
-            pos++;
-          }
-        }
-      } else {
-        if (sql.charAt(pos) == '\'' || sql.charAt(pos) == '"') {
-          quoted = true;
-          startQuote = sql.charAt(pos);
-        } else if (sql.charAt(pos) == '(') {
-          parens++;
-        } else if (sql.charAt(pos) == ')') {
-          parens--;
-          if (parens < 0) {
-            break;
-          }
-        } else if (parens == 0 && sql.charAt(pos) == ',') {
-          break;
-        }
-        if (keywords.stream().anyMatch(this::peekKeyword)) {
+    while ((valid = skipCommentsAndLiterals()) && pos < sql.length()) {
+      if (sql.charAt(pos) == '(') {
+        parens++;
+      } else if (sql.charAt(pos) == ')') {
+        parens--;
+        if (parens < 0) {
           break;
         }
+      } else if (parens == 0 && sql.charAt(pos) == ',') {
+        break;
+      }
+      if ((!sameParensLevelAsStart || parens == 0)
+          && keywords.stream().anyMatch(this::peekKeyword)) {
+        break;
       }
       pos++;
     }
-    if (pos == start || quoted || parens > 0) {
+    if (pos == start || !valid || parens > 0) {
       return null;
     }
     return sql.substring(start, pos).trim();
   }
 
+  @Nonnull
+  String readKeyword() {
+    skipWhitespaces();
+    int startPos = pos;
+    while (pos < sql.length() && !isValidEndOfKeyword(pos)) {
+      pos++;
+    }
+    return sql.substring(startPos, pos);
+  }
+
   TableOrIndexName readTableOrIndexName() {
     String nameOrSchema = readIdentifierPart();
     if (nameOrSchema == null) {
@@ -293,9 +397,132 @@ class SimpleParser {
     return !isValidIdentifierChar(sql.charAt(index));
   }
 
+  boolean skipCommentsAndLiterals() {
+    if (pos >= sql.length()) {
+      return true;
+    }
+    if (sql.charAt(pos) == SINGLE_QUOTE || sql.charAt(pos) == DOUBLE_QUOTE) {
+      return skipQuotedString();
+    } else if (sql.charAt(pos) == HYPHEN
+        && sql.length() > (pos + 1)
+        && sql.charAt(pos + 1) == HYPHEN) {
+      return skipSingleLineComment();
+    } else if (sql.charAt(pos) == SLASH
+        && sql.length() > (pos + 1)
+        && sql.charAt(pos + 1) == ASTERISK) {
+      return skipMultiLineComment();
+    } else if (sql.charAt(pos) == DOLLAR
+        && sql.length() > (pos + 1)
+        && (sql.charAt(pos + 1) == DOLLAR || isValidIdentifierFirstChar(sql.charAt(pos + 1)))
+        && sql.indexOf(DOLLAR, pos + 1) > -1) {
+      return skipDollarQuotedString();
+    } else {
+      return true;
+    }
+  }
+
+  boolean skipQuotedString() {
+    char quote = sql.charAt(pos);
+    pos++;
+    while (pos < sql.length()) {
+      if (sql.charAt(pos) == quote) {
+        if (sql.length() > (pos + 1) && sql.charAt(pos + 1) == quote) {
+          // This is an escaped quote. Skip one ahead.
+          pos++;
+        } else {
+          pos++;
+          return true;
+        }
+      }
+      pos++;
+    }
+    pos = sql.length();
+    return false;
+  }
+
   void skipWhitespaces() {
-    while (sql.length() > pos && Character.isWhitespace(sql.charAt(pos))) {
+    while (pos < sql.length()) {
+      if (sql.charAt(pos) == HYPHEN && sql.length() > (pos + 1) && sql.charAt(pos + 1) == HYPHEN) {
+        skipSingleLineComment();
+      } else if (sql.charAt(pos) == SLASH
+          && sql.length() > (pos + 1)
+          && sql.charAt(pos + 1) == ASTERISK) {
+        skipMultiLineComment();
+      } else if (Character.isWhitespace(sql.charAt(pos))) {
+        pos++;
+      } else {
+        break;
+      }
+    }
+  }
+
+  boolean skipSingleLineComment() {
+    int endIndex = sql.indexOf('\n', pos + 2);
+    if (endIndex == -1) {
+      pos = sql.length();
+      return true;
+    }
+    pos = endIndex + 1;
+    return true;
+  }
+
+  boolean skipMultiLineComment() {
+    int level = 1;
+    pos += 2;
+    while (pos < sql.length()) {
+      if (sql.charAt(pos) == SLASH && sql.length() > (pos + 1) && sql.charAt(pos + 1) == ASTERISK) {
+        level++;
+      }
+      if (sql.charAt(pos) == ASTERISK && sql.length() > (pos + 1) && sql.charAt(pos + 1) == SLASH) {
+        level--;
+        if (level == 0) {
+          pos += 2;
+          return true;
+        }
+      }
+      pos++;
+    }
+    pos = sql.length();
+    return false;
+  }
+
+  String parseDollarQuotedTag() {
+    // Look ahead to the next dollar sign (if any). Everything in between is the quote tag.
+    StringBuilder tag = new StringBuilder();
+    while (pos < sql.length()) {
+      char c = sql.charAt(pos);
+      if (c == DOLLAR) {
+        pos++;
+        return tag.toString();
+      }
+      if (!isValidIdentifierChar(c)) {
+        break;
+      }
+      tag.append(c);
       pos++;
     }
+    return null;
+  }
+
+  boolean skipDollarQuotedString() {
+    if (sql.charAt(pos) != DOLLAR) {
+      return false;
+    }
+    pos++;
+    String tag = parseDollarQuotedTag();
+    if (tag == null) {
+      return false;
+    }
+    while (pos < sql.length()) {
+      if (sql.charAt(pos++) == DOLLAR) {
+        int currentPos = pos;
+        String endTag = parseDollarQuotedTag();
+        if (Objects.equals(tag, endTag)) {
+          return true;
+        }
+        pos = currentPos;
+      }
+    }
+    return false;
   }
 }
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/statements/SimpleQueryStatement.java b/src/main/java/com/google/cloud/spanner/pgadapter/statements/SimpleQueryStatement.java
index 9c1a8e38..ea92d942 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/statements/SimpleQueryStatement.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/statements/SimpleQueryStatement.java
@@ -14,7 +14,7 @@
 
 package com.google.cloud.spanner.pgadapter.statements;
 
-import static com.google.cloud.spanner.pgadapter.utils.StatementParser.splitStatements;
+import static com.google.cloud.spanner.pgadapter.statements.SimpleParser.isCommand;
 import static com.google.cloud.spanner.pgadapter.wireprotocol.QueryMessage.COPY;
 
 import com.google.api.core.InternalApi;
@@ -28,7 +28,6 @@ import com.google.cloud.spanner.pgadapter.commands.Command;
 import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata;
 import com.google.cloud.spanner.pgadapter.statements.BackendConnection.ConnectionState;
 import com.google.cloud.spanner.pgadapter.utils.ClientAutoDetector.WellKnownClient;
-import com.google.cloud.spanner.pgadapter.utils.StatementParser;
 import com.google.cloud.spanner.pgadapter.wireprotocol.BindMessage;
 import com.google.cloud.spanner.pgadapter.wireprotocol.ControlMessage.ManuallyCreatedToken;
 import com.google.cloud.spanner.pgadapter.wireprotocol.DescribeMessage;
@@ -85,7 +84,7 @@ public class SimpleQueryStatement {
         // We need to flush the entire pipeline if we encounter a COPY statement, as COPY statements
         // require additional messages to be sent back and forth, and this ensures that we get
         // everything in the correct order.
-        boolean isCopy = StatementParser.isCommand(COPY, parsedStatement.getSqlWithoutComments());
+        boolean isCopy = isCommand(COPY, parsedStatement.getSqlWithoutComments());
         if (!isFirst && isCopy) {
           new FlushMessage(connectionHandler, ManuallyCreatedToken.MANUALLY_CREATED_TOKEN).send();
           if (connectionHandler
@@ -148,7 +147,8 @@ public class SimpleQueryStatement {
   protected static ImmutableList<Statement> parseStatements(Statement statement) {
     Preconditions.checkNotNull(statement);
     ImmutableList.Builder<Statement> builder = ImmutableList.builder();
-    for (String sql : splitStatements(statement.getSql())) {
+    SimpleParser parser = new SimpleParser(statement.getSql());
+    for (String sql : parser.splitStatements()) {
       builder.add(Statement.of(sql));
     }
     return builder.build();
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/utils/StatementParser.java b/src/main/java/com/google/cloud/spanner/pgadapter/utils/StatementParser.java
deleted file mode 100644
index 9e1232fe..00000000
--- a/src/main/java/com/google/cloud/spanner/pgadapter/utils/StatementParser.java
+++ /dev/null
@@ -1,194 +0,0 @@
-// Copyright 2020 Google LLC
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//      http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-package com.google.cloud.spanner.pgadapter.utils;
-
-import com.google.api.core.InternalApi;
-import com.google.common.base.Preconditions;
-import com.google.common.collect.ImmutableList;
-import java.util.Objects;
-
-// TODO: Replace everything here with parsing in the client library.
-@InternalApi
-public class StatementParser {
-  private static final char STATEMENT_DELIMITER = ';';
-  private static final char SINGLE_QUOTE = '\'';
-  private static final char DOUBLE_QUOTE = '"';
-  private static final char HYPHEN = '-';
-  private static final char SLASH = '/';
-  private static final char ASTERISK = '*';
-  private static final char DOLLAR = '$';
-
-  /**
-   * Simple method to escape SQL. Ultimately it is preferred that a user uses PreparedStatements but
-   * for the case of psql emulation, we apply this to provide a simple layer of protection to the
-   * user. This method simply duplicates all single quotes (i.e. ' becomes '').
-   */
-  public static String singleQuoteEscape(String sql) {
-    return sql.replace("'", "''");
-  }
-
-  /** Determines the (update) command that was received from the sql string. */
-  public static String parseCommand(String sql) {
-    Preconditions.checkNotNull(sql);
-    for (int i = 0; i < sql.length(); i++) {
-      if (Character.isSpaceChar(sql.charAt(i))) {
-        return sql.substring(0, i).toUpperCase();
-      }
-    }
-    return sql.toUpperCase();
-  }
-
-  /** Returns true if the given sql string is the given command. */
-  public static boolean isCommand(String command, String query) {
-    Preconditions.checkNotNull(command);
-    Preconditions.checkNotNull(query);
-    if (query.equalsIgnoreCase(command)) {
-      return true;
-    }
-    if (query.length() <= command.length()) {
-      return false;
-    }
-    return Character.isSpaceChar(query.charAt(command.length()))
-        && query.substring(0, command.length()).equalsIgnoreCase(command);
-  }
-
-  /**
-   * Splits the given sql string into multiple sql statements. A semi-colon (;) indicates the end of
-   * a statement.
-   */
-  public static ImmutableList<String> splitStatements(String sql) {
-    // First check trivial cases with only one statement.
-    int firstIndexOfDelimiter = sql.indexOf(STATEMENT_DELIMITER);
-    if (firstIndexOfDelimiter == -1) {
-      return ImmutableList.of(sql);
-    }
-    if (firstIndexOfDelimiter == sql.length() - 1) {
-      return ImmutableList.of(sql.substring(0, sql.length() - 1));
-    }
-
-    ImmutableList.Builder<String> builder = ImmutableList.builder();
-    int index = 0;
-    int i = 0;
-    while (i < sql.length()) {
-      if (sql.charAt(i) == SINGLE_QUOTE || sql.charAt(i) == DOUBLE_QUOTE) {
-        i = skipQuotedString(sql, i);
-      } else if (sql.charAt(i) == HYPHEN && sql.length() > (i + 1) && sql.charAt(i + 1) == HYPHEN) {
-        i = skipSingleLineComment(sql, i);
-      } else if (sql.charAt(i) == SLASH
-          && sql.length() > (i + 1)
-          && sql.charAt(i + 1) == ASTERISK) {
-        i = skipMultiLineComment(sql, i);
-      } else if (sql.charAt(i) == DOLLAR
-          && sql.length() > (i + 1)
-          && (sql.charAt(i + 1) == DOLLAR || Character.isJavaIdentifierPart(sql.charAt(i + 1)))
-          && sql.indexOf(DOLLAR, i + 1) > -1) {
-        i = skipDollarQuotedString(sql, i);
-      } else {
-        if (sql.charAt(i) == STATEMENT_DELIMITER) {
-          String stmt = sql.substring(index, i).trim();
-          builder.add(stmt);
-          index = i + 1;
-        }
-        i++;
-      }
-    }
-
-    if (index < sql.length()) {
-      String trimmed = sql.substring(index).trim();
-      if (trimmed.length() > 0) {
-        builder.add(trimmed);
-      }
-    }
-    return builder.build();
-  }
-
-  static int skipQuotedString(String sql, int startIndex) {
-    char quote = sql.charAt(startIndex);
-    int i = startIndex + 1;
-    while (i < sql.length()) {
-      if (sql.charAt(i) == quote) {
-        if (sql.length() > (i + 1) && sql.charAt(i + 1) == quote) {
-          // This is an escaped quote. Skip one ahead.
-          i++;
-        } else {
-          return i + 1;
-        }
-      }
-      i++;
-    }
-    return sql.length();
-  }
-
-  static int skipSingleLineComment(String sql, int startIndex) {
-    int endIndex = sql.indexOf('\n', startIndex + 2);
-    if (endIndex == -1) {
-      return sql.length();
-    }
-    return endIndex + 1;
-  }
-
-  static int skipMultiLineComment(String sql, int startIndex) {
-    int level = 1;
-    int i = startIndex + 2;
-    while (i < sql.length()) {
-      if (sql.charAt(i) == SLASH && sql.length() > (i + 1) && sql.charAt(i + 1) == ASTERISK) {
-        level++;
-      }
-      if (sql.charAt(i) == ASTERISK && sql.length() > (i + 1) && sql.charAt(i + 1) == SLASH) {
-        level--;
-        if (level == 0) {
-          return i + 2;
-        }
-      }
-      i++;
-    }
-    return sql.length();
-  }
-
-  static String parseDollarQuotedTag(String sql, int startIndex) {
-    // Look ahead to the next dollar sign (if any). Everything in between is the quote tag.
-    StringBuilder tag = new StringBuilder();
-    while (startIndex < sql.length()) {
-      char c = sql.charAt(startIndex);
-      if (c == DOLLAR) {
-        return tag.toString();
-      }
-      if (!Character.isJavaIdentifierPart(c)) {
-        break;
-      }
-      tag.append(c);
-      startIndex++;
-    }
-    return null;
-  }
-
-  static int skipDollarQuotedString(String sql, int startIndex) {
-    String tag = parseDollarQuotedTag(sql, startIndex + 1);
-    if (tag == null) {
-      return startIndex + 1;
-    }
-    int i = startIndex + 1;
-    while (i < sql.length()) {
-      if (sql.charAt(i) == DOLLAR) {
-        String endTag = parseDollarQuotedTag(sql, i + 1);
-        if (Objects.equals(tag, endTag)) {
-          return i + endTag.length() + 2;
-        }
-      }
-      i++;
-    }
-    return sql.length();
-  }
-}
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/ParseMessage.java b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/ParseMessage.java
index c652ff40..bae025f9 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/ParseMessage.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/ParseMessage.java
@@ -15,6 +15,7 @@
 package com.google.cloud.spanner.pgadapter.wireprotocol;
 
 import static com.google.cloud.spanner.pgadapter.parsers.copy.Copy.parse;
+import static com.google.cloud.spanner.pgadapter.statements.SimpleParser.isCommand;
 import static com.google.cloud.spanner.pgadapter.wireprotocol.QueryMessage.COPY;
 
 import com.google.api.core.InternalApi;
@@ -33,7 +34,6 @@ import com.google.cloud.spanner.pgadapter.statements.BackendConnection;
 import com.google.cloud.spanner.pgadapter.statements.CopyStatement;
 import com.google.cloud.spanner.pgadapter.statements.CopyToStatement;
 import com.google.cloud.spanner.pgadapter.statements.IntermediatePreparedStatement;
-import com.google.cloud.spanner.pgadapter.utils.StatementParser;
 import com.google.cloud.spanner.pgadapter.wireoutput.ParseCompleteResponse;
 import com.google.common.base.Strings;
 import java.text.MessageFormat;
@@ -84,7 +84,7 @@ public class ParseMessage extends AbstractQueryProtocolMessage {
       ParsedStatement parsedStatement,
       Statement originalStatement,
       int[] parameterDataTypes) {
-    if (StatementParser.isCommand(COPY, parsedStatement.getSqlWithoutComments())) {
+    if (isCommand(COPY, originalStatement.getSql())) {
       CopyOptions copyOptions = parseCopyStatement(parsedStatement);
       if (copyOptions.getFromTo() == FromTo.FROM) {
         return new CopyStatement(


Test patch:
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/statements/IntermediateStatement.java b/src/main/java/com/google/cloud/spanner/pgadapter/statements/IntermediateStatement.java
index 9320403f..a7e7a603 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/statements/IntermediateStatement.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/statements/IntermediateStatement.java
@@ -14,6 +14,8 @@
 
 package com.google.cloud.spanner.pgadapter.statements;
 
+import static com.google.cloud.spanner.pgadapter.statements.SimpleParser.parseCommand;
+
 import com.google.api.core.InternalApi;
 import com.google.cloud.spanner.Dialect;
 import com.google.cloud.spanner.ResultSet;
@@ -32,7 +34,6 @@ import com.google.cloud.spanner.pgadapter.ConnectionHandler.QueryMode;
 import com.google.cloud.spanner.pgadapter.metadata.DescribeMetadata;
 import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata;
 import com.google.cloud.spanner.pgadapter.statements.BackendConnection.NoResult;
-import com.google.cloud.spanner.pgadapter.utils.StatementParser;
 import com.google.cloud.spanner.pgadapter.wireoutput.DataRowResponse;
 import com.google.cloud.spanner.pgadapter.wireoutput.WireOutput;
 import java.io.DataOutputStream;
@@ -105,7 +106,7 @@ public class IntermediateStatement {
     }
     this.parsedStatement = potentiallyReplacedStatement;
     this.connection = connectionHandler.getSpannerConnection();
-    this.command = StatementParser.parseCommand(this.parsedStatement.getSqlWithoutComments());
+    this.command = parseCommand(this.parsedStatement.getSqlWithoutComments());
     this.commandTag = this.command;
     this.outputStream = connectionHandler.getConnectionMetadata().peekOutputStream();
   }
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/JdbcMockServerTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/JdbcMockServerTest.java
index c453ff1c..c4b58e1c 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/JdbcMockServerTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/JdbcMockServerTest.java
@@ -71,6 +71,7 @@ import org.junit.runners.JUnit4;
 import org.postgresql.PGConnection;
 import org.postgresql.PGStatement;
 import org.postgresql.jdbc.PgStatement;
+import org.postgresql.util.PSQLException;
 
 @RunWith(JUnit4.class)
 public class JdbcMockServerTest extends AbstractMockServerTest {
@@ -126,6 +127,21 @@ public class JdbcMockServerTest extends AbstractMockServerTest {
     }
   }
 
+  @Test
+  public void testInvalidQuery() throws SQLException {
+    String sql = "/ not a valid comment / SELECT 1";
+
+    try (Connection connection = DriverManager.getConnection(createUrl())) {
+      PSQLException exception =
+          assertThrows(PSQLException.class, () -> connection.createStatement().executeQuery(sql));
+      assertEquals(
+          "ERROR: Unknown statement: / not a valid comment / SELECT 1", exception.getMessage());
+    }
+
+    // The statement is not sent to the mock server.
+    assertEquals(0, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
+  }
+
   @Test
   public void testSelectCurrentSchema() throws SQLException {
     String sql = "SELECT current_schema";
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/parsers/ParserTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/parsers/ParserTest.java
index 9b197aeb..454814a5 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/parsers/ParserTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/parsers/ParserTest.java
@@ -15,6 +15,7 @@
 package com.google.cloud.spanner.pgadapter.parsers;
 
 import static com.google.cloud.spanner.pgadapter.parsers.copy.Copy.parse;
+import static com.google.cloud.spanner.pgadapter.statements.SimpleParser.parseCommand;
 import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
@@ -40,7 +41,6 @@ import com.google.cloud.spanner.pgadapter.parsers.copy.CopyTreeParser.CopyOption
 import com.google.cloud.spanner.pgadapter.parsers.copy.CopyTreeParser.CopyOptions.FromTo;
 import com.google.cloud.spanner.pgadapter.parsers.copy.ParseException;
 import com.google.cloud.spanner.pgadapter.parsers.copy.TokenMgrError;
-import com.google.cloud.spanner.pgadapter.utils.StatementParser;
 import com.google.common.collect.ImmutableList;
 import java.math.BigDecimal;
 import java.nio.charset.StandardCharsets;
@@ -427,7 +427,7 @@ public class ParserTest {
       String sql = "COPY users FROM STDIN;";
       parse(sql, options);
 
-      assertEquals(StatementParser.parseCommand(PARSER.removeCommentsAndTrim(sql)), "COPY");
+      assertEquals(parseCommand(PARSER.removeCommentsAndTrim(sql)), "COPY");
       assertEquals("users", options.getTableName());
       assertEquals(FromTo.FROM, options.getFromTo());
       assertEquals(Format.TEXT, options.getFormat());
@@ -438,7 +438,7 @@ public class ParserTest {
       String sql = "COPY users TO STDOUT;";
       parse(sql, options);
 
-      assertEquals(StatementParser.parseCommand(PARSER.removeCommentsAndTrim(sql)), "COPY");
+      assertEquals(parseCommand(PARSER.removeCommentsAndTrim(sql)), "COPY");
       assertEquals("users", options.getTableName());
       assertEquals(FromTo.TO, options.getFromTo());
       assertEquals("STDOUT", options.getFilename());
@@ -450,7 +450,7 @@ public class ParserTest {
       String sql = "COPY Users FROM STDIN;";
       parse(sql, options);
 
-      assertEquals(StatementParser.parseCommand(PARSER.removeCommentsAndTrim(sql)), "COPY");
+      assertEquals(parseCommand(PARSER.removeCommentsAndTrim(sql)), "COPY");
       assertEquals("users", options.getTableName());
       assertEquals(FromTo.FROM, options.getFromTo());
       assertEquals(Format.TEXT, options.getFormat());
@@ -462,7 +462,7 @@ public class ParserTest {
       String sql = "COPY UsErS FROM STDIN;";
       parse(sql, options);
 
-      assertEquals(StatementParser.parseCommand(PARSER.removeCommentsAndTrim(sql)), "COPY");
+      assertEquals(parseCommand(PARSER.removeCommentsAndTrim(sql)), "COPY");
       assertEquals("users", options.getTableName());
       assertEquals(FromTo.FROM, options.getFromTo());
       assertEquals(Format.TEXT, options.getFormat());
@@ -474,7 +474,7 @@ public class ParserTest {
       String sql = "COPY USERS FROM STDIN;";
       parse(sql, options);
 
-      assertEquals(StatementParser.parseCommand(PARSER.removeCommentsAndTrim(sql)), "COPY");
+      assertEquals(parseCommand(PARSER.removeCommentsAndTrim(sql)), "COPY");
       assertEquals("users", options.getTableName());
       assertEquals(FromTo.FROM, options.getFromTo());
       assertEquals(Format.TEXT, options.getFormat());
@@ -486,7 +486,7 @@ public class ParserTest {
       String sql = "COPY 'Users' FROM STDIN;";
       parse(sql, options);
 
-      assertEquals(StatementParser.parseCommand(PARSER.removeCommentsAndTrim(sql)), "COPY");
+      assertEquals(parseCommand(PARSER.removeCommentsAndTrim(sql)), "COPY");
       assertEquals("Users", options.getTableName());
       assertEquals(FromTo.FROM, options.getFromTo());
       assertEquals(Format.TEXT, options.getFormat());
@@ -498,7 +498,7 @@ public class ParserTest {
       String sql = "COPY \"UsErS\" FROM STDIN;";
       parse(sql, options);
 
-      assertEquals(StatementParser.parseCommand(PARSER.removeCommentsAndTrim(sql)), "COPY");
+      assertEquals(parseCommand(PARSER.removeCommentsAndTrim(sql)), "COPY");
       assertEquals("UsErS", options.getTableName());
       assertEquals(FromTo.FROM, options.getFromTo());
       assertEquals(Format.TEXT, options.getFormat());
@@ -510,7 +510,7 @@ public class ParserTest {
       String sql = "COPY 'USERS' FROM STDIN;";
       parse(sql, options);
 
-      assertEquals(StatementParser.parseCommand(PARSER.removeCommentsAndTrim(sql)), "COPY");
+      assertEquals(parseCommand(PARSER.removeCommentsAndTrim(sql)), "COPY");
       assertEquals("USERS", options.getTableName());
       assertEquals(FromTo.FROM, options.getFromTo());
       assertEquals(Format.TEXT, options.getFormat());
@@ -522,7 +522,7 @@ public class ParserTest {
       String sql = "COPY users (id) FROM STDIN;";
       parse(sql, options);
 
-      assertEquals(StatementParser.parseCommand(PARSER.removeCommentsAndTrim(sql)), "COPY");
+      assertEquals(parseCommand(PARSER.removeCommentsAndTrim(sql)), "COPY");
       assertEquals("users", options.getTableName());
       assertEquals(FromTo.FROM, options.getFromTo());
       assertEquals(Format.TEXT, options.getFormat());
@@ -534,7 +534,7 @@ public class ParserTest {
       String sql = "COPY users (id, age) FROM STDIN;";
       parse(sql, options);
 
-      assertEquals(StatementParser.parseCommand(PARSER.removeCommentsAndTrim(sql)), "COPY");
+      assertEquals(parseCommand(PARSER.removeCommentsAndTrim(sql)), "COPY");
       assertEquals("users", options.getTableName());
       assertEquals(FromTo.FROM, options.getFromTo());
       assertEquals(Format.TEXT, options.getFormat());
@@ -546,7 +546,7 @@ public class ParserTest {
       String sql = "COPY users (id, age, name) FROM STDIN;";
       parse(sql, options);
 
-      assertEquals(StatementParser.parseCommand(PARSER.removeCommentsAndTrim(sql)), "COPY");
+      assertEquals(parseCommand(PARSER.removeCommentsAndTrim(sql)), "COPY");
       assertEquals("users", options.getTableName());
       assertEquals(FromTo.FROM, options.getFromTo());
       assertEquals(Format.TEXT, options.getFormat());
@@ -558,7 +558,7 @@ public class ParserTest {
       String sql = "COPY users (ID, AGE, NAME) FROM STDIN;";
       parse(sql, options);
 
-      assertEquals(StatementParser.parseCommand(PARSER.removeCommentsAndTrim(sql)), "COPY");
+      assertEquals(parseCommand(PARSER.removeCommentsAndTrim(sql)), "COPY");
       assertEquals("users", options.getTableName());
       assertEquals(FromTo.FROM, options.getFromTo());
       assertEquals(Format.TEXT, options.getFormat());
@@ -570,7 +570,7 @@ public class ParserTest {
       String sql = "COPY users (Id, Age, Name) FROM STDIN;";
       parse(sql, options);
 
-      assertEquals(StatementParser.parseCommand(PARSER.removeCommentsAndTrim(sql)), "COPY");
+      assertEquals(parseCommand(PARSER.removeCommentsAndTrim(sql)), "COPY");
       assertEquals("users", options.getTableName());
       assertEquals(FromTo.FROM, options.getFromTo());
       assertEquals(Format.TEXT, options.getFormat());
@@ -582,7 +582,7 @@ public class ParserTest {
       String sql = "COPY users (iD, aGe, nAMe) FROM STDIN;";
       parse(sql, options);
 
-      assertEquals(StatementParser.parseCommand(PARSER.removeCommentsAndTrim(sql)), "COPY");
+      assertEquals(parseCommand(PARSER.removeCommentsAndTrim(sql)), "COPY");
       assertEquals("users", options.getTableName());
       assertEquals(FromTo.FROM, options.getFromTo());
       assertEquals(Format.TEXT, options.getFormat());
@@ -594,7 +594,7 @@ public class ParserTest {
       String sql = "COPY users ('ID', \"AGE\", 'NAME') FROM STDIN;";
       parse(sql, options);
 
-      assertEquals(StatementParser.parseCommand(PARSER.removeCommentsAndTrim(sql)), "COPY");
+      assertEquals(parseCommand(PARSER.removeCommentsAndTrim(sql)), "COPY");
       assertEquals("users", options.getTableName());
       assertEquals(FromTo.FROM, options.getFromTo());
       assertEquals(Format.TEXT, options.getFormat());
@@ -606,7 +606,7 @@ public class ParserTest {
       String sql = "COPY users ('Id', 'Age', 'Name') FROM STDIN;";
       parse(sql, options);
 
-      assertEquals(StatementParser.parseCommand(PARSER.removeCommentsAndTrim(sql)), "COPY");
+      assertEquals(parseCommand(PARSER.removeCommentsAndTrim(sql)), "COPY");
       assertEquals("users", options.getTableName());
       assertEquals(FromTo.FROM, options.getFromTo());
       assertEquals(Format.TEXT, options.getFormat());
@@ -618,7 +618,7 @@ public class ParserTest {
       String sql = "COPY users (\"iD\", \"aGe\", \"nAMe\") FROM STDIN;";
       parse(sql, options);
 
-      assertEquals(StatementParser.parseCommand(PARSER.removeCommentsAndTrim(sql)), "COPY");
+      assertEquals(parseCommand(PARSER.removeCommentsAndTrim(sql)), "COPY");
       assertEquals("users", options.getTableName());
       assertEquals(FromTo.FROM, options.getFromTo());
       assertEquals(Format.TEXT, options.getFormat());
@@ -633,7 +633,7 @@ public class ParserTest {
       String sql = "COPY users FROM STDIN WITH FORMAT CSV;";
       parse(sql, options);
 
-      assertEquals(StatementParser.parseCommand(PARSER.removeCommentsAndTrim(sql)), "COPY");
+      assertEquals(parseCommand(PARSER.removeCommentsAndTrim(sql)), "COPY");
       assertEquals("users", options.getTableName());
       assertEquals(FromTo.FROM, options.getFromTo());
       assertEquals(Format.CSV, options.getFormat());
@@ -644,7 +644,7 @@ public class ParserTest {
       String sql = "COPY users FROM STDIN WITH FORMAT BINARY;";
       parse(sql, options);
 
-      assertEquals(StatementParser.parseCommand(PARSER.removeCommentsAndTrim(sql)), "COPY");
+      assertEquals(parseCommand(PARSER.removeCommentsAndTrim(sql)), "COPY");
       assertEquals("users", options.getTableName());
       assertEquals(FromTo.FROM, options.getFromTo());
       assertEquals(Format.BINARY, options.getFormat());
@@ -655,7 +655,7 @@ public class ParserTest {
       String sql = "COPY users FROM STDIN WITH FORMAT TEXT;";
       parse(sql, options);
 
-      assertEquals(StatementParser.parseCommand(PARSER.removeCommentsAndTrim(sql)), "COPY");
+      assertEquals(parseCommand(PARSER.removeCommentsAndTrim(sql)), "COPY");
       assertEquals("users", options.getTableName());
       assertEquals(FromTo.FROM, options.getFromTo());
       assertEquals(Format.TEXT, options.getFormat());
@@ -666,7 +666,7 @@ public class ParserTest {
       String sql = "COPY users FROM STDIN CSV;";
       parse(sql, options);
 
-      assertEquals(StatementParser.parseCommand(PARSER.removeCommentsAndTrim(sql)), "COPY");
+      assertEquals(parseCommand(PARSER.removeCommentsAndTrim(sql)), "COPY");
       assertEquals("users", options.getTableName());
       assertEquals(FromTo.FROM, options.getFromTo());
       assertEquals(Format.CSV, options.getFormat());
@@ -677,7 +677,7 @@ public class ParserTest {
       String sql = "COPY users FROM STDIN BINARY;";
       parse(sql, options);
 
-      assertEquals(StatementParser.parseCommand(PARSER.removeCommentsAndTrim(sql)), "COPY");
+      assertEquals(parseCommand(PARSER.removeCommentsAndTrim(sql)), "COPY");
       assertEquals("users", options.getTableName());
       assertEquals(FromTo.FROM, options.getFromTo());
       assertEquals(Format.BINARY, options.getFormat());
@@ -688,7 +688,7 @@ public class ParserTest {
       String sql = "COPY users FROM STDIN TEXT;";
       parse(sql, options);
 
-      assertEquals(StatementParser.parseCommand(PARSER.removeCommentsAndTrim(sql)), "COPY");
+      assertEquals(parseCommand(PARSER.removeCommentsAndTrim(sql)), "COPY");
       assertEquals("users", options.getTableName());
       assertEquals(FromTo.FROM, options.getFromTo());
       assertEquals(Format.TEXT, options.getFormat());
@@ -699,7 +699,7 @@ public class ParserTest {
       String sql = "COPY users FROM STDIN DELIMITER '*';";
       parse(sql, options);
 
-      assertEquals(StatementParser.parseCommand(PARSER.removeCommentsAndTrim(sql)), "COPY");
+      assertEquals(parseCommand(PARSER.removeCommentsAndTrim(sql)), "COPY");
       assertEquals("users", options.getTableName());
       assertEquals(FromTo.FROM, options.getFromTo());
       assertEquals(Format.TEXT, options.getFormat());
@@ -711,7 +711,7 @@ public class ParserTest {
       String sql = "COPY users FROM STDIN ESCAPE '\\';";
       parse(sql, options);
 
-      assertEquals(StatementParser.parseCommand(PARSER.removeCommentsAndTrim(sql)), "COPY");
+      assertEquals(parseCommand(PARSER.removeCommentsAndTrim(sql)), "COPY");
       assertEquals("users", options.getTableName());
       assertEquals(FromTo.FROM, options.getFromTo());
       assertEquals(Format.TEXT, options.getFormat());
@@ -725,8 +725,7 @@ public class ParserTest {
 
       SpannerException exception =
           assertThrows(
-              SpannerException.class,
-              () -> StatementParser.parseCommand(PARSER.removeCommentsAndTrim(sql)));
+              SpannerException.class, () -> parseCommand(PARSER.removeCommentsAndTrim(sql)));
       assertEquals(ErrorCode.INVALID_ARGUMENT, exception.getErrorCode());
       assertEquals(
           "INVALID_ARGUMENT: SQL statement contains an unclosed literal: COPY users FROM STDIN QUOTE ''';",
@@ -742,7 +741,7 @@ public class ParserTest {
       String sql = "COPY users FROM STDIN NULL 'null';";
       parse(sql, options);
 
-      assertEquals(StatementParser.parseCommand(PARSER.removeCommentsAndTrim(sql)), "COPY");
+      assertEquals(parseCommand(PARSER.removeCommentsAndTrim(sql)), "COPY");
       assertEquals("users", options.getTableName());
       assertEquals(FromTo.FROM, options.getFromTo());
       assertEquals(Format.TEXT, options.getFormat());
@@ -758,11 +757,10 @@ public class ParserTest {
         if (value == '\'') {
           SpannerException exception =
               assertThrows(
-                  SpannerException.class,
-                  () -> StatementParser.parseCommand(PARSER.removeCommentsAndTrim(sql)));
+                  SpannerException.class, () -> parseCommand(PARSER.removeCommentsAndTrim(sql)));
           assertEquals(ErrorCode.INVALID_ARGUMENT, exception.getErrorCode());
         } else {
-          assertEquals(StatementParser.parseCommand(PARSER.removeCommentsAndTrim(sql)), "COPY");
+          assertEquals(parseCommand(PARSER.removeCommentsAndTrim(sql)), "COPY");
         }
         assertEquals("users", options.getTableName());
         assertEquals(FromTo.FROM, options.getFromTo());
@@ -789,11 +787,10 @@ public class ParserTest {
         if (value >= '%' && value <= '\'') {
           SpannerException exception =
               assertThrows(
-                  SpannerException.class,
-                  () -> StatementParser.parseCommand(PARSER.removeCommentsAndTrim(sql)));
+                  SpannerException.class, () -> parseCommand(PARSER.removeCommentsAndTrim(sql)));
           assertEquals(ErrorCode.INVALID_ARGUMENT, exception.getErrorCode());
         } else {
-          assertEquals(StatementParser.parseCommand(PARSER.removeCommentsAndTrim(sql)), "COPY");
+          assertEquals(parseCommand(PARSER.removeCommentsAndTrim(sql)), "COPY");
         }
         assertEquals("users", options.getTableName());
         assertEquals(FromTo.FROM, options.getFromTo());
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/statements/IntermediateStatementTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/statements/IntermediateStatementTest.java
index 0ba71151..5a523b41 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/statements/IntermediateStatementTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/statements/IntermediateStatementTest.java
@@ -18,7 +18,7 @@ import static com.google.cloud.spanner.pgadapter.statements.IntermediatePrepared
 import static com.google.cloud.spanner.pgadapter.statements.IntermediatePreparedStatement.transformDeleteToSelectParams;
 import static com.google.cloud.spanner.pgadapter.statements.IntermediatePreparedStatement.transformInsertToSelectParams;
 import static com.google.cloud.spanner.pgadapter.statements.IntermediatePreparedStatement.transformUpdateToSelectParams;
-import static com.google.cloud.spanner.pgadapter.utils.StatementParser.splitStatements;
+import static com.google.cloud.spanner.pgadapter.statements.SimpleParserTest.splitStatements;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
@@ -267,6 +267,11 @@ public class IntermediateStatementTest {
         transformUpdate(
                 "update foo set col1=$1, col2=$2, col3=$3, col4=$4, col5=$5, col6=$6, col7=$7, col8=$8, col9=$9  where id=$10")
             .getSql());
+    assertEquals(
+        "select $1, $2, $3 from (select col1=(select col2 from bar where col3=$1), col2=$2 from foo where id=$3) p",
+        transformUpdate(
+                "update foo set col1=(select col2 from bar where col3=$1), col2=$2 where id=$3")
+            .getSql());
   }
 
   @Test
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/statements/SimpleParserTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/statements/SimpleParserTest.java
index e5f77426..9616db5a 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/statements/SimpleParserTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/statements/SimpleParserTest.java
@@ -14,6 +14,7 @@
 
 package com.google.cloud.spanner.pgadapter.statements;
 
+import static com.google.cloud.spanner.pgadapter.statements.SimpleParser.parseCommand;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNull;
@@ -30,6 +31,39 @@ import org.junit.runners.JUnit4;
 @RunWith(JUnit4.class)
 public class SimpleParserTest {
 
+  /** Helper methods to reuse existing tests */
+  public static ImmutableList<String> splitStatements(String sql) {
+    return new SimpleParser(sql).splitStatements();
+  }
+
+  public static int skipQuotedString(String sql, int startIndex) {
+    SimpleParser parser = new SimpleParser(sql);
+    parser.setPos(startIndex);
+    parser.skipQuotedString();
+    return parser.getPos();
+  }
+
+  public static int skipSingleLineComment(String sql, int startIndex) {
+    SimpleParser parser = new SimpleParser(sql);
+    parser.setPos(startIndex);
+    parser.skipSingleLineComment();
+    return parser.getPos();
+  }
+
+  public static int skipMultiLineComment(String sql, int startIndex) {
+    SimpleParser parser = new SimpleParser(sql);
+    parser.setPos(startIndex);
+    parser.skipMultiLineComment();
+    return parser.getPos();
+  }
+
+  public static int skipDollarQuotedString(String sql, int startIndex) {
+    SimpleParser parser = new SimpleParser(sql);
+    parser.setPos(startIndex);
+    parser.skipDollarQuotedString();
+    return parser.getPos();
+  }
+
   @Test
   public void testEatKeyword() {
     assertTrue(new SimpleParser("insert").eatKeyword("insert"));
@@ -308,4 +342,71 @@ public class SimpleParserTest {
         new SimpleParser("insert into foo (\"\"\"\") select * from bar")
             .parseExpressionUntilKeyword(ImmutableList.of("select")));
   }
+
+  @Test
+  public void testParseCommand() {
+    assertEquals("SELECT", parseCommand("select * from foo"));
+    assertEquals("INSERT", parseCommand("insert into foo values (1, 'One')"));
+    assertEquals("SELECT", parseCommand("/* this is a comment */ select * from foo"));
+    assertEquals("CREATE", parseCommand("/* ddl statements are also supported */create table foo"));
+
+    assertEquals(
+        "UPDATE", parseCommand("with my_cte as (select * from foo) update bar set col1='one'"));
+    assertEquals(
+        "UPDATE",
+        parseCommand(
+            "with my_cte as (select * from foo), my_cte2 as (select * from bar) update bar set col1='one'"));
+    assertEquals(
+        "UPDATE",
+        parseCommand(
+            "with my_cte as (select * from foo), my_cte2 as (select * from bar) /* this is a comment*/update bar set col1='one'"));
+    assertEquals(
+        "UPDATE",
+        parseCommand(
+            "with my_cte as (select * from foo) -- also a comment\n, my_cte2 as (select * from bar) /* this is a comment*/update bar set col1='one'"));
+    assertEquals(
+        "UPSERT", parseCommand("with my_cte as (select * from foo) upsert bar set col1='one'"));
+    assertEquals(
+        "SELECT",
+        parseCommand(
+            "WITH temp (value) as (SELECT avg(id) FROM users) SELECT id FROM users, temp ORDER BY id"));
+    assertEquals(
+        "SELECT",
+        parseCommand(
+            "WITH temp (col1,/*comment*/col2, col3) as (SELECT avg(id) FROM users) SELECT id FROM users, temp ORDER BY id"));
+
+    // The parser returns 'WITH' as the command tag if the CTE appears to be invalid.
+    assertEquals(
+        "WITH", parseCommand("with my_cte as (select * from foo update bar set col1='one'"));
+    assertEquals(
+        "WITH", parseCommand("with my_cte as (select * from foo), update bar set col1='one'"));
+    assertEquals("WITH", parseCommand("with my_cte (select * from foo) update bar set col1='one'"));
+    assertEquals(
+        "WITH", parseCommand("with my_cte.bar as (select * from foo) update bar set col1='one'"));
+    assertEquals("WITH", parseCommand("with (select * from foo update bar set col1='one'"));
+    assertEquals("WITH", parseCommand("with 1 as (select * from foo update bar set col1='one'"));
+    assertEquals(
+        "WITH", parseCommand("with my_cte as select * from foo update bar set col1='one'"));
+    assertEquals("WITH", parseCommand("with my_cte () as (select * from foo)"));
+    assertEquals("WITH", parseCommand("with my_cte (,) as (select * from foo)"));
+    assertEquals("WITH", parseCommand("with my_cte (bar"));
+    assertEquals("WITH", parseCommand("with my_cte (bar-"));
+    assertEquals("WITH", parseCommand("with my_cte (bar/"));
+
+    assertEquals("WITH", parseCommand("with my_cte as (select * from foo)"));
+    assertEquals("", parseCommand("/*only a comment*/"));
+
+    assertEquals(
+        "UPDATE", parseCommand("with my_cte as (select a/b from foo) update bar set col1='one'"));
+    assertEquals(
+        "UPDATE", parseCommand("with my_cte as (select a*b from foo) update bar set col1='one'"));
+    assertEquals(
+        "UPDATE", parseCommand("with my_cte as (select a-b from foo) update bar set col1='one'"));
+    assertEquals(
+        "UPDATE", parseCommand("with my_cte as (select a / b from foo) update bar set col1='one'"));
+    assertEquals(
+        "UPDATE", parseCommand("with my_cte as (select a * b from foo) update bar set col1='one'"));
+    assertEquals(
+        "UPDATE", parseCommand("with my_cte as (select a - b from foo) update bar set col1='one'"));
+  }
 }
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/utils/StatementParserTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/utils/StatementParserTest.java
index da13b524..09142ef5 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/utils/StatementParserTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/utils/StatementParserTest.java
@@ -14,12 +14,12 @@
 
 package com.google.cloud.spanner.pgadapter.utils;
 
-import static com.google.cloud.spanner.pgadapter.utils.StatementParser.singleQuoteEscape;
-import static com.google.cloud.spanner.pgadapter.utils.StatementParser.skipDollarQuotedString;
-import static com.google.cloud.spanner.pgadapter.utils.StatementParser.skipMultiLineComment;
-import static com.google.cloud.spanner.pgadapter.utils.StatementParser.skipQuotedString;
-import static com.google.cloud.spanner.pgadapter.utils.StatementParser.skipSingleLineComment;
-import static com.google.cloud.spanner.pgadapter.utils.StatementParser.splitStatements;
+import static com.google.cloud.spanner.pgadapter.commands.DynamicCommand.singleQuoteEscape;
+import static com.google.cloud.spanner.pgadapter.statements.SimpleParserTest.skipDollarQuotedString;
+import static com.google.cloud.spanner.pgadapter.statements.SimpleParserTest.skipMultiLineComment;
+import static com.google.cloud.spanner.pgadapter.statements.SimpleParserTest.skipQuotedString;
+import static com.google.cloud.spanner.pgadapter.statements.SimpleParserTest.skipSingleLineComment;
+import static com.google.cloud.spanner.pgadapter.statements.SimpleParserTest.splitStatements;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertThrows;
 
@@ -194,6 +194,9 @@ public class StatementParserTest {
     assertEquals(15, skipDollarQuotedString("$bar$foo$$, bar", 0));
 
     assertEquals(11, skipDollarQuotedString("$$foo$bar$$, bar", 0));
+    assertEquals(5, skipDollarQuotedString("$$foo", 0));
+    assertEquals(0, skipDollarQuotedString("foo", 0));
+    assertEquals(1, skipDollarQuotedString("$-$bar$-$", 0));
   }
 
   @Test