Download small files from GCS (cruxinformatics#20)

We want all files downloaded from GCS, unless the user chooses to
restrict to Crux domains. Even small files should come from GCS, just
not using ChunkedDownload.

Change file downloads to use plain requests for downloading small files
from signed URLs.

There is some additional work that still needs to be done to make the
file downloads respect the client proxy settings, and throw crux-python
exceptions instead requests exceptions.

Also:

- I got bothered that `file_pointer` wasn't technically a C file
  pointer, change it to `file_obj` instead.
- Then I got really OCD about `local_path` for `download()` and
  `upload()` methods, because it it took either a file-like object or a
  file path. Change to `dest` for downloads and `src` for uploads.
- Update docs with arg name changes, and remove args names when they
  are simple and required.
- Simplify label docs.

The integration tests were very unreliable. Use random names more, and
change to longer random names (I actually hit duplicate random names).
Skip a `test_folder_add_delete_permission()` because it is way too
slow, or maybe doesn't even work. Skip `test_set_datasets_provenance()`
because it is either really flaky or doesn't work.

It isn't ideal to be skipping testing, they will be fixed another PR.

**Test Plan:**

All tests pass, lint, integration, type checking, formatting.

I manually tested with a download script, with urllib3 DEBUG logging
enabled, and was able to see the downloads coming from GCS. And I set
`only_use_crux_domains=True` and it worked, and I was able to see the
download coming from the API.

.gitignore

@@ -32,3 +32,4 @@ __pycache__/
 /.nox/
 .pytest_cache/
 .mypy_cache/
+/pip-wheel-metadata/

crux/models/dataset.py

@@ -533,7 +533,7 @@ def download_files(self, folder, local_path):
             elif resource.type == "file":
                 file_resource = File.from_dict(resource.to_dict())
                 file_resource.connection = self.connection
-                file_resource.download(local_path=resource_local_path)
+                file_resource.download(resource_local_path)
                 local_file_list.append(resource_local_path)
 
         return local_file_list
@@ -592,9 +592,9 @@ def upload_files(
 
             elif os.path.isfile(content_local_path):
                 fil_o = self.upload_file(
+                    content_local_path,
+                    content_path,
                     media_type=media_type,
-                    path=content_path,
-                    local_path=content_local_path,
                     tags=tags,
                     description=description,
                 )
@@ -716,16 +716,14 @@ def load_table_from_file(self, source_file, dest_table, append=False):
             headers=headers,
         )
 
-    def upload_file(
-        self, local_path, path, media_type=None, description=None, tags=None
-    ):
+    def upload_file(self, src, dest, media_type=None, description=None, tags=None):
         # type: (Union[IO, str], str, str, str, List[str]) -> File
         """Uploads the File.
 
         Args:
-            local_path (str or file): Local OS path whose content
+            src (str or file): Local OS path whose content
                 is to be uploaded to file resource.
-            path (str): File resource path.
+            dest (str): File resource path.
             media_type (str): Content type of the file. Defaults to None.
             description (str): Description of the file. Defaults to None.
             tags (:obj:`list` of :obj:`str`): Tags to be attached to the file resource.
@@ -734,20 +732,20 @@ def upload_file(
             crux.models.File: File Object.
 
         Raises:
-            TypeError: If local_path is not file or string object.
+            TypeError: If src is not file or string object.
             LookupError: If media type is not a valid type.
             CruxClientError: If error occurs in api or in client.
         """
 
         tags = tags if tags else []
 
-        file_resource = self.create_file(tags=tags, description=description, path=path)
+        file_resource = self.create_file(tags=tags, description=description, path=dest)
 
-        if hasattr(local_path, "write"):
+        if hasattr(src, "write"):
 
             if media_type is None:
                 try:
-                    media_type = MediaType.detect(getattr(local_path, "name"))
+                    media_type = MediaType.detect(getattr(src, "name"))
                 except LookupError as err:
                     file_resource.delete()
                     raise LookupError(err)
@@ -758,27 +756,27 @@ def upload_file(
                 return self.connection.api_call(
                     "PUT",
                     ["resources", file_resource.id, "content"],
-                    data=local_path,
+                    data=src,
                     headers=headers,
                     model=File,
                 )
             except (CruxClientError, CruxAPIError) as err:
                 file_resource.delete()
                 raise CruxClientError(err.message)
 
-        elif isinstance(local_path, str):
+        elif isinstance(src, str):
 
             if media_type is None:
                 try:
-                    media_type = MediaType.detect(local_path)
+                    media_type = MediaType.detect(src)
                 except LookupError as err:
                     file_resource.delete()
                     raise LookupError(err)
 
             headers = {"Content-Type": media_type, "Accept": "application/json"}
 
             try:
-                with open(local_path, mode="rb") as data:
+                with open(src, mode="rb") as data:
                     return self.connection.api_call(
                         "PUT",
                         ["resources", file_resource.id, "content"],
@@ -791,7 +789,7 @@ def upload_file(
                 raise CruxClientError(str(err))
 
         else:
-            raise TypeError("Invalid Data Type for local_path")
+            raise TypeError("Invalid Data Type for src")
 
     def create_query(self, path, config, tags=None, description=None):
         # type: (str, Dict[str, Any], List[str], str) -> Query

crux/models/file.py

@@ -51,10 +51,21 @@ def _get_signed_url(self):
 
         return url
 
-    def _dl_signed_url_resumable(self, file_pointer, chunk_size=DEFAULT_CHUNK_SIZE):
-
+    def _dl_signed_url(self, file_obj, chunk_size=DEFAULT_CHUNK_SIZE):
+        """Download from signed URL using requests directly, not google-resumable-media."""
         signed_url = self._get_signed_url()
+        transport = get_signed_url_session()
+
+        with transport.get(signed_url, stream=True) as response:
+            response.raise_for_status()
+            for chunk in response.iter_content(chunk_size=chunk_size):
+                file_obj.write(chunk)
+
+        return True
 
+    def _dl_signed_url_resumable(self, file_obj, chunk_size=DEFAULT_CHUNK_SIZE):
+        """Download from signed URL using google-resumable-media."""
+        signed_url = self._get_signed_url()
         transport = get_signed_url_session()
 
         # Track how many bytes the client has downloaded since the last time they
@@ -69,7 +80,7 @@ def _dl_signed_url_resumable(self, file_pointer, chunk_size=DEFAULT_CHUNK_SIZE):
         max_url_refreshes_without_progress = 5
         max_url_refreshes = 100
 
-        download = ChunkedDownload(signed_url, chunk_size, file_pointer)
+        download = ChunkedDownload(signed_url, chunk_size, file_obj)
 
         while not download.finished:
             try:
@@ -108,7 +119,7 @@ def _dl_signed_url_resumable(self, file_pointer, chunk_size=DEFAULT_CHUNK_SIZE):
                 download = ChunkedDownload(
                     new_signed_url,
                     chunk_size,
-                    file_pointer,
+                    file_obj,
                     start=sum_total_bytes_from_urls,
                 )
         return True
@@ -139,91 +150,94 @@ def iter_content(self, chunk_size=DEFAULT_CHUNK_SIZE, decode_unicode=False):
 
         return data.iter_content(chunk_size=chunk_size, decode_unicode=decode_unicode)
 
-    def _download_file(self, file_pointer, chunk_size=DEFAULT_CHUNK_SIZE):
+    def _download_file(self, file_obj, chunk_size=DEFAULT_CHUNK_SIZE):
         # google-resumable-media has a bug where is expects the 'content-range' even
         # for 200 OK responses, which happens when the range is larger than the size.
         # There isn't much point in using resumable media for small files.
         # Make sure size is greater than 2x chunk_size if Google Resumable media is
         # to be used.
         small_enough = self.size < (chunk_size * 2)
 
-        if self.connection.crux_config.only_use_crux_domains or small_enough:
+        # If we must use only Crux domains, download via the API.
+        if self.connection.crux_config.only_use_crux_domains:
             return self._download(
-                file_pointer=file_pointer, media_type=None, chunk_size=chunk_size
+                file_obj=file_obj, media_type=None, chunk_size=chunk_size
             )
+        # Use requests directly for small files.
+        elif small_enough:
+            return self._dl_signed_url(file_obj=file_obj, chunk_size=chunk_size)
+        # Use google-resumable-media for large files
         else:
             return self._dl_signed_url_resumable(
-                file_pointer=file_pointer, chunk_size=chunk_size
+                file_obj=file_obj, chunk_size=chunk_size
             )
 
-    def download(self, local_path, chunk_size=DEFAULT_CHUNK_SIZE):
+    def download(self, dest, chunk_size=DEFAULT_CHUNK_SIZE):
         # type: (str, int) -> bool
         """Downloads the file resource.
 
         Args:
-            local_path (str or file): Local OS path at which file resource will be downloaded.
+            dest (str or file): Local OS path at which file resource will be downloaded.
             chunk_size (int): Number of bytes to be read in memory.
 
         Returns:
             bool: True if it is downloaded.
 
         Raises:
-            TypeError: If local_path is not a file like or string type.
+            TypeError: If dest is not a file like or string type.
         """
-        if hasattr(local_path, "write"):
-            return self._download_file(local_path, chunk_size=chunk_size)
-        elif isinstance(local_path, (str, unicode)):
-            with open(local_path, "wb") as file_pointer:
-                return self._download_file(file_pointer, chunk_size=chunk_size)
+        if not valid_chunk_size(chunk_size):
+            raise ValueError("chunk_size should be multiple of 256 KiB")
+
+        if hasattr(dest, "write"):
+            return self._download_file(dest, chunk_size=chunk_size)
+        elif isinstance(dest, (str, unicode)):
+            with open(dest, "wb") as file_obj:
+                return self._download_file(file_obj, chunk_size=chunk_size)
         else:
-            raise TypeError(
-                "Invalid Data Type for local_path: {}".format(type(local_path))
-            )
+            raise TypeError("Invalid Data Type for dest: {}".format(type(dest)))
 
-    def upload(self, local_path, media_type=None):
+    def upload(self, src, media_type=None):
         # type: (Union[IO, str], str) -> bool
         """Uploads the content to empty file resource.
 
         Args:
-            local_path (str or file): Local OS path whose content is to be uploaded.
+            src (str or file): Local OS path whose content is to be uploaded.
             media_type (str): Content type of the file. Defaults to None.
 
         Returns
             bool: True if it is uploaded.
 
         Raises:
-            TypeError: If local_path type is invalid.
+            TypeError: If src type is invalid.
         """
 
-        if hasattr(local_path, "read"):
+        if hasattr(src, "read"):
 
             if media_type is None:
-                media_type = MediaType.detect(getattr(local_path, "name"))
+                media_type = MediaType.detect(getattr(src, "name"))
 
             headers = {"Content-Type": media_type, "Accept": "application/json"}
 
             resp = self.connection.api_call(
-                "PUT",
-                ["resources", self.id, "content"],
-                data=local_path,
-                headers=headers,
+                "PUT", ["resources", self.id, "content"], data=src, headers=headers
             )
 
             return resp.status_code == 200
 
-        elif isinstance(local_path, str):
+        elif isinstance(src, str):
 
             if media_type is None:
-                media_type = MediaType.detect(local_path)
+                media_type = MediaType.detect(src)
 
             headers = {"Content-Type": media_type, "Accept": "application/json"}
 
-            with open(local_path, mode="rb") as data:
+            with open(src, mode="rb") as data:
                 resp = self.connection.api_call(
                     "PUT", ["resources", self.id, "content"], data=data, headers=headers
                 )
 
             return resp.status_code == 200
 
         else:
-            raise TypeError("Invalid Data Type for local_path")
+            raise TypeError("Invalid Data Type for src")

crux/models/query.py

@@ -70,14 +70,14 @@ def run(
         return data.iter_content(chunk_size=chunk_size, decode_unicode=decode_unicode)
 
     def download(
-        self, local_path, format="csv", params=None
+        self, dest, format="csv", params=None
     ):  # It is by design pylint: disable=redefined-builtin
         # type: (str, str, Dict[Any, Any]) -> bool
         """Method which streams the Query
 
         Args:
-            local_path (str): Local OS path at which resource will be downloaded.
-            format (str): Output format of the query. Defaults to csv.
+            dest (str): Local OS path at which resource will be downloaded.
+            media_type (str): Output format of the query. Defaults to csv.
             params (dict): Run parameters. Defaults to None.
 
         Returns:
@@ -95,7 +95,7 @@ def download(
             headers=headers,
         )
 
-        with open(local_path, "w") as local_file:
+        with open(dest, "w") as local_file:
             for line in data.iter_lines():
                 if line:
                     dcd_line = line.decode("utf-8")

crux/models/resource.py

@@ -452,7 +452,7 @@ def _get_folder(self):
 
         return response.json().get("path")
 
-    def _download(self, file_pointer, media_type, chunk_size=DEFAULT_CHUNK_SIZE):
+    def _download(self, file_obj, media_type, chunk_size=DEFAULT_CHUNK_SIZE):
 
         if media_type is not None:
             headers = {"Accept": media_type}
@@ -464,7 +464,7 @@ def _download(self, file_pointer, media_type, chunk_size=DEFAULT_CHUNK_SIZE):
         )
 
         for chunk in data.iter_content(chunk_size=chunk_size):
-            file_pointer.write(chunk)
+            file_obj.write(chunk)
 
         return True
 

crux/models/table.py

@@ -26,31 +26,27 @@ def to_dict(self):
             "folder": self.folder,
         }
 
-    def download(self, local_path, media_type, chunk_size=DEFAULT_CHUNK_SIZE):
+    def download(self, dest, media_type, chunk_size=DEFAULT_CHUNK_SIZE):
         # type: (str, str, int) -> bool
         """Downloads the table resource.
 
         Args:
-            local_path (str or file): Local OS path at which file resource will be downloaded.
+            dest (str or file): Local OS path at which file resource will be downloaded.
             media_type (str): Content Type for download.
             chunk_size (int): Number of bytes to be read in memory.
 
         Returns:
             bool: True if it is downloaded.
 
         Raises:
-            TypeError: If local_path is not a file like or string type.
+            TypeError: If dest is not a file like or string type.
         """
-        if hasattr(local_path, "write"):
-            return self._download(
-                local_path, media_type=media_type, chunk_size=chunk_size
-            )
-        elif isinstance(local_path, (str, unicode)):
-            with open(local_path, "wb") as file_pointer:
+        if hasattr(dest, "write"):
+            return self._download(dest, media_type=media_type, chunk_size=chunk_size)
+        elif isinstance(dest, (str, unicode)):
+            with open(dest, "wb") as file_obj:
                 return self._download(
-                    file_pointer, media_type=media_type, chunk_size=chunk_size
+                    file_obj, media_type=media_type, chunk_size=chunk_size
                 )
         else:
-            raise TypeError(
-                "Invalid Data Type for local_path: {}".format(type(local_path))
-            )
+            raise TypeError("Invalid Data Type for dest: {}".format(type(dest)))