From e3abe8f6539ca5d859dd1f19e9bc64982ec6a09c Mon Sep 17 00:00:00 2001 From: aspain Date: Thu, 23 Jun 2022 13:57:59 -0400 Subject: [PATCH 1/4] adding option to use non-default gcp_conn_id --- airflow/providers/google/cloud/transfers/gdrive_to_local.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/airflow/providers/google/cloud/transfers/gdrive_to_local.py b/airflow/providers/google/cloud/transfers/gdrive_to_local.py index c61f2db99e116..ba8e1f51d37bc 100644 --- a/airflow/providers/google/cloud/transfers/gdrive_to_local.py +++ b/airflow/providers/google/cloud/transfers/gdrive_to_local.py @@ -35,6 +35,7 @@ class GoogleDriveToLocalOperator(BaseOperator): :param output_file: Path to downloaded file :param folder_id: The folder id of the folder in which the Google Drive file resides :param file_name: The name of the file residing in Google Drive + :param gcp_conn_id: The GCP connection ID to use when fetching connection info. :param drive_id: Optional. The id of the shared Google Drive in which the file resides. :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 @@ -64,6 +65,7 @@ def __init__( file_name: str, folder_id: str, drive_id: Optional[str] = None, + gcp_conn_id: str = "google_cloud_default", delegate_to: Optional[str] = None, impersonation_chain: Optional[Union[str, Sequence[str]]] = None, **kwargs, @@ -73,12 +75,14 @@ def __init__( 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 def execute(self, context: 'Context'): self.log.info('Executing download: %s into %s', self.file_name, self.output_file) gdrive_hook = GoogleDriveHook( + gcp_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to, impersonation_chain=self.impersonation_chain, ) From 37b09c4656fd4a81f8fcd117c3cdd2a29c6ba733 Mon Sep 17 00:00:00 2001 From: aspain Date: Thu, 23 Jun 2022 15:19:58 -0400 Subject: [PATCH 2/4] updating test file --- .../providers/google/cloud/transfers/test_gdrive_to_local.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/providers/google/cloud/transfers/test_gdrive_to_local.py b/tests/providers/google/cloud/transfers/test_gdrive_to_local.py index 65ecbdece17f9..51c2348f80d42 100644 --- a/tests/providers/google/cloud/transfers/test_gdrive_to_local.py +++ b/tests/providers/google/cloud/transfers/test_gdrive_to_local.py @@ -23,6 +23,7 @@ TASK_ID = "test-drive-to-local-operator" FOLDER_ID = "1234567890qwerty" FILE_NAME = "file.pdf" +GCP_CONN_ID = "google_cloud_default" class TestGoogleDriveToLocalOperator(TestCase): @@ -33,13 +34,14 @@ def test_execute(self, hook_mock): task_id=TASK_ID, folder_id=FOLDER_ID, file_name=FILE_NAME, + gcp_conn_id=GCP_CONN_ID, 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.assert_called_once_with(delegate_to=None, gcp_conn_id=GCP_CONN_ID, impersonation_chain=None) hook_mock.return_value.download_file.assert_called_once_with( file_id=meta["id"], file_handle=mock.ANY From bedec302a43c63578bb7a0155e4a699a72a8bf92 Mon Sep 17 00:00:00 2001 From: aspain Date: Thu, 23 Jun 2022 16:55:54 -0400 Subject: [PATCH 3/4] black formatting --- .../google/cloud/transfers/test_gdrive_to_local.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/tests/providers/google/cloud/transfers/test_gdrive_to_local.py b/tests/providers/google/cloud/transfers/test_gdrive_to_local.py index 51c2348f80d42..a493440fe41f9 100644 --- a/tests/providers/google/cloud/transfers/test_gdrive_to_local.py +++ b/tests/providers/google/cloud/transfers/test_gdrive_to_local.py @@ -18,7 +18,9 @@ from tempfile import NamedTemporaryFile from unittest import TestCase, mock -from airflow.providers.google.cloud.transfers.gdrive_to_local import GoogleDriveToLocalOperator +from airflow.providers.google.cloud.transfers.gdrive_to_local import ( + GoogleDriveToLocalOperator, +) TASK_ID = "test-drive-to-local-operator" FOLDER_ID = "1234567890qwerty" @@ -27,7 +29,9 @@ class TestGoogleDriveToLocalOperator(TestCase): - @mock.patch("airflow.providers.google.cloud.transfers.gdrive_to_local.GoogleDriveHook") + @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( @@ -41,7 +45,9 @@ def test_execute(self, hook_mock): hook_mock.return_value.get_file_id.return_value = meta op.execute(context=None) - hook_mock.assert_called_once_with(delegate_to=None, gcp_conn_id=GCP_CONN_ID, impersonation_chain=None) + hook_mock.assert_called_once_with( + delegate_to=None, gcp_conn_id=GCP_CONN_ID, impersonation_chain=None + ) hook_mock.return_value.download_file.assert_called_once_with( file_id=meta["id"], file_handle=mock.ANY From 431a493fd40786d45b0557cbbea5569f2785ab35 Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Fri, 24 Jun 2022 14:44:38 +0800 Subject: [PATCH 4/4] Fix code format --- .../google/cloud/transfers/test_gdrive_to_local.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/tests/providers/google/cloud/transfers/test_gdrive_to_local.py b/tests/providers/google/cloud/transfers/test_gdrive_to_local.py index a493440fe41f9..a2c5c672d7b94 100644 --- a/tests/providers/google/cloud/transfers/test_gdrive_to_local.py +++ b/tests/providers/google/cloud/transfers/test_gdrive_to_local.py @@ -18,9 +18,7 @@ from tempfile import NamedTemporaryFile from unittest import TestCase, mock -from airflow.providers.google.cloud.transfers.gdrive_to_local import ( - GoogleDriveToLocalOperator, -) +from airflow.providers.google.cloud.transfers.gdrive_to_local import GoogleDriveToLocalOperator TASK_ID = "test-drive-to-local-operator" FOLDER_ID = "1234567890qwerty" @@ -29,9 +27,7 @@ class TestGoogleDriveToLocalOperator(TestCase): - @mock.patch( - "airflow.providers.google.cloud.transfers.gdrive_to_local.GoogleDriveHook" - ) + @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(