Add support/tests for loading tables from 'gzip.GzipFile'. (googleapis#5711)

tseaver · tswast · commit fa88745a7291 · 2018-08-01T16:47:51.000-07:00
Closes googleapis#5276.
diff --git a/bigquery/google/cloud/bigquery/client.py b/bigquery/google/cloud/bigquery/client.py
@@ -18,6 +18,7 @@
 
 import collections
 import functools
+import gzip
 import os
 import uuid
 
@@ -1593,10 +1594,16 @@ def _check_mode(stream):
     """
     mode = getattr(stream, 'mode', None)
 
-    if mode is not None and mode not in ('rb', 'r+b', 'rb+'):
-        raise ValueError(
-            "Cannot upload files opened in text mode:  use "
-            "open(filename, mode='rb') or open(filename, mode='r+b')")
+    if isinstance(stream, gzip.GzipFile):
+        if mode != gzip.READ:
+            raise ValueError(
+                "Cannot upload gzip files opened in write mode:  use "
+                "gzip.GzipFile(filename, mode='rb')")
+    else:
+        if mode is not None and mode not in ('rb', 'r+b', 'rb+'):
+            raise ValueError(
+                "Cannot upload files opened in text mode:  use "
+                "open(filename, mode='rb') or open(filename, mode='r+b')")
 
 
 def _get_upload_headers(user_agent):
diff --git a/bigquery/tests/unit/test_client.py b/bigquery/tests/unit/test_client.py
@@ -16,6 +16,7 @@
 import datetime
 import decimal
 import email
+import gzip
 import io
 import json
 import unittest
@@ -3712,6 +3713,12 @@ def _make_do_upload_patch(cls, client, method,
     def _make_file_obj():
         return io.BytesIO(b'hello, is it me you\'re looking for?')
 
+    def _make_gzip_file_obj(self, writable):
+        if writable:
+            return gzip.GzipFile(mode='w', fileobj=io.BytesIO())
+        else:
+            return gzip.GzipFile(mode='r', fileobj=self._make_file_obj())
+
     @staticmethod
     def _make_config():
         from google.cloud.bigquery.job import LoadJobConfig
@@ -3892,6 +3899,33 @@ def test_load_table_from_file_with_rewind(self):
 
         assert file_obj.tell() == 0
 
+    def test_load_table_from_file_with_readable_gzip(self):
+        from google.cloud.bigquery.client import _DEFAULT_NUM_RETRIES
+
+        client = self._make_client()
+        gzip_file = self._make_gzip_file_obj(writable=False)
+
+        do_upload_patch = self._make_do_upload_patch(
+            client, '_do_resumable_upload', self.EXPECTED_CONFIGURATION)
+        with do_upload_patch as do_upload:
+            client.load_table_from_file(
+                gzip_file, self.TABLE_REF, job_id='job_id',
+                job_config=self._make_config())
+
+        do_upload.assert_called_once_with(
+            gzip_file,
+            self.EXPECTED_CONFIGURATION,
+            _DEFAULT_NUM_RETRIES)
+
+    def test_load_table_from_file_with_writable_gzip(self):
+        client = self._make_client()
+        gzip_file = self._make_gzip_file_obj(writable=True)
+
+        with pytest.raises(ValueError):
+            client.load_table_from_file(
+                gzip_file, self.TABLE_REF, job_id='job_id',
+                job_config=self._make_config())
+
     def test_load_table_from_file_failure(self):
         from google.resumable_media import InvalidResponse
         from google.cloud import exceptions
