Repository: GoogleCloudPlatform/pgadapter
Pull request #837 (2023-05-17 08:54:49)

feat: support setting the well-known client
PGAdapter automatically tries to detect the well-known client that is connected to it based on a range of startup parameters, first queries etc., but not all clients can reliably be detected based on that. This change adds support for setting the well-known client as a connection startup parameter.

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/837?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) Report
> Merging [#837](https://app.codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/837?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) (c53471c) into [postgresql-dialect](https://app.codecov.io/gh/GoogleCloudPlatform/pgadapter/commit/af98023a7363d6a2bb2bfb356b747f3ed8499e0e?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) (af98023) will **decrease** coverage by `0.03%`.
> The diff coverage is `85.71%`.

```diff
@@                   Coverage Diff                    @@
##             postgresql-dialect     #837      +/-   ##
========================================================
- Coverage                 89.75%   89.72%   -0.03%     
- Complexity                 2382     2388       +6     
========================================================
  Files                       131      131              
  Lines                      7896     7924      +28     
  Branches                   1133     1140       +7     
========================================================
+ Hits                       7087     7110      +23     
- Misses                      565      567       +2     
- Partials                    244      247       +3     
```

| Flag | Coverage Δ | |
|---|---|---|
| all_tests | `89.72% <85.71%> (-0.03%)` | :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://app.codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/837?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) | Coverage Δ | |
|---|---|---|
| [...ud/spanner/pgadapter/utils/ClientAutoDetector.java](https://app.codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/837?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3V0aWxzL0NsaWVudEF1dG9EZXRlY3Rvci5qYXZh) | `90.55% <63.63%> (-2.56%)` | :arrow_down: |
| [.../cloud/spanner/pgadapter/session/SessionState.java](https://app.codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/837?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3Nlc3Npb24vU2Vzc2lvblN0YXRlLmphdmE=) | `94.94% <100.00%> (+0.02%)` | :arrow_up: |
| [...spanner/pgadapter/wireprotocol/StartupMessage.java](https://app.codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/837?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3dpcmVwcm90b2NvbC9TdGFydHVwTWVzc2FnZS5qYXZh) | `90.62% <100.00%> (+3.12%)` | :arrow_up: |

... and [1 file with indirect coverage changes](https://app.codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/837/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform)

: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/session/SessionState.java b/src/main/java/com/google/cloud/spanner/pgadapter/session/SessionState.java
index 4d8c1919..6daa754f 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/session/SessionState.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/session/SessionState.java
@@ -43,6 +43,7 @@ import java.util.Objects;
 import java.util.Set;
 import java.util.concurrent.Callable;
 import java.util.stream.Collectors;
+import javax.annotation.Nullable;
 
 /** {@link SessionState} contains all session variables for a connection. */
 @InternalApi
@@ -267,6 +268,12 @@ public class SessionState {
     return internalGet(toKey(extension, name), true);
   }
 
+  /** Returns the current value of the specified setting or null if undefined. */
+  @Nullable
+  public PGSetting tryGet(String extension, String name) {
+    return internalGet(toKey(extension, name), false);
+  }
+
   private PGSetting internalGet(String key, boolean throwForUnknownParam) {
     if (localSettings != null && localSettings.containsKey(key)) {
       return localSettings.get(key);
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 7cf1c38c..05a797b8 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
@@ -19,6 +19,7 @@ import com.google.cloud.spanner.Statement;
 import com.google.cloud.spanner.pgadapter.ConnectionHandler;
 import com.google.cloud.spanner.pgadapter.error.PGException;
 import com.google.cloud.spanner.pgadapter.error.SQLState;
+import com.google.cloud.spanner.pgadapter.session.PGSetting;
 import com.google.cloud.spanner.pgadapter.statements.PgCatalog.PgCatalogTable;
 import com.google.cloud.spanner.pgadapter.statements.local.DjangoGetTableNamesStatement;
 import com.google.cloud.spanner.pgadapter.statements.local.ListDatabasesStatement;
@@ -38,8 +39,10 @@ import com.google.common.collect.ImmutableSet;
 import java.time.Duration;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.regex.Pattern;
 import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
 import org.postgresql.core.Oid;
 
 /**
@@ -253,6 +256,10 @@ public class ClientAutoDetector {
         if (!parameters.get("client_encoding").equals("UTF8")) {
           return false;
         }
+        if (parameters.get("options") != null
+            && parameters.get("options").contains("spanner.well_known_client")) {
+          return false;
+        }
         return parameters.get("DateStyle").equals("ISO");
       }
 
@@ -355,26 +362,40 @@ public class ClientAutoDetector {
       boolean isClient(List<String> orderedParameterKeys, Map<String, String> parameters) {
         // Use UNSPECIFIED as default to prevent null checks everywhere and to ease the use of any
         // defaults defined in this enum.
-        return true;
+        return DEFAULT_UNSPECIFIED.get();
       }
 
       @Override
       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;
+        return DEFAULT_UNSPECIFIED.get();
       }
 
       @Override
       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;
+        return DEFAULT_UNSPECIFIED.get();
+      }
+
+      @Override
+      boolean isClient(PGSetting setting) {
+        // Use UNSPECIFIED as default to prevent null checks everywhere and to ease the use of any
+        // defaults defined in this enum.
+        return DEFAULT_UNSPECIFIED.get();
       }
     };
 
+    /** Indicates whether UNSPECIFIED should be used as default (instead of <code>null</code>). */
+    @VisibleForTesting static final AtomicBoolean DEFAULT_UNSPECIFIED = new AtomicBoolean(true);
+
     abstract boolean isClient(List<String> orderedParameterKeys, Map<String, String> parameters);
 
+    boolean isClient(PGSetting setting) {
+      return false;
+    }
+
     /** Resets any cached or temporary settings for the client. */
     @VisibleForTesting
     public void reset() {}
@@ -471,4 +492,18 @@ public class ClientAutoDetector {
     // The following line should never be reached.
     throw new IllegalStateException("UNSPECIFIED.isClient() should have returned true");
   }
+
+  /** Detect the client based on a session state setting. */
+  public static @Nonnull WellKnownClient detectClient(@Nullable PGSetting setting) {
+    if (setting == null || setting.getSetting() == null) {
+      return WellKnownClient.UNSPECIFIED;
+    }
+    for (WellKnownClient client : WellKnownClient.values()) {
+      if (client.name().equalsIgnoreCase(setting.getSetting()) || client.isClient(setting)) {
+        return client;
+      }
+    }
+    // The following line should never be reached.
+    throw new IllegalStateException("UNSPECIFIED.isClient() should have returned true");
+  }
 }
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/StartupMessage.java b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/StartupMessage.java
index be2832f4..327cf7a3 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/StartupMessage.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/StartupMessage.java
@@ -83,6 +83,25 @@ public class StartupMessage extends BootstrapMessage {
           .getBackendConnection()
           .initSessionSetting(parameter.getKey(), parameter.getValue());
     }
+    if (connection.getServer().getOptions().shouldAutoDetectClient()
+        && connection.getWellKnownClient() == WellKnownClient.UNSPECIFIED) {
+      WellKnownClient wellKnownClient =
+          ClientAutoDetector.detectClient(
+              connection
+                  .getExtendedQueryProtocolHandler()
+                  .getBackendConnection()
+                  .getSessionState()
+                  .tryGet("spanner", "well_known_client"));
+      connection.setWellKnownClient(wellKnownClient);
+    }
+    if (connection.getWellKnownClient() != WellKnownClient.UNSPECIFIED) {
+      connection
+          .getExtendedQueryProtocolHandler()
+          .getBackendConnection()
+          .getSessionState()
+          .setConnectionStartupValue(
+              "spanner", "well_known_client", connection.getWellKnownClient().name());
+    }
     sendStartupMessage(
         connection.getConnectionMetadata().getOutputStream(),
         connection.getConnectionId(),
diff --git a/src/main/resources/com/google/cloud/spanner/pgadapter/session/pg_settings.txt b/src/main/resources/com/google/cloud/spanner/pgadapter/session/pg_settings.txt
index 439e4253..59e6334d 100644
--- a/src/main/resources/com/google/cloud/spanner/pgadapter/session/pg_settings.txt
+++ b/src/main/resources/com/google/cloud/spanner/pgadapter/session/pg_settings.txt
@@ -357,3 +357,4 @@ spanner.ddl_transaction_mode	Batch	\N	PGAdapter Options	DDL transaction mode.	\N
 spanner.emulate_pg_class_tables	on	\N	PGAdapter Options Emulate pg_class and related tables using common table expressions and textual OIDs	Generate commont table expressions for pg_class and related tables. The OID of an object is set to its fully qualified name.	\N	user	bool	default	\N	\N	\N	off	off	\N	\N	f
 spanner.force_autocommit	off	\N	PGAdapter Options Execute all statements in autocommit mode	Ignore all transaction statements (BEGIN / COMMIT / ROLLBACK) and execute all statements in autocommit mode.	\N	user	bool	default	\N	\N	\N	off	off	\N	\N	f
 spanner.replace_pg_catalog_tables	on	\N	PGAdapter Options	Automatically replace pg_catalog table references with Common Table Expressions.	\N	user	bool	default	\N	\N	\N	on	on	\N	\N	f
+spanner.well_known_client	UNSPECIFIED	\N	PGAdapter Options	The well-known client that is connected to PGAdapter.	\N	backend	string	default	\N	\N	\N	UNSPECIFIED	UNSPECIFIED	\N	\N	f


Test patch:
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/AbortedMockServerTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/AbortedMockServerTest.java
index d22ccda7..28973ef6 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/AbortedMockServerTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/AbortedMockServerTest.java
@@ -1319,7 +1319,7 @@ public class AbortedMockServerTest extends AbstractMockServerTest {
           }
           count++;
         }
-        assertEquals(360, count);
+        assertEquals(361, count);
       }
     }
   }
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 c4f51505..9bc74318 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/JdbcMockServerTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/JdbcMockServerTest.java
@@ -387,6 +387,44 @@ public class JdbcMockServerTest extends AbstractMockServerTest {
     }
   }
 
+  @Test
+  public void testSetWellKnownClient() throws SQLException {
+    for (String client : new String[] {"pgx", "npgsql", "sqlalchemy2"}) {
+      try (Connection connection =
+          DriverManager.getConnection(
+              String.format(
+                  "jdbc:postgresql://localhost:%d/?options=-c%%20spanner.well_known_client=%s",
+                  pgServer.getLocalPort(), client))) {
+        try (ResultSet resultSet =
+            connection.createStatement().executeQuery("show spanner.well_known_client")) {
+          assertTrue(resultSet.next());
+          assertEquals(client.toUpperCase(), resultSet.getString(1));
+          assertFalse(resultSet.next());
+        }
+        PSQLException exception =
+            assertThrows(
+                PSQLException.class,
+                () -> connection.createStatement().execute("set spanner.well_known_client='foo'"));
+        assertNotNull(exception.getServerErrorMessage());
+        assertEquals(
+            "parameter \"spanner.well_known_client\" cannot be set after connection start",
+            exception.getServerErrorMessage().getMessage());
+      }
+    }
+    try (Connection connection =
+        DriverManager.getConnection(
+            String.format(
+                "jdbc:postgresql://localhost:%d/?options=-c%%20spanner.well_known_client=%s",
+                pgServer.getLocalPort(), "foo"))) {
+      try (ResultSet resultSet =
+          connection.createStatement().executeQuery("show spanner.well_known_client")) {
+        assertTrue(resultSet.next());
+        assertEquals("foo", resultSet.getString(1));
+        assertFalse(resultSet.next());
+      }
+    }
+  }
+
   @Test
   public void testPreparedStatementParameterMetadata() throws SQLException {
     String sql = "SELECT * FROM foo WHERE id=? or value=?";
@@ -3873,7 +3911,7 @@ public class JdbcMockServerTest extends AbstractMockServerTest {
           }
           count++;
         }
-        assertEquals(360, count);
+        assertEquals(361, count);
       }
     }
   }
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/session/SessionStateTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/session/SessionStateTest.java
index e69b5744..e69c9415 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/session/SessionStateTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/session/SessionStateTest.java
@@ -161,7 +161,7 @@ public class SessionStateTest {
   public void testGetAll() {
     SessionState state = new SessionState(mock(OptionsMetadata.class));
     List<PGSetting> allSettings = state.getAll();
-    assertEquals(359, allSettings.size());
+    assertEquals(360, allSettings.size());
   }
 
   @Test
@@ -199,7 +199,7 @@ public class SessionStateTest {
     state.setLocal("spanner", "custom_local_setting", "value2");
 
     List<PGSetting> allSettings = state.getAll();
-    assertEquals(361, allSettings.size());
+    assertEquals(362, allSettings.size());
 
     PGSetting applicationName =
         allSettings.stream()
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 20d30ed5..b4092455 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
@@ -18,6 +18,7 @@ import static com.google.cloud.spanner.pgadapter.utils.ClientAutoDetector.DEFAUL
 import static com.google.cloud.spanner.pgadapter.utils.ClientAutoDetector.EMPTY_LOCAL_STATEMENTS;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotEquals;
+import static org.junit.Assert.assertThrows;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
@@ -27,6 +28,7 @@ import com.google.cloud.spanner.pgadapter.ProxyServer;
 import com.google.cloud.spanner.pgadapter.error.PGException;
 import com.google.cloud.spanner.pgadapter.metadata.ConnectionMetadata;
 import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata;
+import com.google.cloud.spanner.pgadapter.session.PGSetting;
 import com.google.cloud.spanner.pgadapter.statements.local.ListDatabasesStatement;
 import com.google.cloud.spanner.pgadapter.utils.ClientAutoDetector.WellKnownClient;
 import com.google.cloud.spanner.pgadapter.wireoutput.NoticeResponse;
@@ -433,4 +435,30 @@ public class ClientAutoDetectorTest {
     assertEquals("Detected connection from pgbench", startupNotices.get(0).getMessage());
     assertEquals(ClientAutoDetector.PGBENCH_USAGE_HINT + "\n", startupNotices.get(0).getHint());
   }
+
+  @Test
+  public void testSetting() {
+    assertEquals(WellKnownClient.UNSPECIFIED, ClientAutoDetector.detectClient(null));
+    assertEquals(
+        WellKnownClient.UNSPECIFIED, ClientAutoDetector.detectClient(mock(PGSetting.class)));
+
+    PGSetting jdbc = mock(PGSetting.class);
+    when(jdbc.getSetting()).thenReturn("jdbc");
+    assertEquals(WellKnownClient.JDBC, ClientAutoDetector.detectClient(jdbc));
+
+    PGSetting unspecified = mock(PGSetting.class);
+    when(unspecified.getSetting()).thenReturn("unspecified");
+    assertEquals(WellKnownClient.UNSPECIFIED, ClientAutoDetector.detectClient(unspecified));
+
+    PGSetting foo = mock(PGSetting.class);
+    when(foo.getSetting()).thenReturn("foo");
+    assertEquals(WellKnownClient.UNSPECIFIED, ClientAutoDetector.detectClient(foo));
+
+    try {
+      WellKnownClient.DEFAULT_UNSPECIFIED.set(false);
+      assertThrows(IllegalStateException.class, () -> ClientAutoDetector.detectClient(foo));
+    } finally {
+      WellKnownClient.DEFAULT_UNSPECIFIED.set(true);
+    }
+  }
 }