fixup! Prevent large objects from being stored in the RTIF

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"

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",