This is an automated email from the ASF dual-hosted git repository.
potiuk 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 56f6db03a14 Add destination_folder_id to GCSToGoogleDriveOperator
templated fields and update google drive system tests (#66930)
56f6db03a14 is described below
commit 56f6db03a14eaab789452373ff706e62665b559c
Author: Nitochkin <[email protected]>
AuthorDate: Thu Aug 27 12:56:23 2026 +0200
Add destination_folder_id to GCSToGoogleDriveOperator templated fields and
update google drive system tests (#66930)
* Change gdrive system tests
* Remove FOLDER_ID variable from example_gcs_to_drive test
Fix get_media_request to include all drives
* Fix drive test and formatting
---------
Co-authored-by: Anton Nitochkin <[email protected]>
Co-authored-by: Marcin Lubimow <[email protected]>
---
.../airflow/providers/google/suite/hooks/drive.py | 2 +-
.../google/suite/transfers/gcs_to_gdrive.py | 3 +-
.../google/cloud/gcs/example_gcs_to_gdrive.py | 31 +++++++++++++----
.../google/cloud/gcs/example_gdrive_to_gcs.py | 40 +++++++++++++++++-----
.../tests/unit/google/suite/hooks/test_drive.py | 4 ++-
5 files changed, 63 insertions(+), 17 deletions(-)
diff --git a/providers/google/src/airflow/providers/google/suite/hooks/drive.py
b/providers/google/src/airflow/providers/google/suite/hooks/drive.py
index 7139234281a..5f6a21b3e97 100644
--- a/providers/google/src/airflow/providers/google/suite/hooks/drive.py
+++ b/providers/google/src/airflow/providers/google/suite/hooks/drive.py
@@ -141,7 +141,7 @@ class GoogleDriveHook(GoogleBaseHook):
:return: request
"""
service = self.get_conn()
- request = service.files().get_media(fileId=file_id)
+ request = service.files().get_media(fileId=file_id,
supportsAllDrives=True)
return request
def exists(
diff --git
a/providers/google/src/airflow/providers/google/suite/transfers/gcs_to_gdrive.py
b/providers/google/src/airflow/providers/google/suite/transfers/gcs_to_gdrive.py
index e184f00a1bc..00a6051a153 100644
---
a/providers/google/src/airflow/providers/google/suite/transfers/gcs_to_gdrive.py
+++
b/providers/google/src/airflow/providers/google/suite/transfers/gcs_to_gdrive.py
@@ -65,7 +65,7 @@ class GCSToGoogleDriveOperator(BaseOperator):
copied to ``blah/baz``; to retain the prefix write the
destination_object as e.g. ``blah/foo``, in
which case the copied file will be named ``blah/foo/baz``.
:param destination_folder_id: The folder ID where the destination objects
will be placed. It is
- an additive prefix for anything specified in destination_object.
+ an additive prefix for anything specified in destination_object.
(templated)
For example if folder ID ``xXyYzZ`` is called ``foo`` and the
destination is ``bar/baz``, the file
will end up in `foo/bar/baz`.
This can be used to target an existing folder that is already visible
to other users. The credentials
@@ -88,6 +88,7 @@ class GCSToGoogleDriveOperator(BaseOperator):
"source_object",
"destination_object",
"impersonation_chain",
+ "destination_folder_id",
)
ui_color = "#f0eee4"
diff --git
a/providers/google/tests/system/google/cloud/gcs/example_gcs_to_gdrive.py
b/providers/google/tests/system/google/cloud/gcs/example_gcs_to_gdrive.py
index 270231dedf0..a6be20ed1f2 100644
--- a/providers/google/tests/system/google/cloud/gcs/example_gcs_to_gdrive.py
+++ b/providers/google/tests/system/google/cloud/gcs/example_gcs_to_gdrive.py
@@ -41,6 +41,7 @@ else:
from airflow.models.dag import DAG
from airflow.providers.google.cloud.operators.gcs import
GCSCreateBucketOperator, GCSDeleteBucketOperator
from airflow.providers.google.cloud.transfers.gcs_to_gcs import
GCSToGCSOperator
+from airflow.providers.google.common.utils.get_secret import get_secret
from airflow.providers.google.suite.hooks.drive import GoogleDriveHook
from airflow.providers.google.suite.transfers.gcs_to_gdrive import
GCSToGoogleDriveOperator
@@ -71,6 +72,8 @@ CURRENT_FOLDER = Path(__file__).parent
LOCAL_PATH = str(Path("gcs"))
FILE_LOCAL_PATH = str(Path(LOCAL_PATH))
FILE_NAME = "example_upload.txt"
+GDRIVE_SECRET_ID = "gdrive_shared_folder_id"
+
log = logging.getLogger(__name__)
@@ -82,6 +85,12 @@ with DAG(
tags=["example", "gcs", "gdrive"],
) as dag:
+ @task
+ def get_shared_drive_id() -> str:
+ return get_secret(secret_id=GDRIVE_SECRET_ID).strip()
+
+ get_shared_drive_id_task = get_shared_drive_id()
+
@task
def create_connection(connection_id: str):
conn_extra_json = json.dumps(
@@ -124,6 +133,7 @@ with DAG(
source_bucket=BUCKET_NAME,
source_object=f"{TMP_PATH}/{FILE_NAME}",
destination_object=f"{WORK_DIR}/copied_{FILE_NAME}",
+ destination_folder_id=get_shared_drive_id_task,
)
# [END howto_operator_gcs_to_gdrive_copy_single_file]
@@ -134,7 +144,7 @@ with DAG(
source_bucket=BUCKET_NAME,
source_object=f"{TMP_PATH}/{FILE_NAME}",
destination_object=f"{WORK_DIR}/copied_{FILE_NAME}",
- destination_folder_id=FOLDER_ID,
+ destination_folder_id=get_shared_drive_id_task,
)
# [END howto_operator_gcs_to_gdrive_copy_single_file_into_folder]
@@ -145,6 +155,7 @@ with DAG(
source_bucket=BUCKET_NAME,
source_object=f"{TMP_PATH}/*",
destination_object=f"{WORK_DIR}/",
+ destination_folder_id=get_shared_drive_id_task,
)
# [END howto_operator_gcs_to_gdrive_copy_files]
@@ -155,23 +166,31 @@ with DAG(
source_bucket=BUCKET_NAME,
source_object=f"{TMP_PATH}/*.txt",
destination_object=f"{WORK_DIR}/",
+ destination_folder_id=get_shared_drive_id_task,
move_object=True,
)
# [END howto_operator_gcs_to_gdrive_move_files]
@task(trigger_rule=TriggerRule.ALL_DONE)
- def remove_files_from_drive():
+ def remove_files_from_drive(**context):
+ ti = context["ti"]
service = GoogleDriveHook(gcp_conn_id=CONNECTION_ID).get_conn()
root_path = (
service.files()
- .list(q=f"name = '{WORK_DIR}' and mimeType =
'application/vnd.google-apps.folder'")
+ .list(
+ q=f"name = '{WORK_DIR}' and mimeType =
'application/vnd.google-apps.folder'",
+ corpora="drive",
+ driveId=ti.xcom_pull("get_shared_drive_id"),
+ includeItemsFromAllDrives=True,
+ supportsAllDrives=True,
+ )
.execute()
)
if files := root_path["files"]:
batch = service.new_batch_http_request()
for file in files:
- log.info("Preparing to remove file: %s", file)
- batch.add(service.files().delete(fileId=file["id"]))
+ log.info("Deleting file %s...", file["name"])
+ batch.add(service.files().delete(fileId=file["id"],
supportsAllDrives=True))
batch.execute()
log.info("Selected files removed.")
@@ -188,7 +207,7 @@ with DAG(
delete_connection_task = delete_connection(connection_id=CONNECTION_ID)
# TEST SETUP
- create_bucket >> [upload_file_1, upload_file_2]
+ get_shared_drive_id_task >> create_bucket >> [upload_file_1, upload_file_2]
(
[upload_file_1, upload_file_2, create_connection_task]
# TEST BODY
diff --git
a/providers/google/tests/system/google/cloud/gcs/example_gdrive_to_gcs.py
b/providers/google/tests/system/google/cloud/gcs/example_gdrive_to_gcs.py
index ec5f6ab4cca..89c63f9a35c 100644
--- a/providers/google/tests/system/google/cloud/gcs/example_gdrive_to_gcs.py
+++ b/providers/google/tests/system/google/cloud/gcs/example_gdrive_to_gcs.py
@@ -35,6 +35,7 @@ else:
from airflow.providers.google.cloud.operators.gcs import
GCSCreateBucketOperator, GCSDeleteBucketOperator
from airflow.providers.google.cloud.transfers.gcs_to_gcs import
GCSToGCSOperator
from airflow.providers.google.cloud.transfers.gdrive_to_gcs import
GoogleDriveToGCSOperator
+from airflow.providers.google.common.utils.get_secret import get_secret
from airflow.providers.google.suite.hooks.drive import GoogleDriveHook
from airflow.providers.google.suite.sensors.drive import
GoogleDriveFileExistenceSensor
from airflow.providers.google.suite.transfers.gcs_to_gdrive import
GCSToGoogleDriveOperator
@@ -60,10 +61,11 @@ BUCKET_NAME = f"bucket_{DAG_ID}_{ENV_ID}"
CONNECTION_ID = f"connection_{DAG_ID}_{ENV_ID}"
OBJECT = "abc123xyz"
-FOLDER_ID = ""
FILE_NAME = "example_upload.txt"
DRIVE_FILE_NAME = f"example_upload_{DAG_ID}_{ENV_ID}.txt"
LOCAL_PATH = f"gcs/{FILE_NAME}"
+GDRIVE_SECRET_ID = "gdrive_shared_folder_id"
+
log = logging.getLogger(__name__)
@@ -75,6 +77,12 @@ with DAG(
tags=["example", "gcs", "gdrive"],
) as dag:
+ @task
+ def get_shared_drive_id() -> str:
+ return get_secret(secret_id=GDRIVE_SECRET_ID).strip()
+
+ get_shared_drive_id_task = get_shared_drive_id()
+
@task
def create_connection(connection_id: str):
conn_extra_json = json.dumps(
@@ -108,13 +116,15 @@ with DAG(
source_bucket=BUCKET_NAME,
source_object=FILE_NAME,
destination_object=DRIVE_FILE_NAME,
+ destination_folder_id=get_shared_drive_id_task,
)
# [START detect_file]
detect_file = GoogleDriveFileExistenceSensor(
task_id="detect_file",
- folder_id=FOLDER_ID,
+ folder_id=get_shared_drive_id_task,
file_name=DRIVE_FILE_NAME,
+ drive_id=get_shared_drive_id_task,
gcp_conn_id=CONNECTION_ID,
)
# [END detect_file]
@@ -123,7 +133,8 @@ with DAG(
upload_gdrive_to_gcs = GoogleDriveToGCSOperator(
task_id="upload_gdrive_object_to_gcs",
gcp_conn_id=CONNECTION_ID,
- folder_id=FOLDER_ID,
+ folder_id=get_shared_drive_id_task,
+ drive_id=get_shared_drive_id_task,
file_name=DRIVE_FILE_NAME,
bucket_name=BUCKET_NAME,
object_name=OBJECT,
@@ -131,13 +142,26 @@ with DAG(
# [END upload_gdrive_to_gcs]
@task(trigger_rule=TriggerRule.ALL_DONE)
- def remove_files_from_drive():
+ def remove_files_from_drive(**context):
+ ti = context["ti"]
service = GoogleDriveHook(gcp_conn_id=CONNECTION_ID).get_conn()
- response = service.files().list(q=f"name =
'{DRIVE_FILE_NAME}'").execute()
+ response = (
+ service.files()
+ .list(
+ q=f"name = '{DRIVE_FILE_NAME}'",
+ corpora="drive",
+ driveId=ti.xcom_pull("get_shared_drive_id"),
+ includeItemsFromAllDrives=True,
+ supportsAllDrives=True,
+ )
+ .execute()
+ )
if files := response["files"]:
file = files[0]
- log.info("Deleting file %s...", file)
- service.files().delete(fileId=file["id"])
+ log.info("Trashing file %s...", file["name"])
+ service.files().update(
+ fileId=file["id"], body={"trashed": True},
supportsAllDrives=True
+ ).execute()
log.info("Done.")
remove_files_from_drive_task = remove_files_from_drive()
@@ -153,7 +177,7 @@ with DAG(
delete_connection_task = delete_connection(connection_id=CONNECTION_ID)
(
- [create_bucket >> upload_file, create_connection_task]
+ [get_shared_drive_id_task >> create_bucket >> upload_file,
create_connection_task]
>> copy_single_file
# TEST BODY
>> detect_file
diff --git a/providers/google/tests/unit/google/suite/hooks/test_drive.py
b/providers/google/tests/unit/google/suite/hooks/test_drive.py
index 35eb6102081..51da97fa1b2 100644
--- a/providers/google/tests/unit/google/suite/hooks/test_drive.py
+++ b/providers/google/tests/unit/google/suite/hooks/test_drive.py
@@ -279,7 +279,9 @@ class TestGoogleDriveHook:
file_id = "1eC-Ahi4t57pHcLbW3C_xHB3-YrTQLQBa"
self.gdrive_hook.get_media_request(file_id)
-
mock_get_conn.return_value.files.return_value.get_media.assert_called_once_with(fileId=file_id)
+
mock_get_conn.return_value.files.return_value.get_media.assert_called_once_with(
+ fileId=file_id, supportsAllDrives=True
+ )
@mock.patch("airflow.providers.google.suite.hooks.drive.GoogleDriveHook.get_conn")
def test_get_file_id_when_one_file_exists(self, mock_get_conn):
