Skip to content

Avoid including fallback message for S3TaskHandler if streaming logs #29708

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
dstandish merged 3 commits into from
Feb 23, 2023
Merged

Avoid including fallback message for S3TaskHandler if streaming logs #29708

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
ebd2952
Avoid including fallback message for S3TaskHandler if streaming logs
pdebelak Feb 22, 2023
510f4b5
Apply code clean-up suggestions to s3 task handler fallback handling
pdebelak Feb 22, 2023
92aac29
Use fstring rather than string concatenation
pdebelak Feb 22, 2023
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
6 changes: 5 additions & 1 deletion airflow/providers/amazon/aws/log/s3_task_handler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -128,8 +128,12 @@ def _read(self, ti, try_number, metadata=None):
if logs:
return "".join(f"*** {x}\n" for x in messages) + "\n".join(logs), {"end_of_log": True}
else:
if metadata and metadata.get("log_pos", 0) > 0:
log_prefix = ""
else:
log_prefix = "*** Falling back to local log\n"
local_log, metadata = super()._read(ti, try_number, metadata)
return "*** Falling back to local log\n" + local_log, metadata
return f"{log_prefix}{local_log}", metadata

def s3_log_exists(self, remote_log_location: str) -> bool:
"""
Expand Down
27 changes: 27 additions & 0 deletions tests/providers/amazon/aws/log/test_s3_task_handler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -146,6 +146,33 @@ def test_read_when_s3_log_missing(self):
assert actual == expected
assert {"end_of_log": True, "log_pos": 0} == metadata[0]

def test_read_when_s3_log_missing_and_log_pos_missing_pre_26(self):
ti = copy.copy(self.ti)
ti.state = TaskInstanceState.SUCCESS
# mock that super class has no _read_remote_logs method
with mock.patch("airflow.providers.amazon.aws.log.s3_task_handler.hasattr", return_value=False):
log, metadata = self.s3_task_handler.read(ti)
assert 1 == len(log)
assert log[0][0][-1].startswith("*** Falling back to local log")

def test_read_when_s3_log_missing_and_log_pos_zero_pre_26(self):
ti = copy.copy(self.ti)
ti.state = TaskInstanceState.SUCCESS
# mock that super class has no _read_remote_logs method
with mock.patch("airflow.providers.amazon.aws.log.s3_task_handler.hasattr", return_value=False):
log, metadata = self.s3_task_handler.read(ti, metadata={"log_pos": 0})
assert 1 == len(log)
assert log[0][0][-1].startswith("*** Falling back to local log")

def test_read_when_s3_log_missing_and_log_pos_over_zero_pre_26(self):
ti = copy.copy(self.ti)
ti.state = TaskInstanceState.SUCCESS
# mock that super class has no _read_remote_logs method
with mock.patch("airflow.providers.amazon.aws.log.s3_task_handler.hasattr", return_value=False):
log, metadata = self.s3_task_handler.read(ti, metadata={"log_pos": 1})
assert 1 == len(log)
assert not log[0][0][-1].startswith("*** Falling back to local log")

def test_s3_read_when_log_missing(self):
handler = self.s3_task_handler
url = "s3://bucket/foo"
Expand Down