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 1a9091e4c5f contain gemini batch-job results file within
results_folder (#71832)
1a9091e4c5f is described below
commit 1a9091e4c5fcf2d8ee95c7e4f485295a50061849
Author: Samina <[email protected]>
AuthorDate: Mon Sep 21 04:47:29 2026 +0530
contain gemini batch-job results file within results_folder (#71832)
* contain gemini batch-job results file within results_folder
* Resolve symlinks in gemini batch-job results containment check
* Add symlink-escape regression test for gemini batch-job results path
* Name the results-path helper for what it does and cover the happy path
Rename _results_file_path to _build_results_file_path per the action-verb
naming rule, and annotate the job parameter as BatchJob (already imported).
Add one happy-path test per operator. The four existing tests are all
rejection cases, and the file-writing branch had no other coverage, so a
helper that raised unconditionally would have passed the whole suite.
Verified: sabotaging the helper to always raise fails exactly these two.
Generated-by: Claude Opus 5
* Handle a batch job with neither display_name nor name
Both fields are optional on BatchJob, so job.name.replace(...) could raise
AttributeError when neither is set. Annotating the parameter as BatchJob
made
mypy surface it; fail with a clear message instead.
Generated-by: Claude Opus 5
---------
Co-authored-by: Samin061 <[email protected]>
---
.../providers/google/cloud/operators/gen_ai.py | 29 +++-
.../unit/google/cloud/operators/test_gen_ai.py | 162 +++++++++++++++++++++
2 files changed, 187 insertions(+), 4 deletions(-)
diff --git
a/providers/google/src/airflow/providers/google/cloud/operators/gen_ai.py
b/providers/google/src/airflow/providers/google/cloud/operators/gen_ai.py
index 65d8a943374..da8235e07b5 100644
--- a/providers/google/src/airflow/providers/google/cloud/operators/gen_ai.py
+++ b/providers/google/src/airflow/providers/google/cloud/operators/gen_ai.py
@@ -22,6 +22,7 @@ from __future__ import annotations
import os.path
import time
from collections.abc import Sequence
+from pathlib import Path
from typing import TYPE_CHECKING, Any
from google.genai.errors import ClientError
@@ -57,6 +58,28 @@ if TYPE_CHECKING:
from airflow.providers.common.compat.sdk import Context
+def _build_results_file_path(results_folder: str | None, job: BatchJob) -> str:
+ """
+ Build the local results-file path for a batch job, refusing one that
escapes ``results_folder``.
+
+ ``display_name`` and ``name`` come from the batch-job metadata returned by
the Gemini
+ API. ``job.name`` is already sanitized against ``/``, but ``display_name``
(used first)
+ is not, so a ``..`` in it would place the file outside ``results_folder``
(CWE-22).
+ """
+ # Both are optional on BatchJob, so neither being set leaves nothing to
name the file after.
+ file_name = job.display_name or (job.name or "").replace("/", "-")
+ if not file_name:
+ raise ValueError("Batch job has neither display_name nor name; cannot
derive a results file name.")
+ base = Path(f"{results_folder}").resolve()
+ path_to_file = (base / f"{file_name}.jsonl").resolve()
+ if not path_to_file.is_relative_to(base):
+ raise ValueError(
+ f"Refusing to write batch-job results outside results_folder
{results_folder!r}: "
+ f"file name {file_name!r} resolves to {path_to_file!r}."
+ )
+ return str(path_to_file)
+
+
class GenAIGenerateEmbeddingsOperator(GoogleCloudBaseOperator):
"""
Uses the Gemini AI Embeddings API to generate embeddings for words,
phrases, sentences, and code.
@@ -522,8 +545,7 @@ class
GenAIGeminiCreateBatchJobOperator(GoogleCloudBaseOperator):
self._validate_results_folder()
file_content_bytes =
self.hook.download_file(file_name=job.dest.file_name)
file_content = file_content_bytes.decode("utf-8")
- file_name = job.display_name or job.name.replace("/", "-")
- path_to_file =
os.path.abspath(f"{self.results_folder}/{file_name}.jsonl")
+ path_to_file = _build_results_file_path(self.results_folder, job)
with open(path_to_file, "w") as file_with_results:
file_with_results.writelines(file_content.splitlines(True))
results = path_to_file
@@ -972,8 +994,7 @@ class
GenAIGeminiCreateEmbeddingsBatchJobOperator(GoogleCloudBaseOperator):
self._validate_results_folder()
file_content_bytes =
self.hook.download_file(file_name=job.dest.file_name)
file_content = file_content_bytes.decode("utf-8")
- file_name = job.display_name or job.name.replace("/", "-")
- path_to_file =
os.path.abspath(f"{self.results_folder}/{file_name}.jsonl")
+ path_to_file = _build_results_file_path(self.results_folder, job)
with open(path_to_file, "w") as file_with_results:
file_with_results.writelines(file_content.splitlines(True))
results = path_to_file
diff --git a/providers/google/tests/unit/google/cloud/operators/test_gen_ai.py
b/providers/google/tests/unit/google/cloud/operators/test_gen_ai.py
index 1071f2de829..870d4f7e66a 100644
--- a/providers/google/tests/unit/google/cloud/operators/test_gen_ai.py
+++ b/providers/google/tests/unit/google/cloud/operators/test_gen_ai.py
@@ -482,6 +482,87 @@ class TestGenAIGeminiCreateBatchJobOperator:
mock_hook.return_value.download_file.assert_not_called()
+ @mock.patch(GEN_AI_PATH.format("GenAIGeminiAPIHook"))
+ def test_prepare_results_for_xcom_writes_inside_results_folder(self,
mock_hook, tmp_path):
+ results_folder = tmp_path / "results"
+ results_folder.mkdir()
+ op = GenAIGeminiCreateBatchJobOperator(
+ task_id=TASK_ID,
+ project_id=GCP_PROJECT,
+ location=GCP_LOCATION,
+ model=TEST_GEMINI_MODEL,
+ gcp_conn_id=GCP_CONN_ID,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ input_source=TEST_FILE_NAME,
+ gemini_api_key=TEST_GEMINI_API_KEY,
+ results_folder=str(results_folder),
+ )
+ mock_hook.return_value.download_file.return_value = b"data"
+ mock_job = mock.MagicMock()
+ mock_job.dest.inlined_responses = None
+ mock_job.dest.file_name = "results-file"
+ mock_job.display_name = "ok"
+
+ results = op._prepare_results_for_xcom(mock_job)
+
+ assert results == str(results_folder.resolve() / "ok.jsonl")
+ assert (results_folder / "ok.jsonl").read_text() == "data"
+
+ @mock.patch(GEN_AI_PATH.format("GenAIGeminiAPIHook"))
+ def
test_prepare_results_for_xcom_rejects_display_name_path_traversal(self,
mock_hook, tmp_path):
+ op = GenAIGeminiCreateBatchJobOperator(
+ task_id=TASK_ID,
+ project_id=GCP_PROJECT,
+ location=GCP_LOCATION,
+ model=TEST_GEMINI_MODEL,
+ gcp_conn_id=GCP_CONN_ID,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ input_source=TEST_FILE_NAME,
+ gemini_api_key=TEST_GEMINI_API_KEY,
+ results_folder=str(tmp_path / "results"),
+ )
+ (tmp_path / "results").mkdir()
+ mock_hook.return_value.download_file.return_value = b"data"
+ mock_job = mock.MagicMock()
+ mock_job.dest.inlined_responses = None
+ mock_job.dest.file_name = "results-file"
+ mock_job.display_name = "../evil"
+
+ with pytest.raises(ValueError, match="Refusing to write batch-job
results outside"):
+ op._prepare_results_for_xcom(mock_job)
+
+ assert not (tmp_path / "evil.jsonl").exists()
+
+ @mock.patch(GEN_AI_PATH.format("GenAIGeminiAPIHook"))
+ def
test_prepare_results_for_xcom_rejects_display_name_symlink_escape(self,
mock_hook, tmp_path):
+ results_folder = tmp_path / "results"
+ results_folder.mkdir()
+ outside = tmp_path / "outside"
+ outside.mkdir()
+ (results_folder / "link").symlink_to(outside, target_is_directory=True)
+
+ op = GenAIGeminiCreateBatchJobOperator(
+ task_id=TASK_ID,
+ project_id=GCP_PROJECT,
+ location=GCP_LOCATION,
+ model=TEST_GEMINI_MODEL,
+ gcp_conn_id=GCP_CONN_ID,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ input_source=TEST_FILE_NAME,
+ gemini_api_key=TEST_GEMINI_API_KEY,
+ results_folder=str(results_folder),
+ )
+ mock_hook.return_value.download_file.return_value = b"data"
+ mock_job = mock.MagicMock()
+ mock_job.dest.inlined_responses = None
+ mock_job.dest.file_name = "results-file"
+ mock_job.display_name = "link/evil"
+
+ with pytest.raises(ValueError, match="Refusing to write batch-job
results outside"):
+ op._prepare_results_for_xcom(mock_job)
+
+ assert not (outside / "evil.jsonl").exists()
+
@mock.patch(GEN_AI_PATH.format("GenAIGeminiAPIHook"))
def test__wait_until_complete_exception_raises_airflow_exception(self,
mock_hook):
op = GenAIGeminiCreateBatchJobOperator(
@@ -882,6 +963,87 @@ class TestGenAIGeminiCreateEmbeddingsBatchJobOperator:
mock_hook.return_value.download_file.assert_not_called()
+ @mock.patch(GEN_AI_PATH.format("GenAIGeminiAPIHook"))
+ def test_prepare_results_for_xcom_writes_inside_results_folder(self,
mock_hook, tmp_path):
+ results_folder = tmp_path / "results"
+ results_folder.mkdir()
+ op = GenAIGeminiCreateEmbeddingsBatchJobOperator(
+ task_id=TASK_ID,
+ project_id=GCP_PROJECT,
+ location=GCP_LOCATION,
+ input_source=TEST_FILE_NAME,
+ model=EMBEDDING_MODEL,
+ gemini_api_key=TEST_GEMINI_API_KEY,
+ gcp_conn_id=GCP_CONN_ID,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ results_folder=str(results_folder),
+ )
+ mock_hook.return_value.download_file.return_value = b"data"
+ mock_job = mock.MagicMock()
+ mock_job.dest.inlined_embed_content_responses = None
+ mock_job.dest.file_name = "results-file"
+ mock_job.display_name = "ok"
+
+ results = op._prepare_results_for_xcom(mock_job)
+
+ assert results == str(results_folder.resolve() / "ok.jsonl")
+ assert (results_folder / "ok.jsonl").read_text() == "data"
+
+ @mock.patch(GEN_AI_PATH.format("GenAIGeminiAPIHook"))
+ def
test_prepare_results_for_xcom_rejects_display_name_path_traversal(self,
mock_hook, tmp_path):
+ op = GenAIGeminiCreateEmbeddingsBatchJobOperator(
+ task_id=TASK_ID,
+ project_id=GCP_PROJECT,
+ location=GCP_LOCATION,
+ input_source=TEST_FILE_NAME,
+ model=EMBEDDING_MODEL,
+ gemini_api_key=TEST_GEMINI_API_KEY,
+ gcp_conn_id=GCP_CONN_ID,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ results_folder=str(tmp_path / "results"),
+ )
+ (tmp_path / "results").mkdir()
+ mock_hook.return_value.download_file.return_value = b"data"
+ mock_job = mock.MagicMock()
+ mock_job.dest.inlined_embed_content_responses = None
+ mock_job.dest.file_name = "results-file"
+ mock_job.display_name = "../evil"
+
+ with pytest.raises(ValueError, match="Refusing to write batch-job
results outside"):
+ op._prepare_results_for_xcom(mock_job)
+
+ assert not (tmp_path / "evil.jsonl").exists()
+
+ @mock.patch(GEN_AI_PATH.format("GenAIGeminiAPIHook"))
+ def
test_prepare_results_for_xcom_rejects_display_name_symlink_escape(self,
mock_hook, tmp_path):
+ results_folder = tmp_path / "results"
+ results_folder.mkdir()
+ outside = tmp_path / "outside"
+ outside.mkdir()
+ (results_folder / "link").symlink_to(outside, target_is_directory=True)
+
+ op = GenAIGeminiCreateEmbeddingsBatchJobOperator(
+ task_id=TASK_ID,
+ project_id=GCP_PROJECT,
+ location=GCP_LOCATION,
+ input_source=TEST_FILE_NAME,
+ model=EMBEDDING_MODEL,
+ gemini_api_key=TEST_GEMINI_API_KEY,
+ gcp_conn_id=GCP_CONN_ID,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ results_folder=str(results_folder),
+ )
+ mock_hook.return_value.download_file.return_value = b"data"
+ mock_job = mock.MagicMock()
+ mock_job.dest.inlined_embed_content_responses = None
+ mock_job.dest.file_name = "results-file"
+ mock_job.display_name = "link/evil"
+
+ with pytest.raises(ValueError, match="Refusing to write batch-job
results outside"):
+ op._prepare_results_for_xcom(mock_job)
+
+ assert not (outside / "evil.jsonl").exists()
+
@mock.patch(GEN_AI_PATH.format("GenAIGeminiAPIHook"))
def test__wait_until_complete_exception_raises_airflow_exception(self,
mock_hook):
op = GenAIGeminiCreateEmbeddingsBatchJobOperator(
