diff --git a/airflow/callbacks/callback_requests.py b/airflow/callbacks/callback_requests.py index 8112589cd0262..b04a201c08d06 100644 --- a/airflow/callbacks/callback_requests.py +++ b/airflow/callbacks/callback_requests.py @@ -75,7 +75,7 @@ def __init__( def to_json(self) -> str: dict_obj = self.__dict__.copy() - dict_obj["simple_task_instance"] = dict_obj["simple_task_instance"].__dict__ + dict_obj["simple_task_instance"] = self.simple_task_instance.as_dict() return json.dumps(dict_obj) @classmethod diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py index d83ad11b04ba5..e52976e35917e 100644 --- a/airflow/models/taskinstance.py +++ b/airflow/models/taskinstance.py @@ -2631,6 +2631,16 @@ def __eq__(self, other): return self.__dict__ == other.__dict__ return NotImplemented + def as_dict(self): + new_dict = dict(self.__dict__) + for key in new_dict: + if key in ['start_date', 'end_date']: + val = new_dict[key] + if not val or isinstance(val, str): + continue + new_dict.update({key: val.isoformat()}) + return new_dict + @classmethod def from_ti(cls, ti: TaskInstance) -> "SimpleTaskInstance": return cls( diff --git a/tests/callbacks/test_callback_requests.py b/tests/callbacks/test_callback_requests.py index 286d64eaa156e..3764f19c4c4f3 100644 --- a/tests/callbacks/test_callback_requests.py +++ b/tests/callbacks/test_callback_requests.py @@ -15,7 +15,6 @@ # specific language governing permissions and limitations # under the License. -import unittest from datetime import datetime from parameterized import parameterized @@ -29,6 +28,7 @@ from airflow.models.dag import DAG from airflow.models.taskinstance import SimpleTaskInstance, TaskInstance from airflow.operators.bash import BashOperator +from airflow.utils import timezone from airflow.utils.state import State TI = TaskInstance( @@ -38,7 +38,7 @@ ) -class TestCallbackRequest(unittest.TestCase): +class TestCallbackRequest: @parameterized.expand( [ (CallbackRequest(full_filepath="filepath", msg="task_failure"), CallbackRequest), @@ -64,7 +64,20 @@ class TestCallbackRequest(unittest.TestCase): ) def test_from_json(self, input, request_class): json_str = input.to_json() - result = request_class.from_json(json_str=json_str) + assert result == input - self.assertEqual(result, input) + def test_taskcallback_to_json_with_start_date_and_end_date(self, session, create_task_instance): + ti = create_task_instance() + ti.start_date = timezone.utcnow() + ti.end_date = timezone.utcnow() + session.merge(ti) + session.flush() + input = TaskCallbackRequest( + full_filepath="filepath", + simple_task_instance=SimpleTaskInstance.from_ti(ti), + is_failure_callback=True, + ) + json_str = input.to_json() + result = TaskCallbackRequest.from_json(json_str) + assert input == result