From 2b07ed15fd80ff3c43e0cf5fe50ef87df9f92ee2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Kamil=20Bregu=C5=82a?= Date: Wed, 20 Jan 2021 02:05:01 +0100 Subject: [PATCH 1/4] Fix StackdriverTaskHandler --- .../cloud/log/stackdriver_task_handler.py | 24 +++++++++++++------ airflow/utils/log/log_reader.py | 6 ++--- .../test_stackdriver_task_handler_system.py | 4 ++++ 3 files changed, 24 insertions(+), 10 deletions(-) diff --git a/airflow/providers/google/cloud/log/stackdriver_task_handler.py b/airflow/providers/google/cloud/log/stackdriver_task_handler.py index be75fcd973e7f..ac07debaa3307 100644 --- a/airflow/providers/google/cloud/log/stackdriver_task_handler.py +++ b/airflow/providers/google/cloud/log/stackdriver_task_handler.py @@ -99,6 +99,7 @@ def __init__( self.resource: Resource = resource self.labels: Optional[Dict[str, str]] = labels self.task_instance_labels: Optional[Dict[str, str]] = {} + self.task_instance_hostname = 'default-hostname' @cached_property def _client(self) -> gcp_logging.Client: @@ -146,10 +147,11 @@ def set_context(self, task_instance: TaskInstance) -> None: :type task_instance: :class:`airflow.models.TaskInstance` """ self.task_instance_labels = self._task_instance_to_labels(task_instance) + self.task_instance_hostname = task_instance.hostname def read( self, task_instance: TaskInstance, try_number: Optional[int] = None, metadata: Optional[Dict] = None - ) -> Tuple[List[str], List[Dict]]: + ) -> Tuple[List[Tuple[Tuple[str, str]]], List[Dict[str, str]]]: """ Read logs of given task instance from Stackdriver logging. @@ -160,12 +162,15 @@ def read( :type try_number: Optional[int] :param metadata: log metadata. It is used for steaming log reading and auto-tailing. :type metadata: Dict - :return: a tuple of list of logs and list of metadata - :rtype: Tuple[List[str], List[Dict]] + :return: a tuple of ( + list of (one element tuple with two element tuple - hostname and logs) + and list of metadata + ) + :rtype: Tuple[List[Tuple[Tuple[str, str]]], List[Dict[str, str]]] """ if try_number is not None and try_number < 1: - logs = [f"Error fetching the logs. Try number {try_number} is invalid."] - return logs, [{"end_of_log": "true"}] + logs = f"Error fetching the logs. Try number {try_number} is invalid." + return [((self.task_instance_hostname, logs),)], [{"end_of_log": "true"}] if not metadata: metadata = {} @@ -188,7 +193,7 @@ def read( if next_page_token: new_metadata['next_page_token'] = next_page_token - return [messages], [new_metadata] + return [((self.task_instance_hostname, messages),)], [new_metadata] def _prepare_log_filter(self, ti_labels: Dict[str, str]) -> str: """ @@ -271,7 +276,9 @@ def _read_single_logs_page(self, log_filter: str, page_token: Optional[str] = No :return: Downloaded logs and next page token :rtype: Tuple[str, str] """ - entries = self._client.list_entries(filter_=log_filter, page_token=page_token) + entries = self._client.list_entries( + filter_=log_filter, page_token=page_token, order_by='timestamp asc', page_size=1000 + ) page = next(entries.pages) next_page_token = entries.next_page_token messages = [] @@ -331,3 +338,6 @@ def get_external_log_url(self, task_instance: TaskInstance, try_number: int) -> url = f"{self.LOG_VIEWER_BASE_URL}?{urlencode(url_query_string)}" return url + + def close(self) -> None: + self._transport.flush() diff --git a/airflow/utils/log/log_reader.py b/airflow/utils/log/log_reader.py index bef39a0b25af2..7ed6287d3ee8a 100644 --- a/airflow/utils/log/log_reader.py +++ b/airflow/utils/log/log_reader.py @@ -16,7 +16,7 @@ # under the License. import logging -from typing import Any, Dict, Iterator, List, Optional, Tuple +from typing import Dict, Iterator, List, Optional, Tuple from cached_property import cached_property @@ -31,7 +31,7 @@ class TaskLogReader: def read_log_chunks( self, ti: TaskInstance, try_number: Optional[int], metadata - ) -> Tuple[List[str], Dict[str, Any]]: + ) -> Tuple[List[Tuple[Tuple[str, str]]], Dict[str, str]]: """ Reads chunks of Task Instance logs. @@ -42,7 +42,7 @@ def read_log_chunks( :type try_number: Optional[int] :param metadata: A dictionary containing information about how to read the task log :type metadata: dict - :rtype: Tuple[List[str], Dict[str, Any]] + :rtype: Tuple[List[Tuple[Tuple[str, str]]], Dict[str, str]] The following is an example of how to use this method to read log: diff --git a/tests/providers/google/cloud/log/test_stackdriver_task_handler_system.py b/tests/providers/google/cloud/log/test_stackdriver_task_handler_system.py index 46964bb0bfae0..1d094ca2d2fe2 100644 --- a/tests/providers/google/cloud/log/test_stackdriver_task_handler_system.py +++ b/tests/providers/google/cloud/log/test_stackdriver_task_handler_system.py @@ -94,4 +94,8 @@ def assert_remote_logs(self, expected_message, ti): task_log_reader = TaskLogReader() logs = "\n".join(task_log_reader.read_log_stream(ti, try_number=None, metadata={})) + # Preview content + print("=" * 80) + print(logs) + print("=" * 80) assert expected_message in logs From affd8dbf28b2a53e144842c0b6d4ab2834c39b8d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Kamil=20Bregu=C5=82a?= Date: Wed, 20 Jan 2021 21:01:46 +0100 Subject: [PATCH 2/4] fixup! Fix StackdriverTaskHandler --- .../cloud/log/stackdriver_task_handler.py | 5 +- .../log/test_stackdriver_task_handler.py | 98 ++++++++++++------- .../test_stackdriver_task_handler_system.py | 4 +- 3 files changed, 70 insertions(+), 37 deletions(-) diff --git a/airflow/providers/google/cloud/log/stackdriver_task_handler.py b/airflow/providers/google/cloud/log/stackdriver_task_handler.py index ac07debaa3307..445806c6e5952 100644 --- a/airflow/providers/google/cloud/log/stackdriver_task_handler.py +++ b/airflow/providers/google/cloud/log/stackdriver_task_handler.py @@ -164,8 +164,7 @@ def read( :type metadata: Dict :return: a tuple of ( list of (one element tuple with two element tuple - hostname and logs) - and list of metadata - ) + and list of metadata) :rtype: Tuple[List[Tuple[Tuple[str, str]]], List[Dict[str, str]]] """ if try_number is not None and try_number < 1: @@ -257,6 +256,8 @@ def _read_logs( log_filter=log_filter, page_token=next_page_token ) messages.append(new_messages) + if not messages: + break end_of_log = True next_page_token = None diff --git a/tests/providers/google/cloud/log/test_stackdriver_task_handler.py b/tests/providers/google/cloud/log/test_stackdriver_task_handler.py index 4159e9e0f4e54..81de22a86548f 100644 --- a/tests/providers/google/cloud/log/test_stackdriver_task_handler.py +++ b/tests/providers/google/cloud/log/test_stackdriver_task_handler.py @@ -35,10 +35,21 @@ def _create_list_response(messages, token): return mock.MagicMock(pages=(n for n in [page]), next_page_token=token) +def _remove_stackdriver_handlers(): + for handler_ref in logging._handlerList: + handler = handler_ref() + if not isinstance(handler, StackdriverTaskHandler): + continue + logging._removeHandlerRef(handler_ref) + del handler + + class TestStackdriverLoggingHandlerStandalone(unittest.TestCase): @mock.patch('airflow.providers.google.cloud.log.stackdriver_task_handler.get_credentials_and_project_id') @mock.patch('airflow.providers.google.cloud.log.stackdriver_task_handler.gcp_logging.Client') def test_should_pass_message_to_client(self, mock_client, mock_get_creds_and_project_id): + self.addCleanup(_remove_stackdriver_handlers) + mock_get_creds_and_project_id.return_value = ('creds', 'project_id') transport_type = mock.MagicMock() @@ -69,6 +80,7 @@ def setUp(self) -> None: self.ti.try_number = 1 self.ti.state = State.RUNNING self.addCleanup(self.dag.clear) + self.addCleanup(_remove_stackdriver_handlers) @mock.patch('airflow.providers.google.cloud.log.stackdriver_task_handler.get_credentials_and_project_id') @mock.patch('airflow.providers.google.cloud.log.stackdriver_task_handler.gcp_logging.Client') @@ -128,14 +140,18 @@ def test_should_read_logs_for_all_try(self, mock_client, mock_get_creds_and_proj logs, metadata = self.stackdriver_task_handler.read(self.ti) mock_client.return_value.list_entries.assert_called_once_with( - filter_='resource.type="global"\n' - 'logName="projects/asf-project/logs/airflow"\n' - 'labels.task_id="task_for_testing_file_log_handler"\n' - 'labels.dag_id="dag_for_testing_file_task_handler"\n' - 'labels.execution_date="2016-01-01T00:00:00+00:00"', + filter_=( + 'resource.type="global"\n' + 'logName="projects/asf-project/logs/airflow"\n' + 'labels.task_id="task_for_testing_file_log_handler"\n' + 'labels.dag_id="dag_for_testing_file_task_handler"\n' + 'labels.execution_date="2016-01-01T00:00:00+00:00"' + ), + order_by='timestamp asc', + page_size=1000, page_token=None, ) - assert ['MSG1\nMSG2'] == logs + assert [(('default-hostname', 'MSG1\nMSG2'),)] == logs assert [{'end_of_log': True}] == metadata @mock.patch('airflow.providers.google.cloud.log.stackdriver_task_handler.get_credentials_and_project_id') @@ -149,14 +165,18 @@ def test_should_read_logs_for_task_with_quote(self, mock_client, mock_get_creds_ self.ti.task_id = "K\"OT" logs, metadata = self.stackdriver_task_handler.read(self.ti) mock_client.return_value.list_entries.assert_called_once_with( - filter_='resource.type="global"\n' - 'logName="projects/asf-project/logs/airflow"\n' - 'labels.task_id="K\\"OT"\n' - 'labels.dag_id="dag_for_testing_file_task_handler"\n' - 'labels.execution_date="2016-01-01T00:00:00+00:00"', + filter_=( + 'resource.type="global"\n' + 'logName="projects/asf-project/logs/airflow"\n' + 'labels.task_id="K\\"OT"\n' + 'labels.dag_id="dag_for_testing_file_task_handler"\n' + 'labels.execution_date="2016-01-01T00:00:00+00:00"' + ), + order_by='timestamp asc', + page_size=1000, page_token=None, ) - assert ['MSG1\nMSG2'] == logs + assert [(('default-hostname', 'MSG1\nMSG2'),)] == logs assert [{'end_of_log': True}] == metadata @mock.patch('airflow.providers.google.cloud.log.stackdriver_task_handler.get_credentials_and_project_id') @@ -170,15 +190,19 @@ def test_should_read_logs_for_single_try(self, mock_client, mock_get_creds_and_p logs, metadata = self.stackdriver_task_handler.read(self.ti, 3) mock_client.return_value.list_entries.assert_called_once_with( - filter_='resource.type="global"\n' - 'logName="projects/asf-project/logs/airflow"\n' - 'labels.task_id="task_for_testing_file_log_handler"\n' - 'labels.dag_id="dag_for_testing_file_task_handler"\n' - 'labels.execution_date="2016-01-01T00:00:00+00:00"\n' - 'labels.try_number="3"', + filter_=( + 'resource.type="global"\n' + 'logName="projects/asf-project/logs/airflow"\n' + 'labels.task_id="task_for_testing_file_log_handler"\n' + 'labels.dag_id="dag_for_testing_file_task_handler"\n' + 'labels.execution_date="2016-01-01T00:00:00+00:00"\n' + 'labels.try_number="3"' + ), + order_by='timestamp asc', + page_size=1000, page_token=None, ) - assert ['MSG1\nMSG2'] == logs + assert [(('default-hostname', 'MSG1\nMSG2'),)] == logs assert [{'end_of_log': True}] == metadata @mock.patch('airflow.providers.google.cloud.log.stackdriver_task_handler.get_credentials_and_project_id') @@ -190,14 +214,18 @@ def test_should_read_logs_with_pagination(self, mock_client, mock_get_creds_and_ ] mock_get_creds_and_project_id.return_value = ('creds', 'project_id') logs, metadata1 = self.stackdriver_task_handler.read(self.ti, 3) - mock_client.return_value.list_entries.assert_called_once_with(filter_=mock.ANY, page_token=None) - assert ['MSG1\nMSG2'] == logs + mock_client.return_value.list_entries.assert_called_once_with( + filter_=mock.ANY, order_by='timestamp asc', page_size=1000, page_token=None + ) + assert [(('default-hostname', 'MSG1\nMSG2'),)] == logs assert [{'end_of_log': False, 'next_page_token': 'TOKEN1'}] == metadata1 mock_client.return_value.list_entries.return_value.next_page_token = None logs, metadata2 = self.stackdriver_task_handler.read(self.ti, 3, metadata1[0]) - mock_client.return_value.list_entries.assert_called_with(filter_=mock.ANY, page_token="TOKEN1") - assert ['MSG3\nMSG4'] == logs + mock_client.return_value.list_entries.assert_called_with( + filter_=mock.ANY, order_by='timestamp asc', page_size=1000, page_token="TOKEN1" + ) + assert [(('default-hostname', 'MSG3\nMSG4'),)] == logs assert [{'end_of_log': True}] == metadata2 @mock.patch('airflow.providers.google.cloud.log.stackdriver_task_handler.get_credentials_and_project_id') @@ -211,7 +239,7 @@ def test_should_read_logs_with_download(self, mock_client, mock_get_creds_and_pr logs, metadata1 = self.stackdriver_task_handler.read(self.ti, 3, {'download_logs': True}) - assert ['MSG1\nMSG2\nMSG3\nMSG4'] == logs + assert [(('default-hostname', 'MSG1\nMSG2\nMSG3\nMSG4'),)] == logs assert [{'end_of_log': True}] == metadata1 @mock.patch('airflow.providers.google.cloud.log.stackdriver_task_handler.get_credentials_and_project_id') @@ -240,17 +268,21 @@ def test_should_read_logs_with_custom_resources(self, mock_client, mock_get_cred logs, metadata = self.stackdriver_task_handler.read(self.ti) mock_client.return_value.list_entries.assert_called_once_with( - filter_='resource.type="cloud_composer_environment"\n' - 'logName="projects/asf-project/logs/airflow"\n' - 'resource.labels."environment.name"="test-instancce"\n' - 'resource.labels.location="europpe-west-3"\n' - 'resource.labels.project_id="asf-project"\n' - 'labels.task_id="task_for_testing_file_log_handler"\n' - 'labels.dag_id="dag_for_testing_file_task_handler"\n' - 'labels.execution_date="2016-01-01T00:00:00+00:00"', + filter_=( + 'resource.type="cloud_composer_environment"\n' + 'logName="projects/asf-project/logs/airflow"\n' + 'resource.labels."environment.name"="test-instancce"\n' + 'resource.labels.location="europpe-west-3"\n' + 'resource.labels.project_id="asf-project"\n' + 'labels.task_id="task_for_testing_file_log_handler"\n' + 'labels.dag_id="dag_for_testing_file_task_handler"\n' + 'labels.execution_date="2016-01-01T00:00:00+00:00"' + ), + order_by='timestamp asc', + page_size=1000, page_token=None, ) - assert ['TEXT\nTEXT'] == logs + assert [(('default-hostname', 'TEXT\nTEXT'),)] == logs assert [{'end_of_log': True}] == metadata @mock.patch('airflow.providers.google.cloud.log.stackdriver_task_handler.get_credentials_and_project_id') diff --git a/tests/providers/google/cloud/log/test_stackdriver_task_handler_system.py b/tests/providers/google/cloud/log/test_stackdriver_task_handler_system.py index 1d094ca2d2fe2..c4827dacb9140 100644 --- a/tests/providers/google/cloud/log/test_stackdriver_task_handler_system.py +++ b/tests/providers/google/cloud/log/test_stackdriver_task_handler_system.py @@ -62,7 +62,7 @@ def test_should_support_key_auth(self, session): assert 0 == subprocess.Popen(["airflow", "scheduler", "--num-runs", "1"]).wait() ti = session.query(TaskInstance).filter(TaskInstance.task_id == "create_entry_group").first() - self.assert_remote_logs("INFO - Task exited with return code 0", ti) + self.assert_remote_logs("terminated with exit code 0", ti) @provide_session def test_should_support_adc(self, session): @@ -78,7 +78,7 @@ def test_should_support_adc(self, session): assert 0 == subprocess.Popen(["airflow", "scheduler", "--num-runs", "1"]).wait() ti = session.query(TaskInstance).filter(TaskInstance.task_id == "create_entry_group").first() - self.assert_remote_logs("INFO - Task exited with return code 0", ti) + self.assert_remote_logs("terminated with exit code 0", ti) def assert_remote_logs(self, expected_message, ti): with provide_gcp_context(GCP_STACKDRIVER), conf_vars( From 0944b983d65d29fa0636c924578c1c90a227c687 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Kamil=20Bregu=C5=82a?= Date: Tue, 2 Feb 2021 07:23:05 +0100 Subject: [PATCH 3/4] fixup! fixup! Fix StackdriverTaskHandler --- .../cloud/log/test_stackdriver_task_handler.py | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/tests/providers/google/cloud/log/test_stackdriver_task_handler.py b/tests/providers/google/cloud/log/test_stackdriver_task_handler.py index 81de22a86548f..30587b7586784 100644 --- a/tests/providers/google/cloud/log/test_stackdriver_task_handler.py +++ b/tests/providers/google/cloud/log/test_stackdriver_task_handler.py @@ -30,13 +30,28 @@ from airflow.utils.state import State +def _monkey_patch_shutdown(): + orig_shutdown = logging.shutdown + + def new_shutdown(*args, **kwargs): + for handler_no, weak_ref in enumerate(logging._handlerList[:]): + handler = weak_ref() + print("Handler", handler_no, ":", handler, ":", type(handler)) + return orig_shutdown(*args, **kwargs) + + logging.shutdown = new_shutdown + + +_monkey_patch_shutdown() + + def _create_list_response(messages, token): page = [mock.MagicMock(payload={"message": message}) for message in messages] return mock.MagicMock(pages=(n for n in [page]), next_page_token=token) def _remove_stackdriver_handlers(): - for handler_ref in logging._handlerList: + for handler_ref in reversed(logging._handlerList[:]): handler = handler_ref() if not isinstance(handler, StackdriverTaskHandler): continue From 008668d63981e08e18f38e704487ff2bfc5eb606 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Kamil=20Bregu=C5=82a?= Date: Tue, 2 Feb 2021 11:59:45 +0100 Subject: [PATCH 4/4] fixup! fixup! fixup! Fix StackdriverTaskHandler --- tests/cli/commands/test_info_command.py | 3 +++ .../cloud/log/test_stackdriver_task_handler.py | 15 --------------- 2 files changed, 3 insertions(+), 15 deletions(-) diff --git a/tests/cli/commands/test_info_command.py b/tests/cli/commands/test_info_command.py index 7fad6e83dcbfa..3d4c60fc04e8f 100644 --- a/tests/cli/commands/test_info_command.py +++ b/tests/cli/commands/test_info_command.py @@ -18,6 +18,7 @@ import contextlib import importlib import io +import logging import os import unittest from unittest import mock @@ -129,6 +130,8 @@ def test_should_read_logging_configuration(self): assert "stackdriver" in text def tearDown(self) -> None: + for handler_ref in logging._handlerList[:]: + logging._removeHandlerRef(handler_ref) importlib.reload(airflow_local_settings) configure_logging() diff --git a/tests/providers/google/cloud/log/test_stackdriver_task_handler.py b/tests/providers/google/cloud/log/test_stackdriver_task_handler.py index 30587b7586784..b465d79fd97c9 100644 --- a/tests/providers/google/cloud/log/test_stackdriver_task_handler.py +++ b/tests/providers/google/cloud/log/test_stackdriver_task_handler.py @@ -30,21 +30,6 @@ from airflow.utils.state import State -def _monkey_patch_shutdown(): - orig_shutdown = logging.shutdown - - def new_shutdown(*args, **kwargs): - for handler_no, weak_ref in enumerate(logging._handlerList[:]): - handler = weak_ref() - print("Handler", handler_no, ":", handler, ":", type(handler)) - return orig_shutdown(*args, **kwargs) - - logging.shutdown = new_shutdown - - -_monkey_patch_shutdown() - - def _create_list_response(messages, token): page = [mock.MagicMock(payload={"message": message}) for message in messages] return mock.MagicMock(pages=(n for n in [page]), next_page_token=token)