Repository: GoogleCloudPlatform/pgadapter
Pull request #720 (2023-03-13 14:34:57)

chore: allow client-detection to skip the first N messages
Some clients might need to skip the first N Parse messages that we receive before doing the actual client detection. This is for example the case if the client first executes a number of very generic statements, like BEGIN or SHOW ISOLATION LEVEL.

codecov[bot]: ## [Codecov](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/720?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) Report
> Merging [#720](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/720?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) (c6822c1) into [postgresql-dialect](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/commit/08959132e983ad6deebfc9716398d980c6d2b0ca?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) (0895913) will **decrease** coverage by `0.01%`.
> The diff coverage is `88.88%`.

```diff
@@                   Coverage Diff                    @@
##             postgresql-dialect     #720      +/-   ##
========================================================
- Coverage                 89.49%   89.49%   -0.01%     
- Complexity                 2349     2350       +1     
========================================================
  Files                       129      129              
  Lines                      7719     7728       +9     
  Branches                   1105     1107       +2     
========================================================
+ Hits                       6908     6916       +8     
  Misses                      563      563              
- Partials                    248      249       +1     
```

| Flag | Coverage Δ | |
|---|---|---|
| all_tests | `89.49% <88.88%> (-0.01%)` | :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/720?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) | Coverage Δ | |
|---|---|---|
| [...gle/cloud/spanner/pgadapter/ConnectionHandler.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/720?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL0Nvbm5lY3Rpb25IYW5kbGVyLmphdmE=) | `88.10% <87.50%> (+0.01%)` | :arrow_up: |
| [...ud/spanner/pgadapter/utils/ClientAutoDetector.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/720?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3V0aWxzL0NsaWVudEF1dG9EZXRlY3Rvci5qYXZh) | `91.91% <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/ConnectionHandler.java b/src/main/java/com/google/cloud/spanner/pgadapter/ConnectionHandler.java
index 5f3fc7d5..94a0bcbc 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/ConnectionHandler.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/ConnectionHandler.java
@@ -30,6 +30,7 @@ import com.google.cloud.spanner.SpannerException;
 import com.google.cloud.spanner.SpannerException.ResourceNotFoundException;
 import com.google.cloud.spanner.SpannerExceptionFactory;
 import com.google.cloud.spanner.Statement;
+import com.google.cloud.spanner.connection.AbstractStatementParser.StatementType;
 import com.google.cloud.spanner.connection.Connection;
 import com.google.cloud.spanner.connection.ConnectionOptions;
 import com.google.cloud.spanner.connection.ConnectionOptionsHelper;
@@ -71,6 +72,8 @@ import java.text.MessageFormat;
 import java.time.Duration;
 import java.util.HashMap;
 import java.util.HashSet;
+import java.util.LinkedList;
+import java.util.List;
 import java.util.Locale;
 import java.util.Map;
 import java.util.Map.Entry;
@@ -125,6 +128,13 @@ public class ConnectionHandler extends Thread {
   private DatabaseId databaseId;
   private WellKnownClient wellKnownClient = WellKnownClient.UNSPECIFIED;
   private boolean hasDeterminedClientUsingQuery;
+  /**
+   * List of PARSE messages that we received before auto-detecting the client. This list can be used
+   * by the detector to determine which client is connected, and is cleared after the detection is
+   * done.
+   */
+  private final LinkedList<ParseMessage> skippedAutoDetectParseMessages = new LinkedList<>();
+
   private ExtendedQueryProtocolHandler extendedQueryProtocolHandler;
   private CopyStatement activeCopyStatement;
 
@@ -727,12 +737,15 @@ public class ConnectionHandler extends Thread {
     if (!this.hasDeterminedClientUsingQuery) {
       if (this.wellKnownClient == WellKnownClient.UNSPECIFIED
           && getServer().getOptions().shouldAutoDetectClient()) {
-        setWellKnownClient(ClientAutoDetector.detectClient(ImmutableList.of(statement)));
+        setWellKnownClient(
+            ClientAutoDetector.detectClient(
+                skippedAutoDetectParseMessages, ImmutableList.of(statement)));
       }
       maybeSetApplicationName();
+      skippedAutoDetectParseMessages.clear();
+      // Make sure that we only try to detect the client once.
+      this.hasDeterminedClientUsingQuery = true;
     }
-    // Make sure that we only try to detect the client once.
-    this.hasDeterminedClientUsingQuery = true;
   }
 
   /**
@@ -743,14 +756,28 @@ public class ConnectionHandler extends Thread {
    */
   public void maybeDetermineWellKnownClient(ParseMessage parseMessage) {
     if (!this.hasDeterminedClientUsingQuery) {
-      if (this.wellKnownClient == WellKnownClient.UNSPECIFIED
-          && getServer().getOptions().shouldAutoDetectClient()) {
-        setWellKnownClient(ClientAutoDetector.detectClient(parseMessage));
+      // We skip up to 10 Parse messages before forcing the connection to detect a client (or not).
+      // This is a safety measure against clients that might send a very large number of Parse
+      // messages with client-side statements directly after connecting. This could in worst case
+      // cause PGAdapter to buffer an unreasonable number of messages. The number 10 was randomly
+      // chosen as a reasonable number of statements that should be enough. It can safely be
+      // increased if we encounter clients that send more than 10 client-side statements before
+      // sending anything that we can use to automatically recognize them.
+      if (parseMessage.getStatement().getStatementType() == StatementType.CLIENT_SIDE
+          && skippedAutoDetectParseMessages.size() < 10) {
+        skippedAutoDetectParseMessages.add(parseMessage);
+      } else {
+        if (this.wellKnownClient == WellKnownClient.UNSPECIFIED
+            && getServer().getOptions().shouldAutoDetectClient()) {
+          setWellKnownClient(
+              ClientAutoDetector.detectClient(skippedAutoDetectParseMessages, parseMessage));
+        }
+        maybeSetApplicationName();
+        skippedAutoDetectParseMessages.clear();
+        // Make sure that we only try to detect the client once.
+        this.hasDeterminedClientUsingQuery = true;
       }
-      maybeSetApplicationName();
     }
-    // Make sure that we only try to detect the client once.
-    this.hasDeterminedClientUsingQuery = true;
   }
 
   private void maybeSetApplicationName() {
@@ -776,6 +803,16 @@ public class ConnectionHandler extends Thread {
     }
   }
 
+  @VisibleForTesting
+  List<ParseMessage> getSkippedAutoDetectParseMessages() {
+    return this.skippedAutoDetectParseMessages;
+  }
+
+  @VisibleForTesting
+  boolean isHasDeterminedClientUsingQuery() {
+    return this.hasDeterminedClientUsingQuery;
+  }
+
   /** Status of a {@link ConnectionHandler} */
   public enum ConnectionStatus {
     UNAUTHENTICATED,
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/utils/ClientAutoDetector.java b/src/main/java/com/google/cloud/spanner/pgadapter/utils/ClientAutoDetector.java
index 3d9671e9..08b6646a 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/utils/ClientAutoDetector.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/utils/ClientAutoDetector.java
@@ -174,7 +174,7 @@ public class ClientAutoDetector {
       }
 
       @Override
-      boolean isClient(ParseMessage parseMessage) {
+      boolean isClient(List<ParseMessage> skippedParseMessages, ParseMessage parseMessage) {
         // pgx uses a relatively unique naming scheme for prepared statements (and uses prepared
         // statements for everything by default).
         return parseMessage.getName() != null && parseMessage.getName().startsWith("lrupsc_");
@@ -207,7 +207,7 @@ public class ClientAutoDetector {
       }
 
       @Override
-      boolean isClient(List<Statement> statements) {
+      boolean isClient(List<ParseMessage> skippedParseMessages, List<Statement> statements) {
         // The npgsql client always starts with sending a query that contains multiple statements
         // and that starts with the following prefix.
         return statements.size() == 1
@@ -244,14 +244,14 @@ public class ClientAutoDetector {
       }
 
       @Override
-      boolean isClient(List<Statement> statements) {
+      boolean isClient(List<ParseMessage> skippedParseMessages, List<Statement> statements) {
         // Use UNSPECIFIED as default to prevent null checks everywhere and to ease the use of any
         // defaults defined in this enum.
         return true;
       }
 
       @Override
-      boolean isClient(ParseMessage parseMessage) {
+      boolean isClient(List<ParseMessage> skippedParseMessages, ParseMessage parseMessage) {
         // Use UNSPECIFIED as default to prevent null checks everywhere and to ease the use of any
         // defaults defined in this enum.
         return true;
@@ -264,11 +264,11 @@ public class ClientAutoDetector {
     @VisibleForTesting
     public void reset() {}
 
-    boolean isClient(List<Statement> statements) {
+    boolean isClient(List<ParseMessage> skippedParseMessages, List<Statement> statements) {
       return false;
     }
 
-    boolean isClient(ParseMessage parseMessage) {
+    boolean isClient(List<ParseMessage> skippedParseMessages, ParseMessage parseMessage) {
       return false;
     }
 
@@ -331,9 +331,10 @@ public class ClientAutoDetector {
    * Returns the {@link WellKnownClient} that the detector thinks is connected to PGAdapter based on
    * the given list of SQL statements that have been executed.
    */
-  public static @Nonnull WellKnownClient detectClient(List<Statement> statements) {
+  public static @Nonnull WellKnownClient detectClient(
+      List<ParseMessage> skippedParseMessages, List<Statement> statements) {
     for (WellKnownClient client : WellKnownClient.values()) {
-      if (client.isClient(statements)) {
+      if (client.isClient(skippedParseMessages, statements)) {
         return client;
       }
     }
@@ -345,9 +346,10 @@ public class ClientAutoDetector {
    * Returns the {@link WellKnownClient} that the detector thinks is connected to PGAdapter based on
    * the Parse message that has been received.
    */
-  public static @Nonnull WellKnownClient detectClient(ParseMessage parseMessage) {
+  public static @Nonnull WellKnownClient detectClient(
+      List<ParseMessage> skippedParseMessages, ParseMessage parseMessage) {
     for (WellKnownClient client : WellKnownClient.values()) {
-      if (client.isClient(parseMessage)) {
+      if (client.isClient(skippedParseMessages, parseMessage)) {
         return client;
       }
     }


Test patch:
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/ConnectionHandlerTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/ConnectionHandlerTest.java
index e5a5446e..b44c9679 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/ConnectionHandlerTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/ConnectionHandlerTest.java
@@ -27,6 +27,7 @@ import static org.mockito.Mockito.when;
 import com.google.cloud.spanner.ErrorCode;
 import com.google.cloud.spanner.SpannerExceptionFactory;
 import com.google.cloud.spanner.Statement;
+import com.google.cloud.spanner.connection.AbstractStatementParser.StatementType;
 import com.google.cloud.spanner.connection.Connection;
 import com.google.cloud.spanner.pgadapter.ConnectionHandler.ConnectionStatus;
 import com.google.cloud.spanner.pgadapter.error.PGException;
@@ -38,7 +39,9 @@ import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata.SslMode;
 import com.google.cloud.spanner.pgadapter.statements.IntermediatePortalStatement;
 import com.google.cloud.spanner.pgadapter.statements.IntermediatePreparedStatement;
 import com.google.cloud.spanner.pgadapter.utils.ClientAutoDetector.WellKnownClient;
+import com.google.cloud.spanner.pgadapter.wireprotocol.ParseMessage;
 import com.google.cloud.spanner.pgadapter.wireprotocol.WireMessage;
+import com.google.common.collect.ImmutableList;
 import java.io.ByteArrayOutputStream;
 import java.io.DataOutputStream;
 import java.io.EOFException;
@@ -592,4 +595,64 @@ public class ConnectionHandlerTest {
 
     assertEquals(WellKnownClient.UNSPECIFIED, connectionHandler.getWellKnownClient());
   }
+
+  @Test
+  public void testMaybeDetermineWellKnownClient_skipsClientSideStatements() {
+    OptionsMetadata options = mock(OptionsMetadata.class);
+    ProxyServer server = mock(ProxyServer.class);
+    when(server.getOptions()).thenReturn(options);
+    Socket socket = mock(Socket.class);
+    InetAddress address = mock(InetAddress.class);
+    when(socket.getInetAddress()).thenReturn(address);
+    ConnectionHandler connectionHandler = new ConnectionHandler(server, socket);
+
+    when(options.shouldAutoDetectClient()).thenReturn(true);
+    connectionHandler.setWellKnownClient(WellKnownClient.UNSPECIFIED);
+
+    ParseMessage parseMessage = mock(ParseMessage.class);
+    IntermediatePreparedStatement statement = mock(IntermediatePreparedStatement.class);
+    when(statement.getStatementType()).thenReturn(StatementType.CLIENT_SIDE);
+    when(parseMessage.getStatement()).thenReturn(statement);
+
+    connectionHandler.maybeDetermineWellKnownClient(parseMessage);
+
+    assertEquals(WellKnownClient.UNSPECIFIED, connectionHandler.getWellKnownClient());
+    assertFalse(connectionHandler.isHasDeterminedClientUsingQuery());
+    assertEquals(
+        ImmutableList.of(parseMessage), connectionHandler.getSkippedAutoDetectParseMessages());
+  }
+
+  @Test
+  public void testMaybeDetermineWellKnownClient_stopsSkippingParseMessagesAfter10Messages() {
+    OptionsMetadata options = mock(OptionsMetadata.class);
+    ProxyServer server = mock(ProxyServer.class);
+    when(server.getOptions()).thenReturn(options);
+    Socket socket = mock(Socket.class);
+    InetAddress address = mock(InetAddress.class);
+    when(socket.getInetAddress()).thenReturn(address);
+    ConnectionHandler connectionHandler = new ConnectionHandler(server, socket);
+
+    when(options.shouldAutoDetectClient()).thenReturn(true);
+    connectionHandler.setWellKnownClient(WellKnownClient.UNSPECIFIED);
+
+    ParseMessage parseMessage = mock(ParseMessage.class);
+    IntermediatePreparedStatement statement = mock(IntermediatePreparedStatement.class);
+    when(statement.getStatementType()).thenReturn(StatementType.CLIENT_SIDE);
+    when(parseMessage.getStatement()).thenReturn(statement);
+
+    for (int i = 0; i < 10; i++) {
+      connectionHandler.maybeDetermineWellKnownClient(parseMessage);
+    }
+
+    assertEquals(WellKnownClient.UNSPECIFIED, connectionHandler.getWellKnownClient());
+    assertFalse(connectionHandler.isHasDeterminedClientUsingQuery());
+    assertEquals(10, connectionHandler.getSkippedAutoDetectParseMessages().size());
+
+    // This should cause the connection to stop buffering Parse messages and just set the client to
+    // UNSPECIFIED.
+    connectionHandler.maybeDetermineWellKnownClient(parseMessage);
+
+    assertTrue(connectionHandler.isHasDeterminedClientUsingQuery());
+    assertTrue(connectionHandler.getSkippedAutoDetectParseMessages().isEmpty());
+  }
 }
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/utils/ClientAutoDetectorTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/utils/ClientAutoDetectorTest.java
index 790258ed..30241516 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/utils/ClientAutoDetectorTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/utils/ClientAutoDetectorTest.java
@@ -309,6 +309,7 @@ public class ClientAutoDetectorTest {
     assertEquals(
         WellKnownClient.NPGSQL,
         ClientAutoDetector.detectClient(
+            ImmutableList.of(),
             ImmutableList.of(
                 Statement.of(
                     "SELECT version();\n"
@@ -316,10 +317,12 @@ public class ClientAutoDetectorTest {
                         + "SELECT ns.nspname, t.oid, t.typname, t.typtype, t.typnotnull, t.elemtypoid\n"))));
     assertEquals(
         WellKnownClient.UNSPECIFIED,
-        ClientAutoDetector.detectClient(ImmutableList.of(Statement.of("SELECT version()"))));
+        ClientAutoDetector.detectClient(
+            ImmutableList.of(), ImmutableList.of(Statement.of("SELECT version()"))));
     assertEquals(
         WellKnownClient.UNSPECIFIED,
         ClientAutoDetector.detectClient(
+            ImmutableList.of(),
             ImmutableList.of(
                 Statement.of(
                     "SELECT version();\n"