You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by ur...@apache.org on 2022/06/24 09:40:15 UTC
[airflow] branch main updated: Add gcp_conn_id argument to GoogleDriveToLocalOperator (#24622)
This is an automated email from the ASF dual-hosted git repository.
uranusjr 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 ded22eb5b6 Add gcp_conn_id argument to GoogleDriveToLocalOperator (#24622)
ded22eb5b6 is described below
commit ded22eb5b65bbc789c3f1842402e343070b96f19
Author: aspain <ad...@gmail.com>
AuthorDate: Fri Jun 24 05:40:05 2022 -0400
Add gcp_conn_id argument to GoogleDriveToLocalOperator (#24622)
Co-authored-by: Tzu-ping Chung <ur...@gmail.com>
---
airflow/providers/google/cloud/transfers/gdrive_to_local.py | 4 ++++
tests/providers/google/cloud/transfers/test_gdrive_to_local.py | 6 +++++-
2 files changed, 9 insertions(+), 1 deletion(-)
diff --git a/airflow/providers/google/cloud/transfers/gdrive_to_local.py b/airflow/providers/google/cloud/transfers/gdrive_to_local.py
index c61f2db99e..ba8e1f51d3 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 @@ class GoogleDriveToLocalOperator(BaseOperator):
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 @@ class GoogleDriveToLocalOperator(BaseOperator):
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,
)
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 65ecbdece1..a2c5c672d7 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 @@ from airflow.providers.google.cloud.transfers.gdrive_to_local import GoogleDrive
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,16 @@ class TestGoogleDriveToLocalOperator(TestCase):
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