Repository: GoogleCloudPlatform/pgadapter
Pull request #3275 (2025-05-08 11:57:21)

feat: allow Spanner startup parameters in 'options=...'
PostgreSQL allows clients to set connection variables in the connection string by adding an `options=-c name=value` element to the connection string. PGAdapter also accepted these, but only for variables that could be set with a `SET` statement. Startup properties in the Spanner Connection API were not supported, which meant that properties like `options=-c enableEndToEndTracing=true` would have no effect.

This change adds support for setting startup parameters with the `options` argument in the PostgreSQL connection string.


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 a97dc230..8c92327d 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/ConnectionHandler.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/ConnectionHandler.java
@@ -38,6 +38,7 @@ import com.google.cloud.spanner.connection.AbstractStatementParser.ParsedStateme
 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.ConnectionProperties;
 import com.google.cloud.spanner.connection.PGAdapterConnectionOptionsHelper;
 import com.google.cloud.spanner.connection.SavepointSupport;
 import com.google.cloud.spanner.pgadapter.error.PGException;
@@ -200,9 +201,10 @@ public class ConnectionHandler implements Runnable {
   }
 
   @InternalApi
-  public void connectToSpanner(String database, @Nullable Credentials credentials) {
+  public void connectToSpanner(
+      String database, @Nullable Credentials credentials, Map<String, String> parameters) {
     OptionsMetadata options = getServer().getOptions();
-    String uri = buildConnectionURL(database, options, getServer().getProperties());
+    String uri = buildConnectionURL(database, options, getServer().getProperties(), parameters);
     ConnectionOptions.Builder connectionOptionsBuilder = ConnectionOptions.newBuilder().setUri(uri);
     connectionOptionsBuilder =
         PGAdapterConnectionOptionsHelper.maybeAddGrpcLogInterceptor(
@@ -279,12 +281,16 @@ public class ConnectionHandler implements Runnable {
 
   @VisibleForTesting
   static String buildConnectionURL(
-      String database, OptionsMetadata options, Properties properties) {
+      String database,
+      OptionsMetadata options,
+      Properties properties,
+      Map<String, String> startupParameters) {
     String uri =
         options.hasDefaultConnectionUrl()
             ? options.getDefaultConnectionUrl()
             : options.buildConnectionURL(database);
     uri = appendPropertiesToUrl(uri, properties);
+    uri = appendStartupParametersToUrl(uri, startupParameters);
     // We add 'dialect=postgresql' here in all cases, although it is only being used if
     // 'autoConfigEmulator=true' has also been set in the connection URL. This dialect property is
     // only used to determine what dialect the database should have that is being created on the
@@ -331,6 +337,40 @@ public class ConnectionHandler implements Runnable {
     return result.toString();
   }
 
+  /**
+   * Adds parameters to the connection URL that have been added as 'options=-c ...' to the
+   * PostgreSQL connection string.
+   */
+  static String appendStartupParametersToUrl(String url, Map<String, String> startupParameters) {
+    if (startupParameters == null || startupParameters.isEmpty()) {
+      return url;
+    }
+    StringBuilder result = new StringBuilder(url);
+    for (Map.Entry<String, String> entry : startupParameters.entrySet()) {
+      if ("options".equalsIgnoreCase(entry.getKey()) && entry.getValue() != null) {
+        String[] commands = entry.getValue().split("-c[\\s+]+");
+        for (String command : commands) {
+          String[] keyValue = command.split("=", 2);
+          if (keyValue.length == 2
+              && isValidConnectionProperty(keyValue[0])
+              && keyValue[1] != null) {
+            result.append(";").append(keyValue[0]).append("=").append(keyValue[1].trim());
+          }
+        }
+      }
+    }
+    return result.toString();
+  }
+
+  static boolean isValidConnectionProperty(String name) {
+    if (Strings.isNullOrEmpty(name)) {
+      return false;
+    }
+    String key = name.toLowerCase();
+    return ConnectionProperties.VALID_CONNECTION_PROPERTIES.stream()
+        .anyMatch(property -> property.getKey().equals(key));
+  }
+
   /**
    * Simple runner starts a loop which keeps taking inbound messages, processing them, sending them
    * to Spanner, getting a result, processing that result, and replying to the client (in that
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 45273a26..e087ff5d 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
@@ -72,7 +72,7 @@ public class StartupMessage extends BootstrapMessage {
       Map<String, String> parameters,
       @Nullable Credentials credentials)
       throws Exception {
-    connection.connectToSpanner(database, credentials);
+    connection.connectToSpanner(database, credentials, parameters);
     for (Entry<String, String> parameter : parameters.entrySet()) {
       connection
           .getExtendedQueryProtocolHandler()


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 d6c5508a..6057f910 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/ConnectionHandlerTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/ConnectionHandlerTest.java
@@ -665,13 +665,15 @@ public class ConnectionHandlerTest {
         buildConnectionURL(
             "projects/my-project/instances/my-instance/databases/my-database",
             options,
-            buildProperties(ImmutableMap.of())));
+            buildProperties(ImmutableMap.of()),
+            ImmutableMap.of()));
     assertEquals(
         "cloudspanner:/projects/my-project/instances/my-instance/databases/my-database;userAgent=pg-adapter;key1=value1;dialect=postgresql",
         buildConnectionURL(
             "projects/my-project/instances/my-instance/databases/my-database",
             options,
-            buildProperties(ImmutableMap.of("key1", "value1"))));
+            buildProperties(ImmutableMap.of("key1", "value1")),
+            ImmutableMap.of()));
 
     // If the options contain a full database specification, then the database in the connection
     // request is ignored.
@@ -685,7 +687,8 @@ public class ConnectionHandlerTest {
                 .setInstance("test-instance")
                 .setDatabase("test-database")
                 .build(),
-            buildProperties(ImmutableMap.of())));
+            buildProperties(ImmutableMap.of()),
+            ImmutableMap.of()));
     // Enable the autoConfigEmulator flag through the options builder.
     OptionsMetadata emulatorOptions = OptionsMetadata.newBuilder().autoConfigureEmulator().build();
     assertEquals(
@@ -693,7 +696,8 @@ public class ConnectionHandlerTest {
         buildConnectionURL(
             "projects/my-project/instances/my-instance/databases/my-database",
             emulatorOptions,
-            buildProperties(emulatorOptions.getPropertyMap())));
+            buildProperties(emulatorOptions.getPropertyMap()),
+            ImmutableMap.of()));
 
     // Set a channel provider.
     String currentChannelProvider = System.getProperty("CHANNEL_PROVIDER");
@@ -708,7 +712,8 @@ public class ConnectionHandlerTest {
           buildConnectionURL(
               "projects/my-project/instances/my-instance/databases/my-database",
               options,
-              buildProperties(ImmutableMap.of())));
+              buildProperties(ImmutableMap.of()),
+              ImmutableMap.of()));
       assertEquals("true", System.getProperty("ENABLE_CHANNEL_PROVIDER"));
 
       // Set an invalid channel provider.
@@ -720,7 +725,8 @@ public class ConnectionHandlerTest {
               buildConnectionURL(
                   "projects/my-project/instances/my-instance/databases/my-database",
                   options,
-                  buildProperties(ImmutableMap.of())));
+                  buildProperties(ImmutableMap.of()),
+                  ImmutableMap.of()));
       assertNull(System.getProperty("ENABLE_CHANNEL_PROVIDER"));
     } finally {
       if (currentChannelProvider == null) {
@@ -741,13 +747,15 @@ public class ConnectionHandlerTest {
         buildConnectionURL(
             "projects/my-project/instances/my-instance/databases/my-database",
             options,
-            buildProperties(ImmutableMap.of("useVirtualThreads", "true"))));
+            buildProperties(ImmutableMap.of("useVirtualThreads", "true")),
+            ImmutableMap.of()));
     assertEquals(
         "cloudspanner:/projects/my-project/instances/my-instance/databases/my-database;userAgent=pg-adapter;useVirtualGrpcTransportThreads=true;dialect=postgresql",
         buildConnectionURL(
             "projects/my-project/instances/my-instance/databases/my-database",
             options,
-            buildProperties(ImmutableMap.of("useVirtualGrpcTransportThreads", "true"))));
+            buildProperties(ImmutableMap.of("useVirtualGrpcTransportThreads", "true")),
+            ImmutableMap.of()));
 
     runWithSystemProperty(
         OptionsMetadata.USE_VIRTUAL_THREADS_SYSTEM_PROPERTY_NAME,
@@ -760,7 +768,8 @@ public class ConnectionHandlerTest {
               buildConnectionURL(
                   "projects/my-project/instances/my-instance/databases/my-database",
                   optionsWithSystemProperty,
-                  buildProperties(optionsWithSystemProperty.getPropertyMap())));
+                  buildProperties(optionsWithSystemProperty.getPropertyMap()),
+                  ImmutableMap.of()));
         });
     runWithSystemProperty(
         OptionsMetadata.USE_VIRTUAL_GRPC_TRANSPORT_THREADS_SYSTEM_PROPERTY_NAME,
@@ -773,8 +782,32 @@ public class ConnectionHandlerTest {
               buildConnectionURL(
                   "projects/my-project/instances/my-instance/databases/my-database",
                   optionsWithSystemProperty,
-                  buildProperties(optionsWithSystemProperty.getPropertyMap())));
+                  buildProperties(optionsWithSystemProperty.getPropertyMap()),
+                  ImmutableMap.of()));
         });
+    assertEquals(
+        "cloudspanner:/projects/my-project/instances/my-instance/databases/my-database;userAgent=pg-adapter;enableEndToEndTracing=true;dialect=postgresql",
+        buildConnectionURL(
+            "projects/my-project/instances/my-instance/databases/my-database",
+            options,
+            buildProperties(ImmutableMap.of()),
+            ImmutableMap.of("options", "-c enableEndToEndTracing=true")));
+    assertEquals(
+        "cloudspanner:/projects/my-project/instances/my-instance/databases/my-database;userAgent=pg-adapter;minSessions=10;numChannels=2;dialect=postgresql",
+        buildConnectionURL(
+            "projects/my-project/instances/my-instance/databases/my-database",
+            options,
+            buildProperties(ImmutableMap.of()),
+            ImmutableMap.of("options", "-c minSessions=10 -c numChannels=2")));
+
+    // Invalid options are not added to the connection URL.
+    assertEquals(
+        "cloudspanner:/projects/my-project/instances/my-instance/databases/my-database;userAgent=pg-adapter;numChannels=2;dialect=postgresql",
+        buildConnectionURL(
+            "projects/my-project/instances/my-instance/databases/my-database",
+            options,
+            buildProperties(ImmutableMap.of()),
+            ImmutableMap.of("options", "-c invalid_option=some-value -c numChannels=2")));
   }
 
   @Test
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/JdbcSimpleModeMockServerTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/JdbcSimpleModeMockServerTest.java
index a1010a4f..653b8ab4 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/JdbcSimpleModeMockServerTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/JdbcSimpleModeMockServerTest.java
@@ -33,6 +33,7 @@ import com.google.common.util.concurrent.MoreExecutors;
 import com.google.protobuf.ListValue;
 import com.google.protobuf.Value;
 import com.google.spanner.admin.database.v1.UpdateDatabaseDdlRequest;
+import com.google.spanner.v1.BatchCreateSessionsRequest;
 import com.google.spanner.v1.CommitRequest;
 import com.google.spanner.v1.ExecuteSqlRequest;
 import com.google.spanner.v1.ExecuteSqlRequest.QueryMode;
@@ -936,4 +937,15 @@ public class JdbcSimpleModeMockServerTest extends AbstractMockServerTest {
       statement.execute("commit");
     }
   }
+
+  @Test
+  public void testStartupConnectionPropertiesInUrl() throws SQLException {
+    try (Connection connection =
+        DriverManager.getConnection(
+            createUrl() + "&options=-c%20minSessions=10-c%20numChannels=1")) {}
+    assertEquals(1, mockSpanner.countRequestsOfType(BatchCreateSessionsRequest.class));
+    assertEquals(
+        10,
+        mockSpanner.getRequestsOfType(BatchCreateSessionsRequest.class).get(0).getSessionCount());
+  }
 }
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/wireprotocol/ProtocolTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/wireprotocol/ProtocolTest.java
index 59f3df9e..43b5d0c3 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/wireprotocol/ProtocolTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/wireprotocol/ProtocolTest.java
@@ -1636,7 +1636,7 @@ public class ProtocolTest {
     message.send();
 
     DataInputStream outputResult = inputStreamFromOutputStream(result);
-    verify(connectionHandler).connectToSpanner("databasename", null);
+    verify(connectionHandler).connectToSpanner("databasename", null, expectedParameters);
 
     // AuthenticationOkResponse
     assertEquals('R', outputResult.readByte());