Repository: GoogleCloudPlatform/pgadapter
Pull request #121 (2022-04-20 18:38:42)

feat: support connecting to different databases
Support connecting to different databases on an instance. This allows
for the use of the `\c` command in psql, as well as using a single
instance of PGAdapter to connect to multiple different databases using
other clients. The `-d` command line parameter is still supported, but
is now interpreted as follows:
1. If omitted, PGAdapter will connect to the database given in the
   connection request.
2. If included, PGAdapter will always connect to the database given in
   the -d command line parameter and ignore the database in the
   connection request.

Fixes b/194320320

aseering: Do I understand correctly that this is a "breaking change" in the sense that, if customers relied on PGAdapter's behavior of ignoring the database specified by psql (which is pretty easy to do, for example with scripts that assume they can connect to the `postgres` default database to do admin operations), previously they might have worked and now they won't?

If so, could this be accompanied by a bump to PGAdapter's published major version?  That would be a minimal way to signal a breaking change.
aseering: Ah, I see:  In your original statement, when reading:

> If included, PGAdapter will always connect to the database given in
the -d command line parameter and ignore the database in the
connection request.

I interpreted that as the `-d` argument _to psql_.  I think you actually meant the `-d` argument _to PGAdapter_?  Which in retrospect should have been evident from context; I just mis-parsed.

Anyway, agreed, SGTM.

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 98d73398..d642d9f2 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/ConnectionHandler.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/ConnectionHandler.java
@@ -15,10 +15,14 @@
 package com.google.cloud.spanner.pgadapter;
 
 import com.google.api.core.InternalApi;
+import com.google.cloud.spanner.Dialect;
+import com.google.cloud.spanner.ErrorCode;
 import com.google.cloud.spanner.SpannerException;
+import com.google.cloud.spanner.SpannerExceptionFactory;
 import com.google.cloud.spanner.connection.Connection;
 import com.google.cloud.spanner.connection.ConnectionOptions;
 import com.google.cloud.spanner.pgadapter.metadata.ConnectionMetadata;
+import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata;
 import com.google.cloud.spanner.pgadapter.statements.IntermediatePortalStatement;
 import com.google.cloud.spanner.pgadapter.statements.IntermediatePreparedStatement;
 import com.google.cloud.spanner.pgadapter.statements.IntermediateStatement;
@@ -77,20 +81,13 @@ public class ConnectionHandler extends Thread {
   private static final AtomicInteger incrementingConnectionId = new AtomicInteger(0);
   private ConnectionMetadata connectionMetadata;
   private WireMessage message;
-  private final Connection spannerConnection;
+  private Connection spannerConnection;
 
   ConnectionHandler(ProxyServer server, Socket socket) {
     super("ConnectionHandler-" + CONNECTION_HANDLER_ID_GENERATOR.incrementAndGet());
     this.server = server;
     this.socket = socket;
     this.secret = new SecureRandom().nextInt();
-    String uri = server.getOptions().getConnectionURL();
-    if (uri.startsWith("jdbc:")) {
-      uri = uri.substring("jdbc:".length());
-    }
-    uri = appendPropertiesToUrl(uri, server.getProperties());
-    ConnectionOptions connectionOptions = ConnectionOptions.newBuilder().setUri(uri).build();
-    this.spannerConnection = connectionOptions.getConnection();
     setDaemon(true);
     logger.log(
         Level.INFO,
@@ -100,6 +97,37 @@ public class ConnectionHandler extends Thread {
                 getName(), socket.getInetAddress().getHostAddress()));
   }
 
+  @InternalApi
+  public void connectToSpanner(String database) {
+    OptionsMetadata options = getServer().getOptions();
+    String uri =
+        options.hasDefaultConnectionUrl()
+            ? options.getDefaultConnectionUrl()
+            : options.buildConnectionURL(database);
+    if (uri.startsWith("jdbc:")) {
+      uri = uri.substring("jdbc:".length());
+    }
+    uri = appendPropertiesToUrl(uri, getServer().getProperties());
+    ConnectionOptions connectionOptions = ConnectionOptions.newBuilder().setUri(uri).build();
+    Connection spannerConnection = connectionOptions.getConnection();
+    try {
+      // Note: Calling getDialect() will cause a SpannerException if the connection itself is
+      // invalid, for example as a result of the credentials being wrong.
+      if (spannerConnection.getDialect() != Dialect.POSTGRESQL) {
+        throw SpannerExceptionFactory.newSpannerException(
+            ErrorCode.INVALID_ARGUMENT,
+            String.format(
+                "The database uses dialect %s. Currently PGAdapter only supports connections to PostgreSQL dialect databases. "
+                    + "These can be created using https://cloud.google.com/spanner/docs/quickstart-console#postgresql",
+                spannerConnection.getDialect()));
+      }
+    } catch (SpannerException e) {
+      spannerConnection.close();
+      throw e;
+    }
+    this.spannerConnection = spannerConnection;
+  }
+
   private String appendPropertiesToUrl(String url, Properties info) {
     if (info == null || info.isEmpty()) {
       return url;
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/ProxyServer.java b/src/main/java/com/google/cloud/spanner/pgadapter/ProxyServer.java
index 93ee8737..bac3ad58 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/ProxyServer.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/ProxyServer.java
@@ -15,10 +15,7 @@
 package com.google.cloud.spanner.pgadapter;
 
 import com.google.api.core.AbstractApiService;
-import com.google.cloud.spanner.Dialect;
-import com.google.cloud.spanner.ErrorCode;
 import com.google.cloud.spanner.SpannerException;
-import com.google.cloud.spanner.SpannerExceptionFactory;
 import com.google.cloud.spanner.pgadapter.ConnectionHandler.QueryMode;
 import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata;
 import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata.TextFormat;
@@ -221,23 +218,8 @@ public class ProxyServer extends AbstractApiService {
    */
   void createConnectionHandler(Socket socket) {
     ConnectionHandler handler = new ConnectionHandler(this, socket);
-    try {
-      // Note: Calling getDialect() will cause a SpannerException if the connection itself is
-      // invalid, for example as a result of the credentials being wrong.
-      if (handler.getSpannerConnection().getDialect() != Dialect.POSTGRESQL) {
-        throw SpannerExceptionFactory.newSpannerException(
-            ErrorCode.INVALID_ARGUMENT,
-            String.format(
-                "The database uses dialect %s. Currently PGAdapter only supports connections to PostgreSQL dialect databases. "
-                    + "These can be created using https://cloud.google.com/spanner/docs/quickstart-console#postgresql",
-                handler.getSpannerConnection().getDialect()));
-      }
-      register(handler);
-      handler.start();
-    } catch (Exception e) {
-      handler.getSpannerConnection().close();
-      throw e;
-    }
+    register(handler);
+    handler.start();
   }
 
   /**
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/metadata/OptionsMetadata.java b/src/main/java/com/google/cloud/spanner/pgadapter/metadata/OptionsMetadata.java
index df025cea..67e319d8 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/metadata/OptionsMetadata.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/metadata/OptionsMetadata.java
@@ -17,6 +17,7 @@ package com.google.cloud.spanner.pgadapter.metadata;
 import com.google.auth.oauth2.GoogleCredentials;
 import com.google.cloud.spanner.SpannerExceptionFactory;
 import com.google.cloud.spanner.pgadapter.Server;
+import com.google.common.base.Preconditions;
 import com.google.common.base.Strings;
 import java.io.IOException;
 import java.util.HashMap;
@@ -60,8 +61,9 @@ public class OptionsMetadata {
   private static final String OPTION_SERVER_VERSION = "v";
   private static final String OPTION_DEBUG_MODE = "debug";
 
+  private final CommandLine commandLine;
   private final CommandMetadataParser commandMetadataParser;
-  private final String connectionURL;
+  private final String defaultConnectionUrl;
   private final int proxyPort;
   private final TextFormat textFormat;
   private final boolean binaryFormat;
@@ -75,9 +77,14 @@ public class OptionsMetadata {
   private final boolean debugMode;
 
   public OptionsMetadata(String[] args) {
-    CommandLine commandLine = buildOptions(args);
+    this.commandLine = buildOptions(args);
     this.commandMetadataParser = new CommandMetadataParser();
-    this.connectionURL = buildConnectionURL(commandLine);
+    if (this.commandLine.hasOption(OPTION_DATABASE_NAME)) {
+      this.defaultConnectionUrl =
+          buildConnectionURL(this.commandLine.getOptionValue(OPTION_DATABASE_NAME));
+    } else {
+      this.defaultConnectionUrl = null;
+    }
     this.proxyPort = buildProxyPort(commandLine);
     this.textFormat = TextFormat.POSTGRESQL;
     this.binaryFormat = commandLine.hasOption(OPTION_BINARY_FORMAT);
@@ -94,7 +101,7 @@ public class OptionsMetadata {
   }
 
   public OptionsMetadata(
-      String connectionURL,
+      String defaultConnectionUrl,
       int proxyPort,
       TextFormat textFormat,
       boolean forceBinary,
@@ -102,8 +109,9 @@ public class OptionsMetadata {
       boolean requiresMatcher,
       boolean replaceJdbcMetadataQueries,
       JSONObject commandMetadata) {
+    this.commandLine = null;
     this.commandMetadataParser = new CommandMetadataParser();
-    this.connectionURL = connectionURL;
+    this.defaultConnectionUrl = defaultConnectionUrl;
     this.proxyPort = proxyPort;
     this.textFormat = textFormat;
     this.binaryFormat = forceBinary;
@@ -151,10 +159,9 @@ public class OptionsMetadata {
    * Get credential file path from either command line or application default. If neither throw
    * error.
    *
-   * @param commandLine The parsed options for CLI
    * @return The absolute path of the credentials file.
    */
-  private String buildCredentialsFile(CommandLine commandLine) {
+  private String buildCredentialsFile() {
     if (!commandLine.hasOption(OPTION_CREDENTIALS_FILE)) {
       try {
         // This will throw an IOException if no default credentials are available.
@@ -169,10 +176,10 @@ public class OptionsMetadata {
   /**
    * Takes user inputs and builds a JDBC connection string from them.
    *
-   * @param commandLine The parsed options for CLI
    * @return The parsed JDBC connection string.
    */
-  private String buildConnectionURL(CommandLine commandLine) {
+  public String buildConnectionURL(String database) {
+    Preconditions.checkNotNull(database);
     String host = commandLine.getOptionValue(OPTION_SPANNER_ENDPOINT, "");
     String jdbcEndpoint;
     if (host.isEmpty()) {
@@ -198,10 +205,10 @@ public class OptionsMetadata {
                 + ";userAgent=%s",
             commandLine.getOptionValue(OPTION_PROJECT_ID),
             commandLine.getOptionValue(OPTION_INSTANCE_ID),
-            commandLine.getOptionValue(OPTION_DATABASE_NAME),
+            database,
             DEFAULT_USER_AGENT);
 
-    String credentials = buildCredentialsFile(commandLine);
+    String credentials = buildCredentialsFile();
     if (!Strings.isNullOrEmpty(credentials)) {
       url = String.format("%s;credentials=%s", url, credentials);
     }
@@ -270,11 +277,15 @@ public class OptionsMetadata {
         "instance",
         true,
         "The id of the Spanner instance within the GCP project.");
-    options.addRequiredOption(
+    options.addOption(
         OPTION_DATABASE_NAME,
         "database",
         true,
-        "The name of the Spanner database within the GCP project.");
+        "The default Spanner database within the GCP project to use. "
+            + "If specified, PGAdapter will always connect to this database. "
+            + "Any database name in the connection request from the client will be ignored. "
+            + "Omit this option to be able to connect to different databases using a single "
+            + "PGAdapter instance.");
     options.addOption(
         OPTION_CREDENTIALS_FILE,
         "credentials-file",
@@ -386,8 +397,32 @@ public class OptionsMetadata {
     return this.commandMetadataJSON;
   }
 
+  /**
+   * @deprecated use {@link #getDefaultConnectionUrl()}
+   * @return the default connection URL that is used by the server.
+   */
+  @Deprecated
   public String getConnectionURL() {
-    return this.connectionURL;
+    return this.defaultConnectionUrl;
+  }
+
+  /**
+   * @return true if the server uses a default connection URL and ignores the database in a
+   *     connection request
+   */
+  public boolean hasDefaultConnectionUrl() {
+    return this.defaultConnectionUrl != null;
+  }
+
+  /**
+   * Returns the default connection URL that is used by the server. If a default connection URL has
+   * been set, the database parameter in a connection request will be ignored, and the database in
+   * this connection URL will be used instead.
+   *
+   * @return the default connection URL that is used by the server.
+   */
+  public String getDefaultConnectionUrl() {
+    return defaultConnectionUrl;
   }
 
   public int getProxyPort() {
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 e0f3bee3..12388ea3 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
@@ -33,7 +33,7 @@ public class StartupMessage extends BootstrapMessage {
   public static final int IDENTIFIER = 196608; // First Hextet: 3 (version), Second Hextet: 0
 
   private final boolean authenticate;
-  private Map<String, String> parameters;
+  private final Map<String, String> parameters;
 
   public StartupMessage(ConnectionHandler connection, int length) throws Exception {
     super(connection, length);
@@ -44,6 +44,7 @@ public class StartupMessage extends BootstrapMessage {
   @Override
   protected void sendPayload() throws Exception {
     if (!authenticate) {
+      this.connection.connectToSpanner(this.parameters.get("database"));
       sendStartupMessage(
           this.outputStream,
           this.connection.getConnectionId(),


Test patch:
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/AbstractMockServerTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/AbstractMockServerTest.java
index 0d9a4eda..84ee8470 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/AbstractMockServerTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/AbstractMockServerTest.java
@@ -251,11 +251,11 @@ public abstract class AbstractMockServerTest {
 
   @BeforeClass
   public static void startMockSpannerAndPgAdapterServers() throws Exception {
-    doStartMockSpannerAndPgAdapterServers(Collections.emptyList());
+    doStartMockSpannerAndPgAdapterServers("d", Collections.emptyList());
   }
 
   protected static void doStartMockSpannerAndPgAdapterServers(
-      Iterable<String> extraPGAdapterOptions) throws Exception {
+      String defaultDatabase, Iterable<String> extraPGAdapterOptions) throws Exception {
     mockSpanner = new MockSpannerServiceImpl();
     mockSpanner.setAbortProbability(0.0D); // We don't want any unpredictable aborted transactions.
     mockSpanner.putStatementResult(StatementResult.query(SELECT1, SELECT1_RESULTSET));
@@ -303,24 +303,21 @@ public abstract class AbstractMockServerTest {
             .start();
 
     ImmutableList.Builder<String> argsListBuilder =
-        ImmutableList.<String>builder()
-            .add(
-                "-p",
-                "p",
-                "-i",
-                "i",
-                "-d",
-                "d",
-                "-jdbc",
-                "-debug",
-                "-c",
-                "", // empty credentials file, as we are using a plain text connection.
-                "-s",
-                "0", // port 0 to let the OS pick an available port
-                "-e",
-                String.format("localhost:%d", spannerServer.getPort()),
-                "-r",
-                "usePlainText=true;");
+        ImmutableList.<String>builder().add("-p", "p", "-i", "i");
+    if (defaultDatabase != null) {
+      argsListBuilder.add("-d", defaultDatabase);
+    }
+    argsListBuilder.add(
+        "-jdbc",
+        "-debug",
+        "-c",
+        "", // empty credentials file, as we are using a plain text connection.
+        "-s",
+        "0", // port 0 to let the OS pick an available port
+        "-e",
+        String.format("localhost:%d", spannerServer.getPort()),
+        "-r",
+        "usePlainText=true;");
     argsListBuilder.addAll(extraPGAdapterOptions);
     String[] args = argsListBuilder.build().toArray(new String[0]);
     pgServer = new ProxyServer(new OptionsMetadata(args));
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/EmulatedPsqlMockServerTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/EmulatedPsqlMockServerTest.java
index a68bbfa1..c1ab88ef 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/EmulatedPsqlMockServerTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/EmulatedPsqlMockServerTest.java
@@ -20,9 +20,12 @@ import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult;
 import com.google.cloud.spanner.Statement;
 import com.google.common.collect.ImmutableList;
 import com.google.spanner.v1.ExecuteBatchDmlRequest;
+import com.google.spanner.v1.ExecuteSqlRequest;
+import com.google.spanner.v1.SessionName;
 import java.sql.Connection;
 import java.sql.DriverManager;
 import java.sql.SQLException;
+import java.util.List;
 import org.junit.BeforeClass;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -41,7 +44,8 @@ public class EmulatedPsqlMockServerTest extends AbstractMockServerTest {
 
   @BeforeClass
   public static void startMockSpannerAndPgAdapterServers() throws Exception {
-    doStartMockSpannerAndPgAdapterServers(ImmutableList.of("-q"));
+    // Start PGAdapter in psql mode without a default database.
+    doStartMockSpannerAndPgAdapterServers(null, ImmutableList.of("-q"));
 
     mockSpanner.putStatementResults(
         StatementResult.update(Statement.of(INSERT1), 1L),
@@ -53,16 +57,31 @@ public class EmulatedPsqlMockServerTest extends AbstractMockServerTest {
    * mode for queries and DML statements. This makes the JDBC driver behave in (much) the same way
    * as psql.
    */
-  private String createUrl() {
+  private String createUrl(String database) {
     return String.format(
-        "jdbc:postgresql://localhost:%d/?preferQueryMode=simple", pgServer.getLocalPort());
+        "jdbc:postgresql://localhost:%d/%s?preferQueryMode=simple",
+        pgServer.getLocalPort(), database);
   }
 
   @Test
-  public void testTwoInserts() throws SQLException {
-    String sql = "insert into foo values (1); insert into foo values (2);";
+  public void testConnectToDifferentDatabases() throws SQLException {
+    final ImmutableList<String> databases = ImmutableList.of("db1", "db2");
+    for (String database : databases) {
+      try (Connection connection = DriverManager.getConnection(createUrl(database))) {
+        connection.createStatement().execute(INSERT1);
+      }
+    }
 
-    try (Connection connection = DriverManager.getConnection(createUrl())) {
+    List<ExecuteSqlRequest> requests = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class);
+    assertEquals(databases.size(), requests.size());
+    for (int i = 0; i < requests.size(); i++) {
+      assertEquals(databases.get(i), SessionName.parse(requests.get(i).getSession()).getDatabase());
+    }
+  }
+
+  @Test
+  public void testTwoInserts() throws SQLException {
+    try (Connection connection = DriverManager.getConnection(createUrl("my-db"))) {
       connection.createStatement().execute(String.format("%s; %s", INSERT1, INSERT2));
     }
 
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/ITJdbcMetadataTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/ITJdbcMetadataTest.java
index 85384038..877f9ad5 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/ITJdbcMetadataTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/ITJdbcMetadataTest.java
@@ -67,7 +67,8 @@ public class ITJdbcMetadataTest implements IntegrationTest {
   public static void setup() {
     testEnv.setUp();
     database = testEnv.createDatabase(getDdlStatements());
-    testEnv.startPGAdapterServer(database.getId(), getAdditionalPGAdapterOptions());
+    testEnv.startPGAdapterServerWithDefaultDatabase(
+        database.getId(), getAdditionalPGAdapterOptions());
   }
 
   @AfterClass
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/ITJdbcTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/ITJdbcTest.java
index 64811f3a..e46dc598 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/ITJdbcTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/ITJdbcTest.java
@@ -71,7 +71,7 @@ public class ITJdbcTest implements IntegrationTest {
 
     testEnv.setUp();
     database = testEnv.createDatabase(PgAdapterTestEnv.DEFAULT_DATA_MODEL);
-    testEnv.startPGAdapterServer(database.getId(), Collections.emptyList());
+    testEnv.startPGAdapterServer(Collections.emptyList());
   }
 
   @AfterClass
@@ -111,8 +111,8 @@ public class ITJdbcTest implements IntegrationTest {
 
   private String getConnectionUrl() {
     return String.format(
-        "jdbc:postgresql://%s/?preferQueryMode=%s",
-        testEnv.getPGAdapterHostAndPort(), preferQueryMode);
+        "jdbc:postgresql://%s/%s?preferQueryMode=%s",
+        testEnv.getPGAdapterHostAndPort(), database.getId().getDatabase(), preferQueryMode);
   }
 
   @Test
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/ITParameterizedQueryTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/ITParameterizedQueryTest.java
index 8f270da7..a199560e 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/ITParameterizedQueryTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/ITParameterizedQueryTest.java
@@ -51,7 +51,7 @@ public final class ITParameterizedQueryTest implements IntegrationTest {
   public static void setup() {
     testEnv.setUp();
     database = testEnv.createDatabase(getDdlStatements());
-    testEnv.startPGAdapterServer(database.getId(), Collections.emptyList());
+    testEnv.startPGAdapterServerWithDefaultDatabase(database.getId(), Collections.emptyList());
   }
 
   @AfterClass
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/ITQueryTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/ITQueryTest.java
index c0247a72..cdf70613 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/ITQueryTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/ITQueryTest.java
@@ -46,7 +46,7 @@ public final class ITQueryTest implements IntegrationTest {
   public static void setup() {
     testEnv.setUp();
     database = testEnv.createDatabase(getDdlStatements());
-    testEnv.startPGAdapterServer(database.getId(), Collections.emptyList());
+    testEnv.startPGAdapterServerWithDefaultDatabase(database.getId(), Collections.emptyList());
   }
 
   @AfterClass
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/PgAdapterTestEnv.java b/src/test/java/com/google/cloud/spanner/pgadapter/PgAdapterTestEnv.java
index 8907f36e..ec5cdb49 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/PgAdapterTestEnv.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/PgAdapterTestEnv.java
@@ -149,21 +149,25 @@ public final class PgAdapterTestEnv {
     options = createSpannerOptions();
   }
 
-  public void startPGAdapterServer(
+  public void startPGAdapterServer(Iterable<String> additionalPGAdapterOptions) {
+    startPGAdapterServer(null, additionalPGAdapterOptions);
+  }
+
+  public void startPGAdapterServerWithDefaultDatabase(
       DatabaseId databaseId, Iterable<String> additionalPGAdapterOptions) {
+    startPGAdapterServer(databaseId.getDatabase(), additionalPGAdapterOptions);
+  }
+
+  private void startPGAdapterServer(
+      String databaseId, Iterable<String> additionalPGAdapterOptions) {
     if (PG_ADAPTER_ADDRESS == null) {
       String credentials = getCredentials();
       ImmutableList.Builder<String> argsListBuilder =
-          ImmutableList.<String>builder()
-              .add(
-                  "-p",
-                  getProjectId(),
-                  "-i",
-                  getInstanceId(),
-                  "-d",
-                  databaseId.getDatabase(),
-                  "-s",
-                  String.valueOf(0));
+          ImmutableList.<String>builder().add("-p", getProjectId(), "-i", getInstanceId());
+      if (databaseId != null) {
+        argsListBuilder.add("-d", databaseId);
+      }
+      argsListBuilder.add("-s", String.valueOf(0));
       if (getSpannerUrl() != null) {
         String host = getSpannerUrl();
         if (host.startsWith("https://")) {
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/ProtocolTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/ProtocolTest.java
index eb0f73b2..3581a362 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/ProtocolTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/ProtocolTest.java
@@ -1604,6 +1604,7 @@ public class ProtocolTest {
     message.send();
 
     DataInputStream outputResult = inputStreamFromOutputStream(result);
+    verify(connectionHandler).connectToSpanner("databasename");
 
     // AuthenticationOkResponse
     assertEquals('R', outputResult.readByte());
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/golang/ITPgxTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/golang/ITPgxTest.java
index fdb8a8ef..5618af9a 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/golang/ITPgxTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/golang/ITPgxTest.java
@@ -56,7 +56,7 @@ public class ITPgxTest implements IntegrationTest {
 
     testEnv.setUp();
     database = testEnv.createDatabase(getDdlStatements());
-    testEnv.startPGAdapterServer(database.getId(), Collections.emptyList());
+    testEnv.startPGAdapterServerWithDefaultDatabase(database.getId(), Collections.emptyList());
   }
 
   @AfterClass