Add GoogleDriveToLocalOperator (#14191)

Add new operator to download file from Google Drive to local filesystem. Co-authored-by: Tomek Urbaszek <[email protected]>
apache · Feb 15, 2021 · 1ab4066 · 1ab4066
1 parent e4629b6
commit 1ab4066
Show file tree

Hide file tree

Showing 6 changed files with 268 additions and 2 deletions.
diff --git a/airflow/providers/google/cloud/example_dags/example_gdrive_to_local.py b/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
diff --git a/airflow/providers/google/cloud/transfers/gdrive_to_local.py b/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)
diff --git a/airflow/providers/google/provider.yaml b/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

diff --git a/airflow/providers/google/suite/hooks/drive.py b/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)
diff --git a/docs/apache-airflow-providers-google/operators/transfer/gdrive_to_local.rst b/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>`__
diff --git a/tests/providers/google/cloud/transfers/test_gdrive_to_local.py b/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
+ )