[airflow] branch main updated: Support Uploading Bigger Files to Google Drive (#22179)

[po...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git


The following commit(s) were added to refs/heads/main by this push:
     new 4543539  Support Uploading Bigger Files to Google Drive (#22179)
4543539 is described below

commit 45435399fe16c48d11584245a04c357242b23f40
Author: Ulas Can Cengiz <72...@users.noreply.github.com>
AuthorDate: Sun Mar 13 15:10:50 2022 +0100

    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
---
 airflow/providers/google/suite/hooks/drive.py | 17 +++++++++++++++--
 1 file changed, 15 insertions(+), 2 deletions(-)

diff --git a/airflow/providers/google/suite/hooks/drive.py b/airflow/providers/google/suite/hooks/drive.py
index 5c14bfd..10144a1 100644
--- a/airflow/providers/google/suite/hooks/drive.py
+++ b/airflow/providers/google/suite/hooks/drive.py
@@ -180,12 +180,25 @@ class GoogleDriveHook(GoogleBaseHook):
             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 @@ class GoogleDriveHook(GoogleBaseHook):
             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)