Support Uploading Bigger Files to Google Drive (#22179)

Add `chunk_size` & `resumable` as parameters to `upload_file` method Change the default `chunk_size` to a clear representation & fix documentation typo
apache · Mar 13, 2022 · 4543539 · 4543539
1 parent 69f5ab1
commit 4543539
Showing 1 changed file with 15 additions and 2 deletions.
diff --git a/airflow/providers/google/suite/hooks/drive.py b/airflow/providers/google/suite/hooks/drive.py
@@ -180,12 +180,25 @@ def get_file_id(self, folder_id: str, file_name: str, drive_id: Optional[str] =
  file_metadata = {"id": files['files'][0]['id'], "mime_type": files['files'][0]['mimeType']}
  return file_metadata
 
- def upload_file(self, local_location: str, remote_location: str) -> str:
+ def upload_file(
+ self,
+ local_location: str,
+ remote_location: str,
+ chunk_size: int = 100 * 1024 * 1024,
+ resumable: bool = False,
+ ) -> str:
  """
  Uploads a file that is available locally to a Google Drive service.
 
  :param local_location: The path where the file is available.
  :param remote_location: The path where the file will be send
+ :param chunk_size: File will be uploaded in chunks of this many bytes. Only
+ used if resumable=True. Pass in a value of -1 if the file is to be
+ uploaded as a single chunk. Note that Google App Engine has a 5MB limit
+ on request size, so you should never set your chunk size larger than 5MB,
+ or to -1.
+ :param resumable: True if this is a resumable upload. False means upload
+ in a single request.
  :return: File ID
  :rtype: str
  """
@@ -197,7 +210,7 @@ def upload_file(self, local_location: str, remote_location: str) -> str:
  parent = "root"
 
  file_metadata = {"name": file_name, "parents": [parent]}
- media = MediaFileUpload(local_location)
+ media = MediaFileUpload(local_location, chunksize=chunk_size, resumable=resumable)
  file = (
  service.files()
  .create(body=file_metadata, media_body=media, fields="id", supportsAllDrives=True)