Skip to content
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

Jump to bottom

Remove remaining Airflow 2.5 backcompat code from GCS Task Handler #36443

Merged
merged 1 commit into from
Dec 27, 2023
Merged

Remove remaining Airflow 2.5 backcompat code from GCS Task Handler #36443

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
41 changes: 2 additions & 39 deletions airflow/providers/google/cloud/log/gcs_task_handler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,6 @@

# not sure why but mypy complains on missing `storage` but it is clearly there and is importable
from google.cloud import storage # type: ignore[attr-defined]
from packaging.version import Version

from airflow.configuration import conf
from airflow.exceptions import AirflowNotFoundException
Expand All @@ -48,18 +47,6 @@
logger = logging.getLogger(__name__)


def get_default_delete_local_copy():
"""Load delete_local_logs conf if Airflow version > 2.6 and return False if not.

TODO: delete this function when min airflow version >= 2.6.
"""
from airflow.version import version

if Version(version) < Version("2.6"):
return False
return conf.getboolean("logging", "delete_local_logs")


class GCSTaskHandler(FileTaskHandler, LoggingMixin):
"""
GCSTaskHandler is a python log handler that handles and reads task instance logs.
Expand Down Expand Up @@ -108,8 +95,8 @@ def __init__(
self.gcp_keyfile_dict = gcp_keyfile_dict
self.scopes = gcp_scopes
self.project_id = project_id
self.delete_local_copy = (
kwargs["delete_local_copy"] if "delete_local_copy" in kwargs else get_default_delete_local_copy()
self.delete_local_copy = kwargs.get(
"delete_local_copy", conf.getboolean("logging", "delete_local_logs")
)

@cached_property
Expand Down Expand Up @@ -218,30 +205,6 @@ def _read_remote_logs(self, ti, try_number, metadata=None) -> tuple[list[str], l
messages.append(f"Unable to read remote log {e}")
return messages, logs

def _read(self, ti, try_number, metadata=None):
"""
Read logs of given task instance and try_number from GCS.

If failed, read the log from task instance host machine.

todo: when min airflow version >= 2.6, remove this method

:param ti: task instance object
:param try_number: task instance try_number to read logs from
:param metadata: log metadata,
can be used for steaming log reading and auto-tailing.
"""
if hasattr(super(), "_read_remote_logs"):
# from Airflow 2.6, we don't implement the `_read` method.
# if parent has _read_remote_logs, we're >= 2.6
return super()._read(ti, try_number, metadata)

messages, logs = self._read_remote_logs(ti, try_number, metadata)
if not logs:
return super()._read(ti, try_number, metadata)

return "".join([f"*** {x}\n" for x in messages]) + "\n".join(logs), {"end_of_log": True}

def gcs_write(self, log, remote_log_location) -> bool:
"""
Write the log to the remote location and return `True`; fail silently and return `False` on error.
Expand Down
9 changes: 3 additions & 6 deletions tests/providers/google/cloud/log/test_gcs_task_handler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -248,8 +248,8 @@ def test_write_to_remote_on_close_failed_read_old_logs(self, mock_blob, mock_cli
)

@pytest.mark.parametrize(
"delete_local_copy, expected_existence_of_local_copy, airflow_version",
[(True, False, "2.6.0"), (False, True, "2.6.0"), (True, True, "2.5.0"), (False, True, "2.5.0")],
"delete_local_copy, expected_existence_of_local_copy",
[(True, False), (False, True)],
)
@mock.patch(
"airflow.providers.google.cloud.log.gcs_task_handler.get_credentials_and_project_id",
Expand All @@ -265,12 +265,9 @@ def test_close_with_delete_local_copy_conf(
local_log_location,
delete_local_copy,
expected_existence_of_local_copy,
airflow_version,
):
mock_blob.from_string.return_value.download_as_bytes.return_value = b"CONTENT"
with conf_vars({("logging", "delete_local_logs"): str(delete_local_copy)}), mock.patch(
"airflow.version.version", airflow_version
):
with conf_vars({("logging", "delete_local_logs"): str(delete_local_copy)}):
handler = GCSTaskHandler(
base_log_folder=local_log_location,
gcs_log_folder="gs://bucket/remote/log/location",
Expand Down
Loading