Add GoogleDriveToLocalOperator (#14191)

Add new operator to download file from Google Drive to local filesystem. 

Co-authored-by: Tomek Urbaszek <turbaszek@gmail.com>

airflow/providers/google/cloud/example_dags/example_gdrive_to_local.py

@@ -0,0 +1,49 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import os
+
+from airflow import models
+from airflow.providers.google.cloud.transfers.gdrive_to_local import GoogleDriveToLocalOperator
+from airflow.providers.google.suite.sensors.drive import GoogleDriveFileExistenceSensor
+from airflow.utils.dates import days_ago
+
+FOLDER_ID = os.environ.get("FILE_ID", "1234567890qwerty")
+FILE_NAME = os.environ.get("FILE_NAME", "file.pdf")
+OUTPUT_FILE = os.environ.get("OUTPUT_FILE", "out_file.pdf")
+
+with models.DAG(
+    "example_gdrive_to_local_with_gdrive_sensor",
+    start_date=days_ago(1),
+    schedule_interval=None,  # Override to match your needs
+    tags=["example"],
+) as dag:
+    # [START detect_file]
+    detect_file = GoogleDriveFileExistenceSensor(
+        task_id="detect_file", folder_id=FOLDER_ID, file_name=FILE_NAME
+    )
+    # [END detect_file]
+    # [START download_from_gdrive_to_local]
+    download_from_gdrive_to_local = GoogleDriveToLocalOperator(
+        task_id="download_from_gdrive_to_local",
+        folder_id=FOLDER_ID,
+        file_name=FILE_NAME,
+        output_file=OUTPUT_FILE,
+    )
+    # [END download_from_gdrive_to_local]
+    detect_file >> download_from_gdrive_to_local

airflow/providers/google/cloud/transfers/gdrive_to_local.py

@@ -0,0 +1,95 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from typing import Optional, Sequence, Union
+
+from airflow.models import BaseOperator
+from airflow.providers.google.suite.hooks.drive import GoogleDriveHook
+from airflow.utils.decorators import apply_defaults
+
+
+class GoogleDriveToLocalOperator(BaseOperator):
+    """
+    Writes a Google Drive file into local Storage.
+
+    .. seealso::
+        For more information on how to use this operator, take a look at the guide:
+        :ref:`howto/operator:GoogleDriveToLocalOperator`
+
+    :param output_file: Path to downloaded file
+    :type output_file: str
+    :param folder_id: The folder id of the folder in which the Google Drive file resides
+    :type folder_id: str
+    :param file_name: The name of the file residing in Google Drive
+    :type file_name: str
+    :param drive_id: Optional. The id of the shared Google Drive in which the file resides.
+    :type drive_id: str
+    :param delegate_to: The account to impersonate using domain-wide delegation of authority,
+        if any. For this to work, the service account making the request must have
+        domain-wide delegation enabled.
+    :type delegate_to: str
+    :param impersonation_chain: Optional service account to impersonate using short-term
+        credentials, or chained list of accounts required to get the access_token
+        of the last account in the list, which will be impersonated in the request.
+        If set as a string, the account must grant the originating account
+        the Service Account Token Creator IAM role.
+        If set as a sequence, the identities from the list must grant
+        Service Account Token Creator IAM role to the directly preceding identity, with first
+        account from the list granting this role to the originating account (templated).
+    :type impersonation_chain: Union[str, Sequence[str]]
+    """
+
+    template_fields = [
+        "output_file",
+        "folder_id",
+        "file_name",
+        "drive_id",
+        "impersonation_chain",
+    ]
+
+    @apply_defaults
+    def __init__(
+        self,
+        *,
+        output_file: str,
+        file_name: str,
+        folder_id: str,
+        drive_id: Optional[str] = None,
+        delegate_to: Optional[str] = None,
+        impersonation_chain: Optional[Union[str, Sequence[str]]] = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.output_file = output_file
+        self.folder_id = folder_id
+        self.drive_id = drive_id
+        self.file_name = file_name
+        self.delegate_to = delegate_to
+        self.impersonation_chain = impersonation_chain
+
+    def execute(self, context):
+        self.log.info('Executing download: %s into %s', self.file_name, self.output_file)
+        gdrive_hook = GoogleDriveHook(
+            delegate_to=self.delegate_to,
+            impersonation_chain=self.impersonation_chain,
+        )
+        file_metadata = gdrive_hook.get_file_id(
+            folder_id=self.folder_id, file_name=self.file_name, drive_id=self.drive_id
+        )
+
+        with open(self.output_file, "wb") as file:
+            gdrive_hook.download_file(file_id=file_metadata["id"], file_handle=file)

airflow/providers/google/provider.yaml

@@ -699,6 +699,10 @@ transfers:
     target-integration-name: Local
     how-to-guide: /docs/apache-airflow-providers-google/operators/transfer/gcs_to_local.rst
     python-module: airflow.providers.google.cloud.transfers.gcs_to_local
+  - source-integration-name: Google Drive
+    target-integration-name: Local
+    how-to-guide: /docs/apache-airflow-providers-google/operators/transfer/gdrive_to_local.rst
+    python-module: airflow.providers.google.cloud.transfers.gdrive_to_local
   - source-integration-name: Salesforce
     target-integration-name: Google Cloud Storage (GCS)
     how-to-guide: /docs/apache-airflow-providers-google/operators/transfer/salesforce_to_gcs.rst

airflow/providers/google/suite/hooks/drive.py

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """Hook for Google Drive service"""
+from io import TextIOWrapper
 from typing import Any, Optional, Sequence, Union
 
 from googleapiclient.discovery import Resource, build
@@ -203,13 +204,13 @@ def upload_file(self, local_location: str, remote_location: str) -> str:
         :rtype: str
         """
         service = self.get_conn()
-        directory_path, _, filename = remote_location.rpartition("/")
+        directory_path, _, file_name = remote_location.rpartition("/")
         if directory_path:
             parent = self._ensure_folders_exists(directory_path)
         else:
             parent = "root"
 
-        file_metadata = {"name": filename, "parents": [parent]}
+        file_metadata = {"name": file_name, "parents": [parent]}
         media = MediaFileUpload(local_location)
         file = (
             service.files()  # pylint: disable=no-member
@@ -218,3 +219,15 @@ def upload_file(self, local_location: str, remote_location: str) -> str:
         )
         self.log.info("File %s uploaded to gdrive://%s.", local_location, remote_location)
         return file.get("id")
+
+    def download_file(self, file_id: str, file_handle: TextIOWrapper, chunk_size: int = 104857600):
+        """
+        Download a file from Google Drive.
+
+        :param file_id: the id of the file
+        :type file_id: str
+        :param file_handle: file handle used to write the content to
+        :type file_handle: io.TextIOWrapper
+        """
+        request = self.get_media_request(file_id=file_id)
+        self.download_content_from_request(file_handle=file_handle, request=request, chunk_size=chunk_size)

docs/apache-airflow-providers-google/operators/transfer/gdrive_to_local.rst

@@ -0,0 +1,58 @@
+ .. Licensed to the Apache Software Foundation (ASF) under one
+    or more contributor license agreements.  See the NOTICE file
+    distributed with this work for additional information
+    regarding copyright ownership.  The ASF licenses this file
+    to you under the Apache License, Version 2.0 (the
+    "License"); you may not use this file except in compliance
+    with the License.  You may obtain a copy of the License at
+
+ ..   http://www.apache.org/licenses/LICENSE-2.0
+
+ .. Unless required by applicable law or agreed to in writing,
+    software distributed under the License is distributed on an
+    "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+    KIND, either express or implied.  See the License for the
+    specific language governing permissions and limitations
+    under the License.
+
+
+Downloads data from Google Drive Storage to Local Filesystem
+============================================================
+The `Google Drive <https://www.google.com/drive/>`__ is
+used to store daily use data, including documents and photos. Google Drive has built-in mechanisms to facilitate group work e.g.
+document editor, file sharing mechanisms.
+
+.. contents::
+  :depth: 1
+  :local:
+
+
+Prerequisite Tasks
+^^^^^^^^^^^^^^^^^^
+
+.. include::/operators/_partials/prerequisite_tasks.rst
+
+.. _howto/operator:GoogleDriveToLocalOperator:
+
+GCSToLocalFilesystemOperator
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+:class:`~airflow.providers.google.cloud.transfers.gdrive_to_local.GoogleDriveToLocalOperator` allows you to download
+data from Google Drive to local filesystem.
+
+
+Below is an example of using this operator to download file from Google Drive to Local Filesystem.
+
+.. exampleinclude:: /../../airflow/providers/google/cloud/example_dags/example_gdrive_to_local.py
+    :language: python
+    :dedent: 4
+    :start-after: [START download_from_gdrive_to_local]
+    :end-before: [END download_from_gdrive_to_local]
+
+
+Reference
+---------
+
+For further information, look at:
+
+* `Google Drive API Documentation <https://developers.google.com/drive/api/v3/about-sdk>`__

tests/providers/google/cloud/transfers/test_gdrive_to_local.py

@@ -0,0 +1,47 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from tempfile import NamedTemporaryFile
+from unittest import TestCase, mock
+
+from airflow.providers.google.cloud.transfers.gdrive_to_local import GoogleDriveToLocalOperator
+
+TASK_ID = "test-drive-to-local-operator"
+FOLDER_ID = "1234567890qwerty"
+FILE_NAME = "file.pdf"
+
+
+class TestGoogleDriveToLocalOperator(TestCase):
+    @mock.patch("airflow.providers.google.cloud.transfers.gdrive_to_local.GoogleDriveHook")
+    def test_execute(self, hook_mock):
+        with NamedTemporaryFile("wb") as temp_file:
+            op = GoogleDriveToLocalOperator(
+                task_id=TASK_ID,
+                folder_id=FOLDER_ID,
+                file_name=FILE_NAME,
+                output_file=temp_file.name,
+            )
+            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
+            )