Skip to content

Prevent premature validation in Gemini batch job operators #70362

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
shahar1 merged 1 commit into from
Jul 24, 2026
Merged

Prevent premature validation in Gemini batch job operators #70362

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
39 changes: 21 additions & 18 deletions providers/google/src/airflow/providers/google/cloud/operators/gen_ai.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -475,15 +475,6 @@ def __init__(
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:
Expand Down Expand Up @@ -542,6 +533,17 @@ def hook(self):
)

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(
Expand Down Expand Up @@ -918,15 +920,6 @@ def __init__(
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:
Expand Down Expand Up @@ -985,6 +978,16 @@ def hook(self):
)

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(
Expand Down
210 changes: 122 additions & 88 deletions providers/google/tests/unit/google/cloud/operators/test_gen_ai.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -395,50 +395,67 @@ def test_execute_complete_return_value(self, mock_hook):
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):
Expand Down Expand Up @@ -753,50 +770,67 @@ def test_execute(self, mock_hook):
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):
Expand Down
2 changes: 0 additions & 2 deletions scripts/ci/prek/validate_operators_init_exemptions.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -60,8 +60,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
Expand Down