Refactor GoogleDriveToGCSOperator to use common methods (#14276)

Scuall1992 · web-flow · commit a7e4266d675d · 2021-02-21T10:40:35.000+01:00
Refactor GoogleDriveToGCSOperator to use common methods implemented
in hooks used by this operator.
diff --git a/airflow/providers/google/cloud/transfers/gdrive_to_gcs.py b/airflow/providers/google/cloud/transfers/gdrive_to_gcs.py
@@ -15,7 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
-from io import BytesIO
+import warnings
 from typing import Optional, Sequence, Union
 
 from airflow.models import BaseOperator
@@ -32,11 +32,15 @@ class GoogleDriveToGCSOperator(BaseOperator):
         For more information on how to use this operator, take a look at the guide:
         :ref:`howto/operator:GoogleDriveToGCSOperator`
 
-    :param destination_bucket: The destination Google cloud storage bucket where the
+    :param bucket_name: The destination Google cloud storage bucket where the
         file should be written to
-    :type destination_bucket: str
-    :param destination_object: The Google Cloud Storage object name for the object created by the operator.
+    :type bucket_name: str
+    :param object_name: The Google Cloud Storage object name for the object created by the operator.
         For example: ``path/to/my/file/file.txt``.
+    :type object_name: str
+    :param destination_bucket: Same as bucket_name, but for backward compatibly
+    :type destination_bucket: str
+    :param destination_object: Same as object_name, but for backward compatibly
     :type destination_object: str
     :param folder_id: The folder id of the folder in which the Google Drive file resides
     :type folder_id: str
@@ -62,6 +66,8 @@ class GoogleDriveToGCSOperator(BaseOperator):
     """
 
     template_fields = [
+        "bucket_name",
+        "object_name",
         "destination_bucket",
         "destination_object",
         "folder_id",
@@ -74,8 +80,10 @@ class GoogleDriveToGCSOperator(BaseOperator):
     def __init__(
         self,
         *,
-        destination_bucket: str,
-        destination_object: str,
+        bucket_name: Optional[str] = None,
+        object_name: Optional[str] = None,
+        destination_bucket: Optional[str] = None,  # deprecated
+        destination_object: Optional[str] = None,  # deprecated
         file_name: str,
         folder_id: str,
         drive_id: Optional[str] = None,
@@ -85,38 +93,18 @@ def __init__(
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
-        self.destination_bucket = destination_bucket
-        self.destination_object = destination_object
+        self.bucket_name = destination_bucket or bucket_name
+        if destination_bucket:
+            warnings.warn("`destination_bucket` is deprecated please use `bucket_name`", DeprecationWarning)
+        self.object_name = destination_object or object_name
+        if destination_object:
+            warnings.warn("`destination_object` is deprecated please use `object_name`", DeprecationWarning)
         self.folder_id = folder_id
         self.drive_id = drive_id
         self.file_name = file_name
         self.gcp_conn_id = gcp_conn_id
         self.delegate_to = delegate_to
         self.impersonation_chain = impersonation_chain
-        self.file_metadata = None
-
-    def _set_file_metadata(self, gdrive_hook):
-        if not self.file_metadata:
-            self.file_metadata = gdrive_hook.get_file_id(
-                folder_id=self.folder_id, file_name=self.file_name, drive_id=self.drive_id
-            )
-        return self.file_metadata
-
-    def _upload_data(self, gcs_hook: GCSHook, gdrive_hook: GoogleDriveHook) -> str:
-        file_handle = BytesIO()
-        self._set_file_metadata(gdrive_hook=gdrive_hook)
-        file_id = self.file_metadata["id"]
-        mime_type = self.file_metadata["mime_type"]
-        request = gdrive_hook.get_media_request(file_id=file_id)
-        gdrive_hook.download_content_from_request(
-            file_handle=file_handle, request=request, chunk_size=104857600
-        )
-        gcs_hook.upload(
-            bucket_name=self.destination_bucket,
-            object_name=self.destination_object,
-            data=file_handle.getvalue(),
-            mime_type=mime_type,
-        )
 
     def execute(self, context):
         gdrive_hook = GoogleDriveHook(
@@ -129,4 +117,10 @@ def execute(self, context):
             delegate_to=self.delegate_to,
             impersonation_chain=self.impersonation_chain,
         )
-        self._upload_data(gdrive_hook=gdrive_hook, gcs_hook=gcs_hook)
+        file_metadata = gdrive_hook.get_file_id(
+            folder_id=self.folder_id, file_name=self.file_name, drive_id=self.drive_id
+        )
+        with gcs_hook.provide_file_and_upload(
+            bucket_name=self.bucket_name, object_name=self.object_name
+        ) as file:
+            gdrive_hook.download_file(file_id=file_metadata["id"], file_handle=file)
diff --git a/tests/providers/google/cloud/transfers/test_gdrive_to_gcs.py b/tests/providers/google/cloud/transfers/test_gdrive_to_gcs.py
@@ -30,48 +30,9 @@
 
 
 class TestGoogleDriveToGCSOperator:
-    @mock.patch("airflow.providers.google.cloud.transfers.gdrive_to_gcs.BytesIO")
-    @mock.patch("airflow.providers.google.cloud.transfers.gdrive_to_gcs.GoogleDriveHook")
-    def test_upload_data(self, mock_gdrive_hook, mock_file_handle):
-        mock_gdrive_hook.return_value.get_media_request.return_value = mock.MagicMock()
-
-        file_id = mock_gdrive_hook.get_file_id.return_value["id"]
-        mime_type = mock_gdrive_hook.get_file_id.return_value["mime_type"]
-
-        mock_gcs_hook = mock.Mock()
-
-        op = GoogleDriveToGCSOperator(
-            task_id="test_task",
-            folder_id=FOLDER_ID,
-            file_name=FILE_NAME,
-            drive_id=DRIVE_ID,
-            destination_bucket=BUCKET,
-            destination_object=OBJECT,
-        )
-
-        op._upload_data(
-            gcs_hook=mock_gcs_hook,
-            gdrive_hook=mock_gdrive_hook,
-        )
-        # Test writing to file
-        mock_gdrive_hook.get_media_request.assert_called_once_with(file_id=file_id)
-        mock_gdrive_hook.download_content_from_request.assert_called_once_with(
-            file_handle=mock_file_handle(),
-            request=mock_gdrive_hook.get_media_request.return_value,
-            chunk_size=104857600,
-        )
-
-        # Test upload
-        mock_gcs_hook.upload.assert_called_once_with(
-            bucket_name=BUCKET, object_name=OBJECT, data=mock_file_handle().getvalue(), mime_type=mime_type
-        )
-
     @mock.patch("airflow.providers.google.cloud.transfers.gdrive_to_gcs.GCSHook")
     @mock.patch("airflow.providers.google.cloud.transfers.gdrive_to_gcs.GoogleDriveHook")
-    @mock.patch(
-        "airflow.providers.google.cloud.transfers.gdrive_to_gcs.GoogleDriveToGCSOperator._upload_data"
-    )
-    def test_execute(self, mock_upload_data, mock_gdrive_hook, mock_gcs_hook):
+    def test_execute(self, mock_gdrive_hook, mock_gcs_hook):
         context = {}
         op = GoogleDriveToGCSOperator(
             task_id="test_task",
@@ -83,15 +44,18 @@ def test_execute(self, mock_upload_data, mock_gdrive_hook, mock_gcs_hook):
             gcp_conn_id=GCP_CONN_ID,
             impersonation_chain=IMPERSONATION_CHAIN,
         )
+        meta = {"id": "123xyz"}
+        mock_gdrive_hook.return_value.get_file_id.return_value = meta
+
         op.execute(context)
+        mock_gdrive_hook.return_value.get_file_id.assert_called_once_with(
+            folder_id=FOLDER_ID, file_name=FILE_NAME, drive_id=DRIVE_ID
+        )
 
-        mock_gdrive_hook.assert_called_once_with(
-            gcp_conn_id=GCP_CONN_ID,
-            delegate_to=None,
-            impersonation_chain=IMPERSONATION_CHAIN,
+        mock_gdrive_hook.return_value.download_file.assert_called_once_with(
+            file_id=meta["id"], file_handle=mock.ANY
         )
-        mock_gcs_hook.assert_called_once_with(
-            gcp_conn_id=GCP_CONN_ID,
-            delegate_to=None,
-            impersonation_chain=IMPERSONATION_CHAIN,
+
+        mock_gcs_hook.return_value.provide_file_and_upload.assert_called_once_with(
+            bucket_name=BUCKET, object_name=OBJECT
         )
diff --git a/tests/providers/google/cloud/transfers/test_gdrive_to_local.py b/tests/providers/google/cloud/transfers/test_gdrive_to_local.py
@@ -35,13 +35,12 @@ def test_execute(self, hook_mock):
                 file_name=FILE_NAME,
                 output_file=temp_file.name,
             )
+            meta = {"id": "123xyz"}
+            hook_mock.return_value.get_file_id.return_value = meta
+
             op.execute(context=None)
             hook_mock.assert_called_once_with(delegate_to=None, impersonation_chain=None)
 
-            hook_mock.return_value.get_file_id.assert_called_once_with(
-                folder_id=FOLDER_ID, file_name=FILE_NAME, drive_id=None
-            )
-
             hook_mock.return_value.download_file.assert_called_once_with(
-                file_id=mock.ANY, file_handle=mock.ANY
+                file_id=meta["id"], file_handle=mock.ANY
             )
