feat: support option for setting client-id

Author: aakashanandg

google-cloud-spanner/src/main/java/com/google/cloud/spanner/Spanner.java

@@ -125,6 +125,29 @@ public interface Spanner extends Service<SpannerOptions>, AutoCloseable {
    */
   DatabaseClient getDatabaseClient(DatabaseId db);
 
+
+  /**
+   * Returns a {@code DatabaseClient} for the given database and given client id. It uses a pool of sessions to talk to
+   * the database.
+   * <!--SNIPPET get_db_client-->
+   *
+   * <pre>{@code
+   * SpannerOptions options = SpannerOptions.newBuilder().build();
+   * Spanner spanner = options.getService();
+   * final String project = "test-project";
+   * final String instance = "test-instance";
+   * final String database = "example-db";
+   * final String client_id = "client_id"
+   * DatabaseId db =
+   *     DatabaseId.of(project, instance, database);
+   *
+   * DatabaseClient dbClient = spanner.getDatabaseClient(db, client_id);
+   * }</pre>
+   *
+   * <!--SNIPPET get_db_client-->
+   */
+  DatabaseClient getDatabaseClient(DatabaseId db, String clientId);
+
   /**
    * Returns a {@code BatchClient} to do batch operations on Cloud Spanner databases. Batch client
    * is useful when one wants to read/query a large amount of data from Cloud Spanner across

google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java

@@ -254,9 +254,12 @@ public InstanceAdminClient getInstanceAdminClient() {
 
   @Override
   public DatabaseClient getDatabaseClient(DatabaseId db) {
+    return getDatabaseClient(db, null);
+  }
+
+  public DatabaseClient getDatabaseClient(DatabaseId db, String clientId) {
     synchronized (this) {
       checkClosed();
-      String clientId = null;
       if (dbClients.containsKey(db) && !dbClients.get(db).isValid()) {
         // Close the invalidated client and remove it.
         dbClients.get(db).closeAsync(new ClosedException());

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionImpl.java

@@ -108,6 +108,7 @@
 import java.util.Iterator;
 import java.util.LinkedList;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 import java.util.Stack;
 import java.util.UUID;
@@ -157,6 +158,7 @@ class ConnectionImpl implements Connection {
   private static final ParsedStatement RELEASE_STATEMENT =
       AbstractStatementParser.getInstance(Dialect.GOOGLE_STANDARD_SQL)
           .parse(Statement.of("RELEASE s1"));
+  private static final String CLIENT_ID = "client_id";
 
   /**
    * Exception that is used to register the stacktrace of the code that opened a {@link Connection}.
@@ -251,8 +253,8 @@ static UnitOfWorkType of(TransactionMode transactionMode) {
     }
   }
 
-  private StatementExecutor.StatementTimeout statementTimeout =
-      new StatementExecutor.StatementTimeout();
+  private StatementTimeout statementTimeout =
+      new StatementTimeout();
   private boolean closed = false;
 
   private final Spanner spanner;
@@ -323,7 +325,28 @@ static UnitOfWorkType of(TransactionMode transactionMode) {
       EmulatorUtil.maybeCreateInstanceAndDatabase(
           spanner, options.getDatabaseId(), options.getDialect());
     }
-    this.dbClient = spanner.getDatabaseClient(options.getDatabaseId());
+    DatabaseClient tempDbClient = null;
+    try {
+      final Map<String, ConnectionPropertyValue<?>> propertyValueMap = options.getInitialConnectionPropertyValues();
+      String clientIdString = null;
+      if (propertyValueMap != null) {
+        final ConnectionPropertyValue<?> clientIdProp = propertyValueMap.get(CLIENT_ID);
+        if (clientIdProp != null) {
+          Object value = clientIdProp.getValue();
+          if (value != null) {
+            clientIdString = value.toString();
+          }
+        }
+        if (clientIdString != null && !clientIdString.isEmpty()) {
+          tempDbClient = spanner.getDatabaseClient(options.getDatabaseId(), clientIdString);
+        }
+      } else {
+        tempDbClient = spanner.getDatabaseClient(options.getDatabaseId());
+      }
+    } catch(Exception e) {
+      tempDbClient = spanner.getDatabaseClient(options.getDatabaseId());
+    }
+    this.dbClient = tempDbClient;
     this.batchClient = spanner.getBatchClient(options.getDatabaseId());
     this.ddlClient = createDdlClient();
     this.connectionState =
@@ -411,7 +434,7 @@ static Attributes createOpenTelemetryAttributes(DatabaseId databaseId) {
   }
 
   @VisibleForTesting
-  ConnectionState.Type getConnectionStateType() {
+  Type getConnectionStateType() {
     return this.connectionState.getType();
   }
 
@@ -500,7 +523,7 @@ private void reset(Context context, boolean inTransaction) {
 
     this.connectionState.resetValue(AUTOCOMMIT_DML_MODE, context, inTransaction);
     this.statementTag = null;
-    this.statementTimeout = new StatementExecutor.StatementTimeout();
+    this.statementTimeout = new StatementTimeout();
     this.connectionState.resetValue(DIRECTED_READ, context, inTransaction);
     this.connectionState.resetValue(SAVEPOINT_SUPPORT, context, inTransaction);
     this.protoDescriptors = null;
@@ -542,7 +565,7 @@ public boolean isClosed() {
   }
 
   private <T> T getConnectionPropertyValue(
-      com.google.cloud.spanner.connection.ConnectionProperty<T> property) {
+      ConnectionProperty<T> property) {
     return this.connectionState.getValue(property).getValue();
   }
 
@@ -563,8 +586,8 @@ private <T> void setConnectionPropertyValue(
    * Sets a connection property value only for the duration of the current transaction. The effects
    * of this will be undone once the transaction ends, regardless whether the transaction is
    * committed or rolled back. 'Local' properties are supported for both {@link
-   * com.google.cloud.spanner.connection.ConnectionState.Type#TRANSACTIONAL} and {@link
-   * com.google.cloud.spanner.connection.ConnectionState.Type#NON_TRANSACTIONAL} connection states.
+   * Type#TRANSACTIONAL} and {@link
+   * Type#NON_TRANSACTIONAL} connection states.
    *
    * <p>NOTE: This feature is not yet exposed in the public API.
    */

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionOptions.java

@@ -46,6 +46,7 @@
 import static com.google.cloud.spanner.connection.ConnectionProperties.RETURN_COMMIT_STATS;
 import static com.google.cloud.spanner.connection.ConnectionProperties.ROUTE_TO_LEADER;
 import static com.google.cloud.spanner.connection.ConnectionProperties.TRACING_PREFIX;
+import static com.google.cloud.spanner.connection.ConnectionProperties.CLIENT_ID;
 import static com.google.cloud.spanner.connection.ConnectionProperties.TRACK_CONNECTION_LEAKS;
 import static com.google.cloud.spanner.connection.ConnectionProperties.TRACK_SESSION_LEAKS;
 import static com.google.cloud.spanner.connection.ConnectionProperties.USER_AGENT;
@@ -539,6 +540,11 @@ public Builder setTracingPrefix(String tracingPrefix) {
       return this;
     }
 
+    public Builder setClientId(String clientId) {
+      setConnectionPropertyValue(CLIENT_ID, clientId);
+      return this;
+    }
+
     /** @return the {@link ConnectionOptions} */
     public ConnectionOptions build() {
       Preconditions.checkState(this.uri != null, "Connection URI is required");
@@ -603,7 +609,6 @@ private ConnectionOptions(Builder builder) {
 
     // Create the initial connection state from the parsed properties in the connection URL.
     this.initialConnectionState = new ConnectionState(connectionPropertyValues);
-
     // Check that at most one of credentials location, encoded credentials, credentials provider and
     // OUAuth token has been specified in the connection URI.
     Preconditions.checkArgument(

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionProperties.java

@@ -135,6 +135,14 @@ public class ConnectionProperties {
 
   private static final Boolean[] BOOLEANS = new Boolean[] {Boolean.TRUE, Boolean.FALSE};
 
+  static final ConnectionProperty<String> CLIENT_ID =
+      create(
+          "client_id",
+          "Client Id to use for this connection. Can only be set at the start up time",
+          null,
+          StringValueConverter.INSTANCE,
+          Context.STARTUP);
+
   static final ConnectionProperty<ConnectionState.Type> CONNECTION_STATE_TYPE =
       create(
           "connection_state_type",

google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerImplTest.java

@@ -358,6 +358,40 @@ public void testCreateInstanceAdminClient_whenMockAdminSettings_assertException(
     assertNotNull(instanceAdminClient);
   }
 
+  @Test
+  public void testGetDatabaseClient_when_clientId_is_not_null() {
+    String dbName =
+        String.format("projects/p1/instances/i1/databases/%s", UUID.randomUUID().toString());
+    DatabaseId db = DatabaseId.of(dbName);
+
+    Mockito.when(spannerOptions.getTransportOptions())
+        .thenReturn(GrpcTransportOptions.newBuilder().build());
+    Mockito.when(spannerOptions.getSessionPoolOptions())
+        .thenReturn(SessionPoolOptions.newBuilder().setMinSessions(0).build());
+    Mockito.when(spannerOptions.getDatabaseRole()).thenReturn("role");
+
+    DatabaseClientImpl databaseClient = (DatabaseClientImpl) impl.getDatabaseClient(db, "clientId-1");
+    assertThat(databaseClient.clientId).isEqualTo("clientId-1");
+
+    // Get same db client again.
+    DatabaseClientImpl databaseClient1 = (DatabaseClientImpl) impl.getDatabaseClient(db, "clientId-1");
+    assertThat(databaseClient1.clientId).isEqualTo(databaseClient.clientId);
+
+    // Get a db client for a different database.
+    String dbName2 =
+        String.format("projects/p1/instances/i1/databases/%s", UUID.randomUUID().toString());
+    DatabaseId db2 = DatabaseId.of(dbName2);
+    DatabaseClientImpl databaseClient2 = (DatabaseClientImpl) impl.getDatabaseClient(db2, "clientId-1");
+    assertThat(databaseClient2.clientId).isEqualTo("clientId-1");
+
+    // Getting a new database client for an invalidated database should use the same client id.
+    databaseClient.pool.setResourceNotFoundException(
+        new DatabaseNotFoundException(DoNotConstructDirectly.ALLOWED, "not found", null, null));
+    DatabaseClientImpl revalidated = (DatabaseClientImpl) impl.getDatabaseClient(db, "clientId-1");
+    assertThat(revalidated).isNotSameInstanceAs(databaseClient);
+    assertThat(revalidated.clientId).isEqualTo(databaseClient.clientId);
+  }
+
   private void closeSpannerAndIncludeStacktrace(Spanner spanner) {
     spanner.close();
   }