diff --git a/airflow/serialization/helpers.py b/airflow/serialization/helpers.py index cbaa9c0bcaf6ff..aa262c2f34a690 100644 --- a/airflow/serialization/helpers.py +++ b/airflow/serialization/helpers.py @@ -43,7 +43,8 @@ def is_jsonable(x): if not is_jsonable(template_field): if isinstance(template_field, (list, tuple)): - if sys.getsizeof(template_field[0]) > max_size: + size = sum(sys.getsizeof(x) for x in template_field) + if size > max_size: return ( "Value redacted as it is too large to be stored in the database. " "You can change this behaviour in [core]max_templated_field_size" @@ -56,7 +57,8 @@ def is_jsonable(x): return str(template_field) else: if template_field and isinstance(template_field, (list, tuple)): - if len(template_field) and sys.getsizeof(template_field[0]) > max_size: + size = sum(sys.getsizeof(x) for x in template_field) + if size > max_size: return ( "Value redacted as it is too large to be stored in the database. " "You can change this behaviour in [core]max_templated_field_size" diff --git a/tests/models/test_renderedtifields.py b/tests/models/test_renderedtifields.py index 8993404d245143..d718ed43a84bb4 100644 --- a/tests/models/test_renderedtifields.py +++ b/tests/models/test_renderedtifields.py @@ -33,7 +33,7 @@ from airflow.utils.task_instance_session import set_current_task_instance_session from airflow.utils.timezone import datetime from tests.test_utils.asserts import assert_queries_count -from tests.test_utils.db import clear_db_dags, clear_db_runs, clear_rendered_ti_fields +from tests.test_utils.db import clear_db_dags, clear_db_runs, clear_db_xcom, clear_rendered_ti_fields pytestmark = pytest.mark.db_test @@ -70,6 +70,7 @@ def clean_db(): clear_db_runs() clear_db_dags() clear_rendered_ti_fields() + clear_db_xcom() def setup_method(self): self.clean_db() @@ -340,7 +341,10 @@ def consumer_task(value): ti, ti2 = dr.task_instances ti.xcom_push(value=large_string, key="return_value") rtif = RTIF(ti=ti2) - rtif.write() + rtif.write(session=session) + session.flush() + rtif = session.query(RTIF).filter(RTIF.dag_id == rtif.dag_id, RTIF.task_id == rtif.task_id).first() + assert rtif.rendered_fields == { "op_args": "Value redacted as it is too large to be stored in the database. " "You can change this behaviour in [core]max_templated_field_size", @@ -371,7 +375,9 @@ def consumer_task(value): ti, ti2 = dr.task_instances ti.xcom_push(value=large_dataframe, key="return_value") rtif = RTIF(ti=ti2) - rtif.write() + rtif.write(session=session) + session.flush() + rtif = session.query(RTIF).filter(RTIF.dag_id == rtif.dag_id, RTIF.task_id == rtif.task_id).first() assert rtif.rendered_fields == { "op_args": "Value redacted as it is too large to be stored in the database. " "You can change this behaviour in [core]max_templated_field_size",