From 7239425a44ddaccc24293631f095676025bb0b8d Mon Sep 17 00:00:00 2001
From: ajaaym <34161822+ajaaym@users.noreply.github.com>
Date: Tue, 18 Jun 2019 11:14:58 -0400
Subject: [PATCH] Add option to return raw input stream for response (#1323)

---
 .../services/AbstractGoogleClientRequest.java | 28 ++++++++++
 .../AbstractGoogleClientRequestTest.java      | 52 +++++++++++++++++++
 2 files changed, 80 insertions(+)

diff --git a/google-api-client/src/main/java/com/google/api/client/googleapis/services/AbstractGoogleClientRequest.java b/google-api-client/src/main/java/com/google/api/client/googleapis/services/AbstractGoogleClientRequest.java
index 1de14b99a..f212ccd9b 100644
--- a/google-api-client/src/main/java/com/google/api/client/googleapis/services/AbstractGoogleClientRequest.java
+++ b/google-api-client/src/main/java/com/google/api/client/googleapis/services/AbstractGoogleClientRequest.java
@@ -94,6 +94,9 @@ public abstract class AbstractGoogleClientRequest<T> extends GenericData {
   /** Whether to disable GZip compression of HTTP content. */
   private boolean disableGZipContent;
 
+  /** Whether to return raw input stream in {@link HttpResponse#getContent()}. */
+  private boolean returnRawInputStream;
+
   /** Response class to parse into. */
   private Class<T> responseClass;
 
@@ -222,6 +225,11 @@ public final boolean getDisableGZipContent() {
     return disableGZipContent;
   }
 
+  /** Returns whether response should return raw input stream. */
+  public final boolean getReturnRawInputSteam() {
+    return returnRawInputStream;
+  }
+
   /**
    * Sets whether to disable GZip compression of HTTP content.
    *
@@ -239,6 +247,25 @@ public AbstractGoogleClientRequest<T> setDisableGZipContent(boolean disableGZipC
     return this;
   }
 
+  /**
+   * Sets whether the response should return raw input stream or not.
+   *
+   * <p>
+   * By default it is {@code false}.
+   * </p>
+   *
+   * <p>
+   * When the response contains a known content-encoding header, the response stream is wrapped
+   * with an InputStream that decodes the content. This fails when we download large files in
+   * chunks (see <a href="https://github.com/googleapis/google-api-java-client/issues/1009">#1009
+   * </a>). Setting this to true will make the response return the raw input stream.
+   * </p>
+   */
+  public AbstractGoogleClientRequest<T> setReturnRawInputStream(boolean returnRawInputStream) {
+    this.returnRawInputStream = returnRawInputStream;
+    return this;
+  }
+
   /** Returns the HTTP method. */
   public final String getRequestMethod() {
     return requestMethod;
@@ -406,6 +433,7 @@ private HttpRequest buildHttpRequest(boolean usingHead) throws IOException {
     if (!disableGZipContent) {
       httpRequest.setEncoding(new GZipEncoding());
     }
+    httpRequest.setResponseReturnRawInputStream(returnRawInputStream);
     final HttpResponseInterceptor responseInterceptor = httpRequest.getResponseInterceptor();
     httpRequest.setResponseInterceptor(new HttpResponseInterceptor() {
 
diff --git a/google-api-client/src/test/java/com/google/api/client/googleapis/services/AbstractGoogleClientRequestTest.java b/google-api-client/src/test/java/com/google/api/client/googleapis/services/AbstractGoogleClientRequestTest.java
index 58c5320da..458e991d5 100644
--- a/google-api-client/src/test/java/com/google/api/client/googleapis/services/AbstractGoogleClientRequestTest.java
+++ b/google-api-client/src/test/java/com/google/api/client/googleapis/services/AbstractGoogleClientRequestTest.java
@@ -18,6 +18,7 @@
 import com.google.api.client.http.EmptyContent;
 import com.google.api.client.http.HttpMethods;
 import com.google.api.client.http.HttpRequest;
+import com.google.api.client.http.HttpRequestInitializer;
 import com.google.api.client.http.HttpResponseException;
 import com.google.api.client.http.HttpStatusCodes;
 import com.google.api.client.http.HttpTransport;
@@ -31,8 +32,12 @@
 import com.google.api.client.testing.http.MockLowLevelHttpRequest;
 import com.google.api.client.testing.http.MockLowLevelHttpResponse;
 import com.google.api.client.util.StringUtils;
+import com.google.common.io.BaseEncoding;
+import java.io.ByteArrayInputStream;
 import java.io.IOException;
+import java.io.InputStream;
 import java.util.Arrays;
+import java.util.zip.GZIPInputStream;
 import junit.framework.TestCase;
 
 /**
@@ -246,6 +251,53 @@ public void testSetsApiClientHeaderWithoutOsVersion() throws Exception {
     assertFalse("Api version should not contain the os version", version.matches(".*my-os.*"));
   }
 
+  public void testReturnRawInputStream_defaultFalse() throws Exception {
+    HttpTransport transport = new MockHttpTransport() {
+      @Override
+      public LowLevelHttpRequest buildRequest(final String method, final String url) {
+        return new MockLowLevelHttpRequest() {
+          @Override
+          public LowLevelHttpResponse execute() {
+            return new MockLowLevelHttpResponse().setContentEncoding("gzip").setContent(new ByteArrayInputStream(
+                BaseEncoding.base64()
+                    .decode("H4sIAAAAAAAAAPNIzcnJV3DPz0/PSVVwzskvTVEILskvSkxPVQQA/LySchsAAAA=")));
+          }
+        };
+      }
+    };
+    MockGoogleClient client = new MockGoogleClient.Builder(transport, ROOT_URL, SERVICE_PATH,
+        JSON_OBJECT_PARSER, null).setApplicationName(
+        "Test Application").build();
+    MockGoogleClientRequest<String> request = new MockGoogleClientRequest<String>(
+        client, HttpMethods.GET, URI_TEMPLATE, null, String.class);
+    InputStream inputStream = request.executeAsInputStream();
+    assertTrue(inputStream instanceof GZIPInputStream);
+  }
+
+  public void testReturnRawInputStream_True() throws Exception {
+    HttpTransport transport = new MockHttpTransport() {
+      @Override
+      public LowLevelHttpRequest buildRequest(final String method, final String url) {
+        return new MockLowLevelHttpRequest() {
+          @Override
+          public LowLevelHttpResponse execute() {
+            return new MockLowLevelHttpResponse().setContentEncoding("gzip").setContent(new ByteArrayInputStream(
+                BaseEncoding.base64()
+                    .decode("H4sIAAAAAAAAAPNIzcnJV3DPz0/PSVVwzskvTVEILskvSkxPVQQA/LySchsAAAA=")));
+          }
+        };
+      }
+    };
+    MockGoogleClient client = new MockGoogleClient.Builder(transport, ROOT_URL, SERVICE_PATH,
+        JSON_OBJECT_PARSER, null).setApplicationName(
+        "Test Application").build();
+    MockGoogleClientRequest<String> request = new MockGoogleClientRequest<String>(
+        client, HttpMethods.GET, URI_TEMPLATE, null, String.class);
+    request.setReturnRawInputStream(true);
+    InputStream inputStream = request.executeAsInputStream();
+    assertFalse(inputStream instanceof GZIPInputStream);
+  }
+
   private class AssertHeaderTransport extends MockHttpTransport {
     String expectedHeader;
     String expectedHeaderValue;