feat: support setting compression option

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

@@ -49,6 +49,8 @@
 import com.google.spanner.admin.database.v1.RestoreDatabaseRequest;
 import com.google.spanner.v1.ExecuteSqlRequest.QueryOptions;
 import io.grpc.CallCredentials;
+import io.grpc.CompressorRegistry;
+import io.grpc.ExperimentalApi;
 import io.grpc.ManagedChannelBuilder;
 import java.io.IOException;
 import java.net.MalformedURLException;
@@ -58,6 +60,7 @@
 import java.util.Map.Entry;
 import java.util.Set;
 import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
 import org.threeten.bp.Duration;
 
 /** Options for the Cloud Spanner service. */
@@ -104,6 +107,7 @@ public class SpannerOptions extends ServiceOptions<Spanner, SpannerOptions> {
   private final Map<DatabaseId, QueryOptions> mergedQueryOptions;
 
   private final CallCredentialsProvider callCredentialsProvider;
+  private final String compressorName;
 
   /**
    * Interface that can be used to provide {@link CallCredentials} instead of {@link Credentials} to
@@ -174,6 +178,7 @@ private SpannerOptions(Builder builder) {
       this.mergedQueryOptions = ImmutableMap.copyOf(merged);
     }
     callCredentialsProvider = builder.callCredentialsProvider;
+    compressorName = builder.compressorName;
   }
 
   /**
@@ -238,6 +243,7 @@ public static class Builder
     private boolean autoThrottleAdministrativeRequests = false;
     private Map<DatabaseId, QueryOptions> defaultQueryOptions = new HashMap<>();
     private CallCredentialsProvider callCredentialsProvider;
+    private String compressorName;
     private String emulatorHost = System.getenv("SPANNER_EMULATOR_HOST");
 
     private Builder() {
@@ -309,6 +315,7 @@ private Builder() {
       this.autoThrottleAdministrativeRequests = options.autoThrottleAdministrativeRequests;
       this.defaultQueryOptions = options.defaultQueryOptions;
       this.callCredentialsProvider = options.callCredentialsProvider;
+      this.compressorName = options.compressorName;
       this.channelProvider = options.channelProvider;
       this.channelConfigurator = options.channelConfigurator;
       this.interceptorProvider = options.interceptorProvider;
@@ -558,6 +565,28 @@ public Builder setCallCredentialsProvider(CallCredentialsProvider callCredential
       return this;
     }
 
+    /**
+     * Sets the compression to use for all gRPC calls. The compressor must be a valid name known in
+     * the {@link CompressorRegistry}.
+     *
+     * <p>Supported values are:
+     *
+     * <ul>
+     *   <li>gzip: Enable gzip compression
+     *   <li>identity: Disable compression
+     *   <li><code>null</code>: Use default comporession
+     *       <ul>
+     */
+    @ExperimentalApi("https://github.com/grpc/grpc-java/issues/1704")
+    public Builder setCompressorName(@Nullable String compressorName) {
+      Preconditions.checkArgument(
+          compressorName == null
+              || CompressorRegistry.getDefaultInstance().lookupCompressor(compressorName) != null,
+          String.format("%s is not a known compressor", compressorName));
+      this.compressorName = compressorName;
+      return this;
+    }
+
     /**
      * Specifying this will allow the client to prefetch up to {@code prefetchChunks} {@code
      * PartialResultSet} chunks for each read and query. The data size of each chunk depends on the
@@ -690,6 +719,10 @@ public CallCredentialsProvider getCallCredentialsProvider() {
     return callCredentialsProvider;
   }
 
+  public String getCompressorName() {
+    return compressorName;
+  }
+
   /** Returns the default query options to use for the specific database. */
   public QueryOptions getDefaultQueryOptions(DatabaseId databaseId) {
     // Use the specific query options for the database if any have been specified. These have

google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java

@@ -226,6 +226,7 @@ private void awaitTermination() throws InterruptedException {
   private final String projectName;
   private final SpannerMetadataProvider metadataProvider;
   private final CallCredentialsProvider callCredentialsProvider;
+  private final String compressorName;
   private final Duration waitTimeout =
       systemProperty(PROPERTY_TIMEOUT_SECONDS, DEFAULT_TIMEOUT_SECONDS);
   private final Duration idleTimeout =
@@ -279,6 +280,7 @@ public GapicSpannerRpc(final SpannerOptions options) {
             mergedHeaderProvider.getHeaders(),
             internalHeaderProviderBuilder.getResourceHeaderKey());
     this.callCredentialsProvider = options.getCallCredentialsProvider();
+    this.compressorName = options.getCompressorName();
 
     // Create a managed executor provider.
     this.executorProvider =
@@ -1239,6 +1241,9 @@ GrpcCallContext newCallContext(@Nullable Map<Option, ?> options, String resource
             context.withCallOptions(context.getCallOptions().withCallCredentials(callCredentials));
       }
     }
+    if (compressorName != null) {
+      context = context.withCallOptions(context.getCallOptions().withCompression(compressorName));
+    }
     return context.withStreamWaitTimeout(waitTimeout).withStreamIdleTimeout(idleTimeout);
   }
 

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

@@ -505,4 +505,21 @@ public String getOptimizerVersion() {
     assertThat(options.getDefaultQueryOptions(DatabaseId.of("p", "i", "o")))
         .isEqualTo(QueryOptions.newBuilder().setOptimizerVersion("2").build());
   }
+
+  @Test
+  public void testCompressorName() {
+    assertThat(SpannerOptions.newBuilder().setCompressorName("gzip").build().getCompressorName())
+        .isEqualTo("gzip");
+    assertThat(
+            SpannerOptions.newBuilder().setCompressorName("identity").build().getCompressorName())
+        .isEqualTo("identity");
+    assertThat(SpannerOptions.newBuilder().setCompressorName(null).build().getCompressorName())
+        .isNull();
+    try {
+      SpannerOptions.newBuilder().setCompressorName("foo");
+      fail("missing expected exception");
+    } catch (IllegalArgumentException e) {
+      // ignore, this is the expected exception.
+    }
+  }
 }

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

@@ -168,6 +168,27 @@ public void testMultipleSpannersFromSameSpannerOptions() throws InterruptedExcep
     assertThat(getNumberOfThreadsWithName(SPANNER_THREAD_NAME)).isAtMost(baseThreadCount);
   }
 
+  @Test
+  public void testCompression() {
+    for (String compressorName : new String[] {"gzip", "identity", null}) {
+      SpannerOptions options =
+          env.getTestHelper().getOptions().toBuilder().setCompressorName(compressorName).build();
+      try (Spanner spanner = options.getService()) {
+        DatabaseClient client = spanner.getDatabaseClient(db.getId());
+        try (ResultSet rs =
+            client
+                .singleUse()
+                .executeQuery(Statement.of("SELECT 1 AS COL1 UNION ALL SELECT 2 AS COL2"))) {
+          assertThat(rs.next()).isTrue();
+          assertThat(rs.getLong(0)).isEqualTo(1L);
+          assertThat(rs.next()).isTrue();
+          assertThat(rs.getLong(0)).isEqualTo(2L);
+          assertThat(rs.next()).isFalse();
+        }
+      }
+    }
+  }
+
   private void waitForStartup() throws InterruptedException {
     // Wait until the IT environment has already started all base worker threads.
     int threadCount;