Add a filter for local files in GoogleDisplayVideo360CreateQueryOpera…

…tor (#35635)
apache · Nov 14, 2023 · 9d5f726 · 9d5f726
1 parent eeae875
commit 9d5f726
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 2 deletions.
diff --git a/airflow/providers/google/marketing_platform/operators/display_video.py b/airflow/providers/google/marketing_platform/operators/display_video.py
@@ -285,13 +285,15 @@ def execute(self, context: Context):
 
  # 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
@@ -23,6 +23,7 @@
 
 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 @@ def teardown_method(self):
  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 @@ def test_execute(
  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 @@ def test_execute(
  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,