This is an automated email from the ASF dual-hosted git repository.
husseinawala 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 9d5f726c8a Add a filter for local files in
GoogleDisplayVideo360CreateQueryOperator (#35635)
9d5f726c8a is described below
commit 9d5f726c8a5705f1b0dec09ff54fe3fea12d4a5f
Author: Hussein Awala <[email protected]>
AuthorDate: Wed Nov 15 01:10:25 2023 +0200
Add a filter for local files in GoogleDisplayVideo360CreateQueryOperator
(#35635)
---
.../google/marketing_platform/operators/display_video.py | 4 +++-
.../marketing_platform/operators/test_display_video.py | 12 +++++++++++-
2 files changed, 14 insertions(+), 2 deletions(-)
diff --git
a/airflow/providers/google/marketing_platform/operators/display_video.py
b/airflow/providers/google/marketing_platform/operators/display_video.py
index 003ad4723c..2df18939b8 100644
--- a/airflow/providers/google/marketing_platform/operators/display_video.py
+++ b/airflow/providers/google/marketing_platform/operators/display_video.py
@@ -285,13 +285,15 @@ class
GoogleDisplayVideo360DownloadReportV2Operator(BaseOperator):
# If no custom report_name provided, use DV360 name
file_url = resource["metadata"]["googleCloudStoragePath"]
+ if urllib.parse.urlparse(file_url).scheme == "file":
+ raise AirflowException("Accessing local file is not allowed in
this operator")
report_name = self.report_name or
urlsplit(file_url).path.split("/")[-1]
report_name = self._resolve_file_name(report_name)
# Download the report
self.log.info("Starting downloading report %s", self.report_id)
with tempfile.NamedTemporaryFile(delete=False) as temp_file:
- with urllib.request.urlopen(file_url) as response:
+ with urllib.request.urlopen(file_url) as response: # nosec
shutil.copyfileobj(response, temp_file, length=self.chunk_size)
temp_file.flush()
diff --git
a/tests/providers/google/marketing_platform/operators/test_display_video.py
b/tests/providers/google/marketing_platform/operators/test_display_video.py
index 66b3ddd809..9eb109f9a5 100644
--- a/tests/providers/google/marketing_platform/operators/test_display_video.py
+++ b/tests/providers/google/marketing_platform/operators/test_display_video.py
@@ -23,6 +23,7 @@ from unittest import mock
import pytest
+from airflow.exceptions import AirflowException
from airflow.models import DAG, TaskInstance as TI
from airflow.providers.google.marketing_platform.operators.display_video
import (
GoogleDisplayVideo360CreateQueryOperator,
@@ -78,6 +79,9 @@ class TestGoogleDisplayVideo360DownloadReportV2Operator:
with create_session() as session:
session.query(TI).delete()
+ @pytest.mark.parametrize(
+ "file_path, should_except", [("https://host/path";, False),
("file:/path/to/file", True)]
+ )
@mock.patch("airflow.providers.google.marketing_platform.operators.display_video.shutil")
@mock.patch("airflow.providers.google.marketing_platform.operators.display_video.urllib.request")
@mock.patch("airflow.providers.google.marketing_platform.operators.display_video.tempfile")
@@ -97,12 +101,14 @@ class TestGoogleDisplayVideo360DownloadReportV2Operator:
mock_temp,
mock_request,
mock_shutil,
+ file_path,
+ should_except,
):
mock_temp.NamedTemporaryFile.return_value.__enter__.return_value.name
= FILENAME
mock_hook.return_value.get_report.return_value = {
"metadata": {
"status": {"state": "DONE"},
- "googleCloudStoragePath": "TEST",
+ "googleCloudStoragePath": file_path,
}
}
op = GoogleDisplayVideo360DownloadReportV2Operator(
@@ -112,6 +118,10 @@ class TestGoogleDisplayVideo360DownloadReportV2Operator:
report_name=REPORT_NAME,
task_id="test_task",
)
+ if should_except:
+ with pytest.raises(AirflowException):
+ op.execute(context=None)
+ return
op.execute(context=None)
mock_hook.assert_called_once_with(
gcp_conn_id=GCP_CONN_ID,
