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 0801bb24a7501..65d8a9433742c 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 @@ -499,6 +499,11 @@ def _wait_until_complete(self, job, polling_interval: int = 30): raise AirflowException("Something went wrong during waiting of the batch job.") return job + def _validate_results_folder(self): + # Re-checked where the results file is written: deferral resume skips execute(). + 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 _prepare_results_for_xcom(self, job): results = [] if job.dest and job.dest.inlined_responses: @@ -514,6 +519,7 @@ def _prepare_results_for_xcom(self, job): self.log.warning("Error found in the inline result") results.append(inline_response.error) elif job.dest and job.dest.file_name: + 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("/", "-") @@ -541,8 +547,7 @@ def execute(self, context: Context): 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") + self._validate_results_folder() if self.deferrable: self.defer( @@ -944,6 +949,11 @@ def _wait_until_complete(self, job, polling_interval: int = 30): raise AirflowException("Something went wrong during waiting of the batch job: %s", e) return job + def _validate_results_folder(self): + # Re-checked where the results file is written: deferral resume skips execute(). + 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 _prepare_results_for_xcom(self, job): results = [] if job.dest and job.dest.inlined_embed_content_responses: @@ -959,6 +969,7 @@ def _prepare_results_for_xcom(self, job): self.log.warning("Error found in the inline result") results.append(inline_embed_response.error) elif job.dest and job.dest.file_name: + 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("/", "-") @@ -986,8 +997,7 @@ def execute(self, context: Context): 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") + self._validate_results_folder() 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 3e814f9bd9a9e..1071f2de82920 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 @@ -457,6 +457,31 @@ def test_execute_results_folder_not_exists_raises_airflow_exception(self): ): op.execute(context={"ti": mock.MagicMock()}) + @mock.patch(GEN_AI_PATH.format("GenAIGeminiAPIHook")) + def test_prepare_results_for_xcom_results_folder_not_exists_raises_airflow_exception(self, mock_hook): + 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=TEST_FILE_PATH, + ) + mock_job = mock.MagicMock() + mock_job.dest.inlined_responses = None + mock_job.dest.file_name = "results-file" + + with pytest.raises( + AirflowException, + match="path to results_folder does not exist, please provide correct path", + ): + op._prepare_results_for_xcom(mock_job) + + mock_hook.return_value.download_file.assert_not_called() + @mock.patch(GEN_AI_PATH.format("GenAIGeminiAPIHook")) def test__wait_until_complete_exception_raises_airflow_exception(self, mock_hook): op = GenAIGeminiCreateBatchJobOperator( @@ -832,6 +857,31 @@ def test_execute_results_folder_not_exists_raises_airflow_exception(self): ): op.execute(context={"ti": mock.MagicMock()}) + @mock.patch(GEN_AI_PATH.format("GenAIGeminiAPIHook")) + def test_prepare_results_for_xcom_results_folder_not_exists_raises_airflow_exception(self, mock_hook): + 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=TEST_FILE_PATH, + ) + mock_job = mock.MagicMock() + mock_job.dest.inlined_embed_content_responses = None + mock_job.dest.file_name = "results-file" + + with pytest.raises( + AirflowException, + match="path to results_folder does not exist, please provide correct path", + ): + op._prepare_results_for_xcom(mock_job) + + mock_hook.return_value.download_file.assert_not_called() + @mock.patch(GEN_AI_PATH.format("GenAIGeminiAPIHook")) def test__wait_until_complete_exception_raises_airflow_exception(self, mock_hook): op = GenAIGeminiCreateEmbeddingsBatchJobOperator(