diff --git a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/log/wasb_task_handler.py b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/log/wasb_task_handler.py index 32cd3ee374ec1..5c631f4d61248 100644 --- a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/log/wasb_task_handler.py +++ b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/log/wasb_task_handler.py @@ -66,6 +66,8 @@ def upload(self, path: str | os.PathLike, ti: RuntimeTI | None = None) -> None: has_uploaded = self.write(log, remote_loc) if has_uploaded and self.delete_local_copy: shutil.rmtree(os.path.dirname(local_loc)) + elif has_uploaded: + local_loc.write_text("") @cached_property def hook(self): @@ -193,7 +195,9 @@ def write(self, log: str, remote_log_location: str, append: bool = True) -> bool """ if append and self.wasb_log_exists(remote_log_location): old_log = self.wasb_read(remote_log_location) - log = f"{old_log}\n{log}" if old_log else log + if old_log: + sep = "" if old_log.endswith("\n") else "\n" + log = f"{old_log}{sep}{log}" try: self.hook.load_string(log, self.wasb_container, remote_log_location, overwrite=True) diff --git a/providers/microsoft/azure/tests/unit/microsoft/azure/log/test_wasb_task_handler.py b/providers/microsoft/azure/tests/unit/microsoft/azure/log/test_wasb_task_handler.py index 57d21664c131f..349ca250cbdb4 100644 --- a/providers/microsoft/azure/tests/unit/microsoft/azure/log/test_wasb_task_handler.py +++ b/providers/microsoft/azure/tests/unit/microsoft/azure/log/test_wasb_task_handler.py @@ -255,6 +255,76 @@ def test_write_on_existing_log(self, mock_log_exists, mock_wasb_read, mock_hook) overwrite=True, ) + @mock.patch("airflow.providers.microsoft.azure.hooks.wasb.WasbHook") + @mock.patch.object(WasbRemoteLogIO, "wasb_read") + @mock.patch.object(WasbRemoteLogIO, "wasb_log_exists") + def test_write_on_existing_log_already_newline_terminated( + self, mock_log_exists, mock_wasb_read, mock_hook + ): + mock_log_exists.return_value = True + mock_wasb_read.return_value = "old log\n" + self.wasb_task_handler.io.write("text", self.remote_log_location) + mock_hook.return_value.load_string.assert_called_once_with( + "old log\ntext", + self.container_name, + self.remote_log_location, + overwrite=True, + ) + + def test_upload_repeated_appends_no_duplication(self, tmp_path): + """Each execution lifecycle of one attempt appends to the same local log and uploads it.""" + blobs: dict[str, str] = {} + + class FakeHook: + def check_for_blob(self, container, blob_name, **kwargs): + return blob_name in blobs + + def read_file(self, container, blob_name, **kwargs): + return blobs.get(blob_name, "") + + def load_string(self, string_data, container, blob_name, **kwargs): + blobs[blob_name] = string_data + + io = WasbRemoteLogIO( + remote_base="remote/log/location", + base_log_folder=str(tmp_path), + delete_local_copy=False, + wasb_container=self.container_name, + ) + io.hook = FakeHook() + + local_log = tmp_path / "attempt=1.log" + for cycle in range(1, 4): + with open(local_log, "a") as f: + f.write(f"cycle {cycle}\n") + io.upload("attempt=1.log") + + assert blobs["remote/log/location/attempt=1.log"] == "cycle 1\ncycle 2\ncycle 3\n" + assert local_log.read_text() == "" + + @pytest.mark.parametrize( + ("has_uploaded", "expected_local_content"), + [(True, ""), (False, "some log\n")], + ) + @mock.patch.object(WasbRemoteLogIO, "write") + def test_upload_truncates_local_log_only_after_successful_write( + self, mock_write, tmp_path, has_uploaded, expected_local_content + ): + """A failed upload must leave the local log intact so it can be retried.""" + mock_write.return_value = has_uploaded + io = WasbRemoteLogIO( + remote_base="remote/log/location", + base_log_folder=str(tmp_path), + delete_local_copy=False, + wasb_container=self.container_name, + ) + local_log = tmp_path / "attempt=1.log" + local_log.write_text("some log\n") + + io.upload("attempt=1.log") + + assert local_log.read_text() == expected_local_content + @mock.patch("airflow.providers.microsoft.azure.hooks.wasb.WasbHook") def test_write_when_append_is_false(self, mock_hook): self.wasb_task_handler.io.write("text", self.remote_log_location, False)