This is an automated email from the ASF dual-hosted git repository.
shahar1 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 a1c27c5d36c Prevent premature validation in Gemini batch job operators
(#70362)
a1c27c5d36c is described below
commit a1c27c5d36c5429303461b4a71f373307535f8b6
Author: daniel.jin <[email protected]>
AuthorDate: Fri Jul 24 04:08:00 2026 -0400
Prevent premature validation in Gemini batch job operators (#70362)
---
.../providers/google/cloud/operators/gen_ai.py | 39 ++--
.../unit/google/cloud/operators/test_gen_ai.py | 210 ++++++++++++---------
.../ci/prek/validate_operators_init_exemptions.txt | 2 -
3 files changed, 143 insertions(+), 108 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 71a51c32419..0801bb24a75 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
@@ -475,15 +475,6 @@ class
GenAIGeminiCreateBatchJobOperator(GoogleCloudBaseOperator):
self.results_folder = results_folder
self.deferrable = deferrable
- if self.retrieve_result and not (self.wait_until_complete or
self.deferrable):
- raise AirflowException(
- "Retrieving results is possible only if wait_until_complete
set to True or in deferrable mode"
- )
- if self.results_folder and not isinstance(self.input_source, str):
- raise AirflowException("results_folder works only when
input_source is file name")
- if self.results_folder and not
os.path.exists(os.path.abspath(self.results_folder)):
- raise AirflowException("path to results_folder does not exist,
please provide correct path")
-
def _wait_until_complete(self, job, polling_interval: int = 30):
try:
while True:
@@ -542,6 +533,17 @@ class
GenAIGeminiCreateBatchJobOperator(GoogleCloudBaseOperator):
)
def execute(self, context: Context):
+ if self.retrieve_result and not (self.wait_until_complete or
self.deferrable):
+ raise AirflowException(
+ "Retrieving results is possible only if wait_until_complete
set to True or in deferrable mode"
+ )
+
+ if self.results_folder and not isinstance(self.input_source, str):
+ raise AirflowException("results_folder works only when
input_source is file name")
+
+ if self.results_folder and not
os.path.exists(os.path.abspath(self.results_folder)):
+ raise AirflowException("path to results_folder does not exist,
please provide correct path")
+
if self.deferrable:
self.defer(
trigger=GenAIGeminiCreateBatchJobTrigger(
@@ -918,15 +920,6 @@ class
GenAIGeminiCreateEmbeddingsBatchJobOperator(GoogleCloudBaseOperator):
self.results_folder = results_folder
self.deferrable = deferrable
- if self.retrieve_result and not (self.wait_until_complete or
self.deferrable):
- raise AirflowException(
- "Retrieving results is possible only if wait_until_complete
set to True or in deferrable mode"
- )
- if self.results_folder and not isinstance(self.input_source, str):
- raise AirflowException("results_folder works only when
input_source is file name")
- if self.results_folder and not
os.path.exists(os.path.abspath(self.results_folder)):
- raise AirflowException("path to results_folder does not exist,
please provide correct path")
-
def _wait_until_complete(self, job, polling_interval: int = 30):
try:
while True:
@@ -985,6 +978,16 @@ class
GenAIGeminiCreateEmbeddingsBatchJobOperator(GoogleCloudBaseOperator):
)
def execute(self, context: Context):
+ if self.retrieve_result and not (self.wait_until_complete or
self.deferrable):
+ raise AirflowException(
+ "Retrieving results is possible only if wait_until_complete
set to True or in deferrable mode"
+ )
+
+ if self.results_folder and not isinstance(self.input_source, str):
+ raise AirflowException("results_folder works only when
input_source is file name")
+
+ if self.results_folder and not
os.path.exists(os.path.abspath(self.results_folder)):
+ raise AirflowException("path to results_folder does not exist,
please provide correct path")
if self.deferrable:
self.defer(
trigger=GenAIGeminiCreateEmbeddingsBatchJobTrigger(
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 cbd2105ef82..3e814f9bd9a 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
@@ -395,50 +395,67 @@ class TestGenAIGeminiCreateBatchJobOperator:
mock_hook.return_value.get_batch_job.assert_called_once_with("test-name")
mock_job.model_dump.assert_called_once_with(mode="json")
- def
test_init_retrieve_result_and_not_wait_until_complete_raises_airflow_exception(self):
- with pytest.raises(AirflowException):
- 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_BATCH_JOB_INLINED_REQUESTS,
- gemini_api_key=TEST_GEMINI_API_KEY,
- wait_until_complete=False,
- retrieve_result=True,
- )
-
- def test_init_input_source_not_string_raises_airflow_exception(self):
- with pytest.raises(AirflowException):
- 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_BATCH_JOB_INLINED_REQUESTS,
- gemini_api_key=TEST_GEMINI_API_KEY,
- wait_until_complete=False,
- results_folder=TEST_FILE_PATH,
- )
-
- def test_init_results_folder_not_exists_raises_airflow_exception(self):
- with pytest.raises(AirflowException):
- 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,
- wait_until_complete=False,
- results_folder=TEST_FILE_PATH,
- )
+ def
test_execute_retrieve_result_and_not_wait_until_complete_raises_airflow_exception(self):
+ 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_BATCH_JOB_INLINED_REQUESTS,
+ gemini_api_key=TEST_GEMINI_API_KEY,
+ wait_until_complete=False,
+ retrieve_result=True,
+ )
+
+ with pytest.raises(
+ AirflowException,
+ match=(
+ "Retrieving results is possible only if wait_until_complete
set to True or in deferrable mode"
+ ),
+ ):
+ op.execute(context={"ti": mock.MagicMock()})
+
+ def test_execute_input_source_not_string_raises_airflow_exception(self):
+ 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_BATCH_JOB_INLINED_REQUESTS,
+ gemini_api_key=TEST_GEMINI_API_KEY,
+ wait_until_complete=False,
+ results_folder=TEST_FILE_PATH,
+ )
+
+ with pytest.raises(
+ AirflowException,
+ match="results_folder works only when input_source is file name",
+ ):
+ op.execute(context={"ti": mock.MagicMock()})
+
+ def test_execute_results_folder_not_exists_raises_airflow_exception(self):
+ 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,
+ wait_until_complete=False,
+ results_folder=TEST_FILE_PATH,
+ )
+
+ with pytest.raises(
+ AirflowException,
+ match="path to results_folder does not exist, please provide
correct path",
+ ):
+ op.execute(context={"ti": mock.MagicMock()})
@mock.patch(GEN_AI_PATH.format("GenAIGeminiAPIHook"))
def test__wait_until_complete_exception_raises_airflow_exception(self,
mock_hook):
@@ -753,50 +770,67 @@ class TestGenAIGeminiCreateEmbeddingsBatchJobOperator:
create_embeddings_config=None,
)
- def
test_init_retrieve_result_and_not_wait_until_complete_raises_airflow_exception(self):
- with pytest.raises(AirflowException):
- GenAIGeminiCreateEmbeddingsBatchJobOperator(
- task_id=TASK_ID,
- project_id=GCP_PROJECT,
- location=GCP_LOCATION,
- input_source=TEST_EMBEDDINGS_JOB_INLINED_REQUESTS,
- model=EMBEDDING_MODEL,
- gemini_api_key=TEST_GEMINI_API_KEY,
- gcp_conn_id=GCP_CONN_ID,
- impersonation_chain=IMPERSONATION_CHAIN,
- wait_until_complete=False,
- retrieve_result=True,
- )
-
- def test_init_input_source_not_string_raises_airflow_exception(self):
- with pytest.raises(AirflowException):
- GenAIGeminiCreateEmbeddingsBatchJobOperator(
- task_id=TASK_ID,
- project_id=GCP_PROJECT,
- location=GCP_LOCATION,
- input_source=TEST_EMBEDDINGS_JOB_INLINED_REQUESTS,
- model=EMBEDDING_MODEL,
- gemini_api_key=TEST_GEMINI_API_KEY,
- gcp_conn_id=GCP_CONN_ID,
- impersonation_chain=IMPERSONATION_CHAIN,
- wait_until_complete=False,
- results_folder=TEST_FILE_PATH,
- )
-
- def test_init_results_folder_not_exists_raises_airflow_exception(self):
- with pytest.raises(AirflowException):
- 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,
- wait_until_complete=False,
- results_folder=TEST_FILE_PATH,
- )
+ def
test_execute_retrieve_result_and_not_wait_until_complete_raises_airflow_exception(self):
+ op = GenAIGeminiCreateEmbeddingsBatchJobOperator(
+ task_id=TASK_ID,
+ project_id=GCP_PROJECT,
+ location=GCP_LOCATION,
+ input_source=TEST_EMBEDDINGS_JOB_INLINED_REQUESTS,
+ model=EMBEDDING_MODEL,
+ gemini_api_key=TEST_GEMINI_API_KEY,
+ gcp_conn_id=GCP_CONN_ID,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ wait_until_complete=False,
+ retrieve_result=True,
+ )
+
+ with pytest.raises(
+ AirflowException,
+ match=(
+ "Retrieving results is possible only if wait_until_complete
set to True or in deferrable mode"
+ ),
+ ):
+ op.execute(context={"ti": mock.MagicMock()})
+
+ def test_execute_input_source_not_string_raises_airflow_exception(self):
+ op = GenAIGeminiCreateEmbeddingsBatchJobOperator(
+ task_id=TASK_ID,
+ project_id=GCP_PROJECT,
+ location=GCP_LOCATION,
+ input_source=TEST_EMBEDDINGS_JOB_INLINED_REQUESTS,
+ model=EMBEDDING_MODEL,
+ gemini_api_key=TEST_GEMINI_API_KEY,
+ gcp_conn_id=GCP_CONN_ID,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ wait_until_complete=False,
+ results_folder=TEST_FILE_PATH,
+ )
+
+ with pytest.raises(
+ AirflowException,
+ match="results_folder works only when input_source is file name",
+ ):
+ op.execute(context={"ti": mock.MagicMock()})
+
+ def test_execute_results_folder_not_exists_raises_airflow_exception(self):
+ 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,
+ wait_until_complete=False,
+ results_folder=TEST_FILE_PATH,
+ )
+
+ with pytest.raises(
+ AirflowException,
+ match="path to results_folder does not exist, please provide
correct path",
+ ):
+ op.execute(context={"ti": mock.MagicMock()})
@mock.patch(GEN_AI_PATH.format("GenAIGeminiAPIHook"))
def test__wait_until_complete_exception_raises_airflow_exception(self,
mock_hook):
diff --git a/scripts/ci/prek/validate_operators_init_exemptions.txt
b/scripts/ci/prek/validate_operators_init_exemptions.txt
index 052da40bfaa..2f3a91c6b04 100644
--- a/scripts/ci/prek/validate_operators_init_exemptions.txt
+++ b/scripts/ci/prek/validate_operators_init_exemptions.txt
@@ -53,8 +53,6 @@
providers/google/src/airflow/providers/google/cloud/operators/functions.py::Clou
providers/google/src/airflow/providers/google/cloud/operators/gcs.py::GCSDeleteObjectsOperator
providers/google/src/airflow/providers/google/cloud/operators/gcs.py::GCSFileTransformOperator
providers/google/src/airflow/providers/google/cloud/operators/gcs.py::GCSListObjectsOperator
-providers/google/src/airflow/providers/google/cloud/operators/gen_ai.py::GenAIGeminiCreateBatchJobOperator
-providers/google/src/airflow/providers/google/cloud/operators/gen_ai.py::GenAIGeminiCreateEmbeddingsBatchJobOperator
providers/google/src/airflow/providers/google/cloud/sensors/bigquery_dts.py::BigQueryDataTransferServiceTransferRunSensor
providers/google/src/airflow/providers/google/cloud/sensors/cloud_composer.py::CloudComposerExternalTaskSensor
providers/google/src/airflow/providers/google/cloud/transfers/azure_fileshare_to_gcs.py::AzureFileShareToGCSOperator
