chore: add getter for database role in DatabaseClient and BatchClient (#2029)

* chore: add getter for database role in DatabaseClient and BatchClient for beam unit tests

* fix clirr

* 🦉 Updates from OwlBot post-processor

See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md

* incorporate requested changes

* incorporate requested changes

* incorporate requested changes

* 🦉 Updates from OwlBot post-processor

See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md

* Update google-cloud-spanner/src/test/java/com/google/cloud/spanner/BatchClientImplTest.java

* Update BatchClientImplTest.java

* Add doc for getDatabaseRole in DatabaseClient

* Update DatabaseClient.java

Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>
Co-authored-by: Rajat Bhatta <[email protected]>

Author: 3 people

google-cloud-spanner/clirr-ignored-differences.xml

@@ -152,4 +152,14 @@
     <className>com/google/cloud/spanner/StructReader</className>
     <method>java.util.List getPgJsonbList(java.lang.String)</method>
   </difference>
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/cloud/spanner/BatchClient</className>
+    <method>java.lang.String getDatabaseRole()</method>
+  </difference>
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/cloud/spanner/DatabaseClient</className>
+    <method>java.lang.String getDatabaseRole()</method>
+  </difference>
 </differences>

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

@@ -59,4 +59,13 @@ public interface BatchClient {
    *     successive read/query.
    */
   BatchReadOnlyTransaction batchReadOnlyTransaction(BatchTransactionId batchTransactionId);
+
+  /**
+   * Returns the {@link DatabaseRole} used by the client connection. The database role that is used
+   * determines the access permissions that a connection has. This can for example be used to create
+   * connections that are only permitted to access certain tables.
+   */
+  default String getDatabaseRole() {
+    throw new UnsupportedOperationException("method should be overwritten");
+  };
 }

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

@@ -33,6 +33,7 @@
 import io.opencensus.trace.Tracing;
 import java.util.List;
 import java.util.Map;
+import javax.annotation.Nullable;
 
 /** Default implementation for Batch Client interface. */
 public class BatchClientImpl implements BatchClient {
@@ -42,6 +43,12 @@ public class BatchClientImpl implements BatchClient {
     this.sessionClient = checkNotNull(sessionClient);
   }
 
+  @Override
+  @Nullable
+  public String getDatabaseRole() {
+    return this.sessionClient.getSpanner().getOptions().getDatabaseRole();
+  }
+
   @Override
   public BatchReadOnlyTransaction batchReadOnlyTransaction(TimestampBound bound) {
     SessionImpl session = sessionClient.createSession();

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

@@ -36,6 +36,17 @@ default Dialect getDialect() {
     throw new UnsupportedOperationException("method should be overwritten");
   }
 
+  /**
+   * Returns the {@link DatabaseRole} used by the client connection. The database role that is used
+   * determines the access permissions that a connection has. This can for example be used to create
+   * connections that are only permitted to access certain tables.
+   *
+   * @return the {@link DatabaseRole} used by the client connection.
+   */
+  default String getDatabaseRole() {
+    throw new UnsupportedOperationException("method should be overwritten");
+  }
+
   /**
    * Writes the given mutations atomically to the database.
    *

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

@@ -28,6 +28,7 @@
 import io.opencensus.trace.Span;
 import io.opencensus.trace.Tracer;
 import io.opencensus.trace.Tracing;
+import javax.annotation.Nullable;
 
 class DatabaseClientImpl implements DatabaseClient {
   private static final String READ_WRITE_TRANSACTION = "CloudSpanner.ReadWriteTransaction";
@@ -58,6 +59,12 @@ public Dialect getDialect() {
     return pool.getDialect();
   }
 
+  @Override
+  @Nullable
+  public String getDatabaseRole() {
+    return pool.getDatabaseRole();
+  }
+
   @Override
   public Timestamp write(final Iterable<Mutation> mutations) throws SpannerException {
     return writeWithOptions(mutations).getCommitTimestamp();

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

@@ -1784,6 +1784,7 @@ private enum Position {
 
   private final SessionPoolOptions options;
   private final SettableFuture<Dialect> dialect = SettableFuture.create();
+  private final String databaseRole;
   private final SessionClient sessionClient;
   private final ScheduledExecutorService executor;
   private final ExecutorFactory<ScheduledExecutorService> executorFactory;
@@ -1858,6 +1859,7 @@ static SessionPool createPool(
       SpannerOptions spannerOptions, SessionClient sessionClient, List<LabelValue> labelValues) {
     return createPool(
         spannerOptions.getSessionPoolOptions(),
+        spannerOptions.getDatabaseRole(),
         ((GrpcTransportOptions) spannerOptions.getTransportOptions()).getExecutorFactory(),
         sessionClient,
         new Clock(),
@@ -1879,6 +1881,7 @@ static SessionPool createPool(
       Clock clock) {
     return createPool(
         poolOptions,
+        null,
         executorFactory,
         sessionClient,
         clock,
@@ -1888,6 +1891,7 @@ static SessionPool createPool(
 
   static SessionPool createPool(
       SessionPoolOptions poolOptions,
+      String databaseRole,
       ExecutorFactory<ScheduledExecutorService> executorFactory,
       SessionClient sessionClient,
       Clock clock,
@@ -1896,6 +1900,7 @@ static SessionPool createPool(
     SessionPool pool =
         new SessionPool(
             poolOptions,
+            databaseRole,
             executorFactory,
             executorFactory.get(),
             sessionClient,
@@ -1908,13 +1913,15 @@ static SessionPool createPool(
 
   private SessionPool(
       SessionPoolOptions options,
+      String databaseRole,
       ExecutorFactory<ScheduledExecutorService> executorFactory,
       ScheduledExecutorService executor,
       SessionClient sessionClient,
       Clock clock,
       MetricRegistry metricRegistry,
       List<LabelValue> labelValues) {
     this.options = options;
+    this.databaseRole = databaseRole;
     this.executorFactory = executorFactory;
     this.executor = executor;
     this.sessionClient = sessionClient;
@@ -1956,6 +1963,11 @@ Dialect getDialect() {
     }
   }
 
+  @Nullable
+  public String getDatabaseRole() {
+    return databaseRole;
+  }
+
   @VisibleForTesting
   int getNumberOfSessionsInUse() {
     synchronized (lock) {

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

@@ -17,6 +17,7 @@
 package com.google.cloud.spanner;
 
 import static com.google.common.truth.Truth.assertThat;
+import static org.junit.Assert.assertEquals;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.Mockito.anyMap;
 import static org.mockito.Mockito.eq;
@@ -115,4 +116,9 @@ public void testBatchReadOnlyTxnWithTxnId() {
     assertThat(batchTxn.getReadTimestamp())
         .isEqualTo(batchTxn.getBatchTransactionId().getTimestamp());
   }
+
+  @Test
+  public void testGetDatabaseRole() {
+    assertEquals(client.getDatabaseRole(), "role");
+  }
 }

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

@@ -94,6 +94,7 @@ public class DatabaseClientImplTest {
   private static final String TEST_PROJECT = "my-project";
   private static final String TEST_INSTANCE = "my-instance";
   private static final String TEST_DATABASE = "my-database";
+  private static final String TEST_DATABASE_ROLE = "my-role";
   private static final String INSTANCE_NAME =
       String.format("projects/%s/instances/%s", TEST_PROJECT, TEST_INSTANCE);
   private static final String DATABASE_NAME =
@@ -149,6 +150,7 @@ public void setUp() {
     spanner =
         SpannerOptions.newBuilder()
             .setProjectId(TEST_PROJECT)
+            .setDatabaseRole(TEST_DATABASE_ROLE)
             .setChannelProvider(channelProvider)
             .setCredentials(NoCredentials.getInstance())
             .setSessionPoolOption(SessionPoolOptions.newBuilder().setFailOnSessionLeak().build())
@@ -2318,4 +2320,11 @@ public void testUntypedNullParameters() {
     assertNotNull(updateCount);
     assertEquals(1L, updateCount.longValue());
   }
+
+  @Test
+  public void testGetDatabaseRole() {
+    DatabaseClient client =
+        spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE));
+    assertEquals(TEST_DATABASE_ROLE, client.getDatabaseRole());
+  }
 }

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

@@ -20,6 +20,7 @@
 import static com.google.cloud.spanner.MetricRegistryConstants.NUM_READ_SESSIONS;
 import static com.google.cloud.spanner.MetricRegistryConstants.NUM_SESSIONS_BEING_PREPARED;
 import static com.google.cloud.spanner.MetricRegistryConstants.NUM_WRITE_SESSIONS;
+import static com.google.cloud.spanner.MetricRegistryConstants.SPANNER_DEFAULT_LABEL_VALUES;
 import static com.google.cloud.spanner.MetricRegistryConstants.SPANNER_LABEL_KEYS;
 import static com.google.cloud.spanner.MetricRegistryConstants.SPANNER_LABEL_KEYS_WITH_TYPE;
 import static com.google.common.truth.Truth.assertThat;
@@ -108,6 +109,7 @@ public class SessionPoolTest extends BaseSessionPoolTest {
   SessionPool pool;
   SessionPoolOptions options;
   private String sessionName = String.format("%s/sessions/s", db.getName());
+  private String TEST_DATABASE_ROLE = "my-role";
 
   @Parameters(name = "min sessions = {0}")
   public static Collection<Object[]> data() {
@@ -127,6 +129,7 @@ private SessionPool createPool(
       Clock clock, MetricRegistry metricRegistry, List<LabelValue> labelValues) {
     return SessionPool.createPool(
         options,
+        TEST_DATABASE_ROLE,
         new TestExecutorFactory(),
         client.getSessionClient(db),
         clock,
@@ -1178,6 +1181,13 @@ public void testSessionMetrics() throws Exception {
     assertThat(writePreparedSessions.value()).isEqualTo(0L);
   }
 
+  @Test
+  public void testGetDatabaseRole() throws Exception {
+    setupMockSessionCreation();
+    pool = createPool(new FakeClock(), new FakeMetricRegistry(), SPANNER_DEFAULT_LABEL_VALUES);
+    assertEquals(TEST_DATABASE_ROLE, pool.getDatabaseRole());
+  }
+
   private void mockKeepAlive(Session session) {
     ReadContext context = mock(ReadContext.class);
     ResultSet resultSet = mock(ResultSet.class);

google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITDatabaseRolePermissionTest.java

@@ -110,6 +110,7 @@ public void grantAndRevokeDatabaseRolePermissions() throws Exception {
     try (ResultSet rs =
         dbClient.singleUse().executeQuery(Statement.of("SELECT COUNT(*) as cnt FROM T"))) {
       assertTrue(rs.next());
+      assertEquals(dbClient.getDatabaseRole(), dbRoleParent);
     } catch (PermissionDeniedException e) {
       // This is not expected
       fail("Got PermissionDeniedException when it should not have occurred.");