This is an automated email from the ASF dual-hosted git repository. vincbeck pushed a commit to branch main in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push: new c01abd1c2e Upgrade watchtower to 3.0.1 (#25019) (#34747) c01abd1c2e is described below commit c01abd1c2eed8f60fec5b9d6cc0232b54efa52de Author: cBiscuitSurprise <36653704+cbiscuitsurpr...@users.noreply.github.com> AuthorDate: Fri Oct 6 09:35:09 2023 -0500 Upgrade watchtower to 3.0.1 (#25019) (#34747) --- .../amazon/aws/log/cloudwatch_task_handler.py | 35 ++++++++++++++- airflow/providers/amazon/provider.yaml | 24 ++++++++--- docs/apache-airflow-providers-amazon/index.rst | 2 +- generated/provider_dependencies.json | 2 +- .../amazon/aws/log/test_cloudwatch_task_handler.py | 50 +++++++++++++++++++++- 5 files changed, 103 insertions(+), 10 deletions(-) diff --git a/airflow/providers/amazon/aws/log/cloudwatch_task_handler.py b/airflow/providers/amazon/aws/log/cloudwatch_task_handler.py index 4d3ccec00c..126ed3ea84 100644 --- a/airflow/providers/amazon/aws/log/cloudwatch_task_handler.py +++ b/airflow/providers/amazon/aws/log/cloudwatch_task_handler.py @@ -17,9 +17,9 @@ # under the License. from __future__ import annotations -from datetime import datetime, timedelta +from datetime import date, datetime, timedelta from functools import cached_property -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any import watchtower @@ -33,6 +33,35 @@ if TYPE_CHECKING: from airflow.models import TaskInstance +def json_serialize_legacy(value: Any) -> str | None: + """ + JSON serializer replicating legacy watchtower behavior. + + The legacy `watchtower@2.0.1` json serializer function that serialized + datetime objects as ISO format and all other non-JSON-serializable to `null`. + + :param value: the object to serialize + :return: string representation of `value` if it is an instance of datetime or `None` otherwise + """ + if isinstance(value, (date, datetime)): + return value.isoformat() + else: + return None + + +def json_serialize(value: Any) -> str | None: + """ + JSON serializer replicating current watchtower behavior. + + This provides customers with an accessible import, + `airflow.providers.amazon.aws.log.cloudwatch_task_handler.json_serialize` + + :param value: the object to serialize + :return: string representation of `value` + """ + return watchtower._json_serialize_default(value) + + class CloudwatchTaskHandler(FileTaskHandler, LoggingMixin): """ CloudwatchTaskHandler is a python log handler that handles and reads task instance logs. @@ -69,11 +98,13 @@ class CloudwatchTaskHandler(FileTaskHandler, LoggingMixin): def set_context(self, ti): super().set_context(ti) + self.json_serialize = conf.getimport("aws", "cloudwatch_task_handler_json_serializer") self.handler = watchtower.CloudWatchLogHandler( log_group_name=self.log_group, log_stream_name=self._render_filename(ti, ti.try_number), use_queues=not getattr(ti, "is_trigger_log_context", False), boto3_client=self.hook.get_conn(), + json_serialize_default=self.json_serialize, ) def close(self): diff --git a/airflow/providers/amazon/provider.yaml b/airflow/providers/amazon/provider.yaml index 8ce738480c..3f5ad39025 100644 --- a/airflow/providers/amazon/provider.yaml +++ b/airflow/providers/amazon/provider.yaml @@ -81,11 +81,7 @@ dependencies: # NOTE!!! BOTOCORE version is always shifted by 3 MINOR VERSIONS from boto3 # See https://github.com/boto/boto3/issues/2702 - botocore>=1.31.0 - # watchtower 3 has been released end Jan and introduced breaking change across the board that might - # change logging behaviour: - # https://github.com/kislyuk/watchtower/blob/develop/Changes.rst#changes-for-v300-2022-01-26 - # TODO: update to watchtower >3 - - watchtower~=2.0.1 + - watchtower~=3.0.1 - jsonpath_ng>=1.5.3 - redshift_connector>=2.0.888 - sqlalchemy_redshift>=0.8.6 @@ -726,3 +722,21 @@ config: example: my_company.aws.MyCustomSessionFactory type: string version_added: 3.1.1 + cloudwatch_task_handler_json_serializer: + description: | + By default, when logging non-string messages, all non-json objects are logged as `null`. + Except `datetime` objects which are ISO formatted. Users can optionally use a `repr` serializer or + provide their own JSON serializer for any non-JSON-serializable objects in the logged message. + + * `airflow.providers.amazon.aws.log.cloudwatch_task_handler.json_serialize` uses `repr` (be aware + there is the potential of logging sensitive data depending on the `repr` method of logged objects) + * `airflow.providers.amazon.aws.log.cloudwatch_task_handler.json_serialize_legacy` uses `null`. + + If a custom serializer is provided, it must adhere to `Callable[[Any], str | None]`, where `None` + serializes to `null` (e.g. `def my_serializer(o: Any) -> str | None`). Since this is on the logging + path and it's possible there's an exception being handled, special care should be taken to fail + gracefully without raising a new exception inside of your serializer. + type: string + version_added: 8.7.2 + example: airflow.providers.amazon.aws.log.cloudwatch_task_handler.json_serialize + default: airflow.providers.amazon.aws.log.cloudwatch_task_handler.json_serialize_legacy diff --git a/docs/apache-airflow-providers-amazon/index.rst b/docs/apache-airflow-providers-amazon/index.rst index d846b73c56..c1f52f9651 100644 --- a/docs/apache-airflow-providers-amazon/index.rst +++ b/docs/apache-airflow-providers-amazon/index.rst @@ -112,7 +112,7 @@ PIP package Version required ``boto3`` ``>=1.28.0`` ``botocore`` ``>=1.31.0`` ``asgiref`` -``watchtower`` ``~=2.0.1`` +``watchtower`` ``~=3.0.1`` ``jsonpath_ng`` ``>=1.5.3`` ``redshift_connector`` ``>=2.0.888`` ``sqlalchemy_redshift`` ``>=0.8.6`` diff --git a/generated/provider_dependencies.json b/generated/provider_dependencies.json index 212f308989..5fc8ed0df0 100644 --- a/generated/provider_dependencies.json +++ b/generated/provider_dependencies.json @@ -30,7 +30,7 @@ "jsonpath_ng>=1.5.3", "redshift_connector>=2.0.888", "sqlalchemy_redshift>=0.8.6", - "watchtower~=2.0.1" + "watchtower~=3.0.1" ], "cross-providers-deps": [ "apache.hive", diff --git a/tests/providers/amazon/aws/log/test_cloudwatch_task_handler.py b/tests/providers/amazon/aws/log/test_cloudwatch_task_handler.py index 00d5cf2f46..b3c1cd6e39 100644 --- a/tests/providers/amazon/aws/log/test_cloudwatch_task_handler.py +++ b/tests/providers/amazon/aws/log/test_cloudwatch_task_handler.py @@ -17,10 +17,12 @@ # under the License. from __future__ import annotations +import contextlib +import logging import time from datetime import datetime as dt, timedelta from unittest import mock -from unittest.mock import call +from unittest.mock import ANY, Mock, call import boto3 import moto @@ -171,6 +173,52 @@ class TestCloudwatchTaskHandler: end_time=expected_end_time, ) + @pytest.mark.parametrize( + "conf_json_serialize, expected_serialized_output", + [ + (None, '{"datetime": "2023-01-01T00:00:00+00:00", "customObject": null}'), + ( + "airflow.providers.amazon.aws.log.cloudwatch_task_handler.json_serialize", + '{"datetime": "2023-01-01T00:00:00+00:00", "customObject": "SomeCustomSerialization(...)"}', + ), + ], + ) + @mock.patch.object(AwsLogsHook, "get_log_events") + def test_write_json_logs(self, mock_get_log_events, conf_json_serialize, expected_serialized_output): + class ToSerialize: + def __init__(self): + pass + + def __repr__(self): + return "SomeCustomSerialization(...)" + + with contextlib.ExitStack() as stack: + if conf_json_serialize: + stack.enter_context( + conf_vars({("aws", "cloudwatch_task_handler_json_serializer"): conf_json_serialize}) + ) + + handler = self.cloudwatch_task_handler + handler.set_context(self.ti) + message = logging.LogRecord( + name="test_log_record", + level=logging.DEBUG, + pathname="fake.path", + lineno=42, + args=None, + exc_info=None, + msg={ + "datetime": datetime(2023, 1, 1), + "customObject": ToSerialize(), + }, + ) + stack.enter_context(mock.patch("watchtower.threading.Thread")) + mock_queue = Mock() + stack.enter_context(mock.patch("watchtower.queue.Queue", return_value=mock_queue)) + handler.handle(message) + + mock_queue.put.assert_called_once_with({"message": expected_serialized_output, "timestamp": ANY}) + def test_close_prevents_duplicate_calls(self): with mock.patch("watchtower.CloudWatchLogHandler.close") as mock_log_handler_close: with mock.patch("airflow.utils.log.file_task_handler.FileTaskHandler.set_context"):