Repository: googleapis/google-api-python-client

Issue:
Cannot set headers when downloading media
Downloading media with the Google API Python Client is frequently done with the googleapiclient.http.MediaIoBaseDownload class. That class does not seem to respect headers set on HttpRequest objects. 

Example:

request = service.objects().get_media(bucket=bucket, object=filename)
request.headers["Cache-Control"] = "no-store"
downloader = http.MediaIoBaseDownload(out_file, request)

What should happen:

Outgoing HTTP requests to download the file should include the header "Cache-Control: no-store".

What does happen:

No Cache-Control header is respected.

One feature of Google Cloud Storage, customer-supplied encryption keys, uses headers to pass those encryption keys around, and so this bug makes downloading encrypted objects difficult.



Discussion:
/cc @jonparrott 

Looks like the same problem is with encryption headers

def setKeyHeaders(req, aes256KeyHash, aes256Key):
    req.headers['x-goog-encryption-algorithm'] = 'AES256'
    req.headers['x-goog-encryption-key'] = aes256Key
    req.headers['x-goog-encryption-key-sha256'] = aes256KeyHash


Resolving patch:
diff --git a/googleapiclient/http.py b/googleapiclient/http.py
--- a/googleapiclient/http.py
+++ b/googleapiclient/http.py
@@ -647,6 +647,14 @@ def __init__(self, fd, request, chunksize=DEFAULT_CHUNK_SIZE):
     self._sleep = time.sleep
     self._rand = random.random
 
+    self._headers = {}
+    for k, v in six.iteritems(request.headers):
+      # allow users to supply custom headers by setting them on the request
+      # but strip out the ones that are set by default on requests generated by
+      # API methods like Drive's files().get(fileId=...)
+      if not k.lower() in ('accept', 'accept-encoding', 'user-agent'):
+        self._headers[k] = v
+
   @util.positional(1)
   def next_chunk(self, num_retries=0):
     """Get the next chunk of the download.
@@ -666,10 +674,9 @@ def next_chunk(self, num_retries=0):
       googleapiclient.errors.HttpError if the response was not a 2xx.
       httplib2.HttpLib2Error if a transport error has occured.
     """
-    headers = {
-        'range': 'bytes=%d-%d' % (
+    headers = self._headers.copy()
+    headers['range'] = 'bytes=%d-%d' % (
             self._progress, self._progress + self._chunksize)
-        }
     http = self._request.http
 
     resp, content = _retry_request(


Test patch:
diff --git a/tests/test_http.py b/tests/test_http.py
--- a/tests/test_http.py
+++ b/tests/test_http.py
@@ -29,6 +29,7 @@
 from six.moves.urllib.parse import urlencode
 
 # Do not remove the httplib2 import
+import json
 import httplib2
 import logging
 import mock
@@ -456,6 +457,41 @@ def test_media_io_base_download(self):
     self.assertEqual(5, download._progress)
     self.assertEqual(5, download._total_size)
 
+  def test_media_io_base_download_custom_request_headers(self):
+    self.request.http = HttpMockSequence([
+      ({'status': '200',
+        'content-range': '0-2/5'}, 'echo_request_headers_as_json'),
+      ({'status': '200',
+        'content-range': '3-4/5'}, 'echo_request_headers_as_json'),
+    ])
+    self.assertEqual(True, self.request.http.follow_redirects)
+
+    self.request.headers['Cache-Control'] = 'no-store'
+
+    download = MediaIoBaseDownload(
+        fd=self.fd, request=self.request, chunksize=3)
+
+    self.assertEqual(download._headers, {'Cache-Control':'no-store'})
+
+    status, done = download.next_chunk()
+
+    result = self.fd.getvalue().decode('utf-8')
+
+    # we abuse the internals of the object we're testing, pay no attention
+    # to the actual bytes= values here; we are just asserting that the
+    # header we added to the original request is sent up to the server
+    # on each call to next_chunk
+
+    self.assertEqual(json.loads(result),
+                     {"Cache-Control": "no-store", "range": "bytes=0-3"})
+
+    download._fd = self.fd = BytesIO()
+    status, done = download.next_chunk()
+
+    result = self.fd.getvalue().decode('utf-8')
+    self.assertEqual(json.loads(result),
+                     {"Cache-Control": "no-store", "range": "bytes=51-54"})
+
   def test_media_io_base_download_handle_redirects(self):
     self.request.http = HttpMockSequence([
       ({'status': '200',