Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Add gcp_conn_id argument to GoogleDriveToLocalOperator #24622

Merged
merged 5 commits into from
Jun 24, 2022
Merged

Add gcp_conn_id argument to GoogleDriveToLocalOperator #24622

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
e3abe8f
adding option to use non-default gcp_conn_id
aspain Jun 23, 2022
37b09c4
updating test file
aspain Jun 23, 2022
12ff933
Merge branch 'main' into patch-2
aspain Jun 23, 2022
bedec30
black formatting
aspain Jun 23, 2022
431a493
Fix code format
uranusjr Jun 24, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions airflow/providers/google/cloud/transfers/gdrive_to_local.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand All @@ -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,
)
Expand Down
6 changes: 5 additions & 1 deletion tests/providers/google/cloud/transfers/test_gdrive_to_local.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -33,13 +34,16 @@ 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
Expand Down