shahar1 commented on code in PR #61284:
URL: https://github.com/apache/airflow/pull/61284#discussion_r2749377017
##########
providers/google/src/airflow/providers/google/cloud/transfers/calendar_to_gcs.py:
##########
@@ -186,6 +189,9 @@ def execute(self, context):
time_zone=self.time_zone,
updated_min=self.updated_min,
)
- gcs_path_to_file = self._upload_data(events)
-
- return gcs_path_to_file
+ gcs_uri = self._upload_data(events)
+ result = [gcs_uri]
+
+ if self.unwrap_single:
+ return gcs_uri
+ return result
Review Comment:
For a slightly better consistency:
```suggestion
if self.unwrap_single:
return result[0]
return result
```
##########
providers/google/tests/unit/google/cloud/transfers/test_calendar_to_gcs.py:
##########
@@ -124,3 +163,53 @@ def test_execute(self, mock_upload_data,
mock_calendar_hook):
)
mock_upload_data.assert_called_once_with(data)
+
+
@mock.patch("airflow.providers.google.cloud.transfers.calendar_to_gcs.GoogleCalendarHook")
+ @mock.patch(
+
"airflow.providers.google.cloud.transfers.calendar_to_gcs.GoogleCalendarToGCSOperator._upload_data"
+ )
+ def test_execute_with_unwrap_single_false(self, mock_upload_data,
mock_calendar_hook):
+ context = {}
+ data = [EVENT]
+ mock_upload_data.side_effect = [f"gs://{BUCKET}/{PATH}"]
+ mock_calendar_hook.return_value.get_events.return_value = data
+
+ op = GoogleCalendarToGCSOperator(
+ api_version=API_VERSION,
+ task_id="test_task",
+ calendar_id=CALENDAR_ID,
+ destination_bucket=BUCKET,
+ destination_path=PATH,
+ gcp_conn_id=GCP_CONN_ID,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ unwrap_single=False,
+ )
+ result = op.execute(context)
+
+ mock_calendar_hook.assert_called_once_with(
+ api_version=API_VERSION,
+ gcp_conn_id=GCP_CONN_ID,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ )
+
+ mock_calendar_hook.return_value.get_events.assert_called_once_with(
+ calendar_id=CALENDAR_ID,
+ i_cal_uid=None,
+ max_attendees=None,
+ max_results=None,
+ order_by=None,
+ private_extended_property=None,
+ q=None,
+ shared_extended_property=None,
+ show_deleted=None,
+ show_hidden_invitation=None,
+ single_events=None,
+ sync_token=None,
+ time_max=None,
+ time_min=None,
+ time_zone=None,
+ updated_min=None,
+ )
+
+ mock_upload_data.assert_called_once_with(data)
+ assert result == [f"gs://{BUCKET}/{PATH}"]
Review Comment:
Add a new line
##########
providers/google/src/airflow/providers/google/marketing_platform/operators/display_video.py:
##########
@@ -133,6 +133,7 @@ class GoogleDisplayVideo360SDFtoGCSOperator(BaseOperator):
"bucket_name",
"object_name",
"impersonation_chain",
+ "unwrap_single",
Review Comment:
Same comments as before regarding removal from templated fields + docstring
##########
providers/google/src/airflow/providers/google/cloud/transfers/calendar_to_gcs.py:
##########
@@ -82,6 +82,7 @@ class GoogleCalendarToGCSOperator(BaseOperator):
"destination_bucket",
"destination_path",
"impersonation_chain",
+ "unwrap_single",
Review Comment:
I prefer not to add `unwrap_single` to the `template_fields` at the moment,
as I don't see a clear use case for that + we might decide to eventually
deprecate the flag in the future after making behavior consistent.
##########
providers/google/src/airflow/providers/google/cloud/transfers/calendar_to_gcs.py:
##########
Review Comment:
Please add a docstring for the `unwrap_single ` flag
##########
providers/google/tests/unit/google/marketing_platform/operators/test_display_video.py:
##########
@@ -144,4 +214,4 @@ def test_execute(self, mock_hook):
mock_hook.return_value.create_sdf_download_operation.assert_called_once_with(
body_request=body_request
)
-
mock_context["task_instance"].xcom_push.assert_called_once_with(key="name",
value=test_name)
+
mock_context["task_instance"].xcom_push.assert_called_once_with(key="name",
value=test_name)
Review Comment:
Add a new line
##########
providers/google/src/airflow/providers/google/marketing_platform/operators/display_video.py:
##########
@@ -155,8 +157,9 @@ def __init__(
self.api_version = api_version
self.gcp_conn_id = gcp_conn_id
self.impersonation_chain = impersonation_chain
+ self.unwrap_single = unwrap_single
- def execute(self, context: Context) -> str:
+ def execute(self, context: Context) -> Union[str, list[str]]:
Review Comment:
```suggestion
def execute(self, context: Context) -> str | list[str]:
```
##########
providers/google/tests/unit/google/marketing_platform/operators/test_display_video.py:
##########
@@ -145,3 +215,5 @@ def test_execute(self, mock_hook):
body_request=body_request
)
mock_context["task_instance"].xcom_push.assert_called_once_with(key="name",
value=test_name)
+
+
Review Comment:
Remove one empty line
##########
providers/google/src/airflow/providers/google/marketing_platform/operators/display_video.py:
##########
@@ -194,4 +197,8 @@ def execute(self, context: Context) -> str:
filename=os.path.join(tmp_dir, fname),
gzip=False,
)
- return f"{self.bucket_name}/{self.object_name}"
+ result = [f"gs://{self.bucket_name}/{self.object_name}"]
+
+ if self.unwrap_single:
+ return f"gs://{self.bucket_name}/{self.object_name}"
+ return result
Review Comment:
For a slightly better consistency:
```suggestion
result = [f"gs://{self.bucket_name}/{self.object_name}"]
if self.unwrap_single:
return result[0]
return result
```
##########
providers/google/src/airflow/providers/google/cloud/transfers/calendar_to_gcs.py:
##########
@@ -108,6 +109,7 @@ def __init__(
destination_path: str | None = None,
gcp_conn_id: str = "google_cloud_default",
impersonation_chain: str | Sequence[str] | None = None,
+ unwrap_single: bool = True,
Review Comment:
In the future we plan to change the default to `False`, which will be a
breaking change if you don't set it up explicitly.
For that reason, please let the value also accept `None` as well, and if it
is `None` (meaning that it is unset) - set `unwrap_single` it to `True` +
raise a `FutureWarning` about the expected change.
If it is explicitly set (either `True` or `False`), there's no need to raise
this warning.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
