Validate Gemini batch results_folder in the worker that writes results

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(

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(