You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by ep...@apache.org on 2024/03/25 08:35:24 UTC
(airflow) branch main updated: Prevent large objects from being stored in the RTIF (#38094)
This is an automated email from the ASF dual-hosted git repository.
ephraimanierobi 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 420709962e Prevent large objects from being stored in the RTIF (#38094)
420709962e is described below
commit 420709962ec864fa62b6e1b6c2f723a2c14cb751
Author: Ephraim Anierobi <sp...@gmail.com>
AuthorDate: Mon Mar 25 09:35:17 2024 +0100
Prevent large objects from being stored in the RTIF (#38094)
* Prevent large objects from being stored in the RTIF
There's no control over the size of objects stored in the rendered taskinstance
field. This PR adds control and enable users to be able to customize the size of
data that can be stored in this field
closes: #28199
* fixup! Prevent large objects from being stored in the RTIF
* Use len and check the size of the serialized
* Add test to already existing test
* Remove xcom db clearing
* Update airflow/config_templates/config.yml
Co-authored-by: Jed Cunningham <66...@users.noreply.github.com>
* Apply review suggestions
* Add test for redacting values and add significant item
* Prefix with Truncated line
* Ensure secrets are masked
* fixup! Ensure secrets are masked
* Check the template field length in separate branches for jsonable and nonjsonable
* add tests for rendered dataframe
* Apply suggestions from code review
* update code and tests
* fixup! update code and tests
---------
Co-authored-by: Jed Cunningham <66...@users.noreply.github.com>
---
airflow/config_templates/config.yml | 8 ++++
airflow/models/renderedtifields.py | 3 +-
airflow/serialization/helpers.py | 22 ++++++++++-
airflow/serialization/serialized_objects.py | 2 +-
newsfragments/38094.significant.rst | 5 +++
tests/models/test_renderedtifields.py | 60 +++++++++++++++++++++++++++++
6 files changed, 97 insertions(+), 3 deletions(-)
diff --git a/airflow/config_templates/config.yml b/airflow/config_templates/config.yml
index 0a1b292fdc..4b086dda5e 100644
--- a/airflow/config_templates/config.yml
+++ b/airflow/config_templates/config.yml
@@ -492,6 +492,14 @@ core:
type: string
example: ~
default: "Disabled"
+ max_templated_field_length:
+ description: |
+ The maximum length of the rendered template field. If the value to be stored in the
+ rendered template field exceeds this size, it's redacted.
+ version_added: 2.9.0
+ type: integer
+ example: ~
+ default: "4096"
database:
description: ~
options:
diff --git a/airflow/models/renderedtifields.py b/airflow/models/renderedtifields.py
index 4354566a47..71c263dce1 100644
--- a/airflow/models/renderedtifields.py
+++ b/airflow/models/renderedtifields.py
@@ -121,7 +121,8 @@ class RenderedTaskInstanceFields(TaskInstanceDependencies):
self.k8s_pod_yaml = render_k8s_pod_yaml(ti)
self.rendered_fields = {
- field: serialize_template_field(getattr(self.task, field)) for field in self.task.template_fields
+ field: serialize_template_field(getattr(self.task, field), field)
+ for field in self.task.template_fields
}
self._redact()
diff --git a/airflow/serialization/helpers.py b/airflow/serialization/helpers.py
index aa3ceb644e..25d9cbacbf 100644
--- a/airflow/serialization/helpers.py
+++ b/airflow/serialization/helpers.py
@@ -20,10 +20,12 @@ from __future__ import annotations
from typing import Any
+from airflow.configuration import conf
from airflow.settings import json
+from airflow.utils.log.secrets_masker import redact
-def serialize_template_field(template_field: Any) -> str | dict | list | int | float:
+def serialize_template_field(template_field: Any, name: str) -> str | dict | list | int | float:
"""Return a serializable representation of the templated field.
If ``templated_field`` contains a class or instance that requires recursive
@@ -38,7 +40,25 @@ def serialize_template_field(template_field: Any) -> str | dict | list | int | f
else:
return True
+ max_length = conf.getint("core", "max_templated_field_length")
+
if not is_jsonable(template_field):
+ serialized = str(template_field)
+ if len(serialized) > max_length:
+ rendered = redact(serialized, name)
+ return (
+ "Truncated. You can change this behaviour in [core]max_templated_field_length. "
+ f"{rendered[:max_length - 79]!r}... "
+ )
return str(template_field)
else:
+ if not template_field:
+ return template_field
+ serialized = str(template_field)
+ if len(serialized) > max_length:
+ rendered = redact(serialized, name)
+ return (
+ "Truncated. You can change this behaviour in [core]max_templated_field_length. "
+ f"{rendered[:max_length - 79]!r}... "
+ )
return template_field
diff --git a/airflow/serialization/serialized_objects.py b/airflow/serialization/serialized_objects.py
index d17c76c0af..f0427f8f18 100644
--- a/airflow/serialization/serialized_objects.py
+++ b/airflow/serialization/serialized_objects.py
@@ -993,7 +993,7 @@ class SerializedBaseOperator(BaseOperator, BaseSerialization):
)
value = getattr(op, template_field, None)
if not cls._is_excluded(value, template_field, op):
- serialize_op[template_field] = serialize_template_field(value)
+ serialize_op[template_field] = serialize_template_field(value, template_field)
if op.params:
serialize_op["params"] = cls._serialize_params_dict(op.params)
diff --git a/newsfragments/38094.significant.rst b/newsfragments/38094.significant.rst
new file mode 100644
index 0000000000..25889c1c3c
--- /dev/null
+++ b/newsfragments/38094.significant.rst
@@ -0,0 +1,5 @@
+Prevent large string objects from being stored in the Rendered Template Fields
+
+There's now a limit to the length of data that can be stored in the Rendered Template Fields.
+The limit is set to 4096 characters. If the data exceeds this limit, it will be truncated. You can change this limit
+by setting the ``[core]max_template_field_length`` configuration option in your airflow config.
diff --git a/tests/models/test_renderedtifields.py b/tests/models/test_renderedtifields.py
index a7b25a3171..0d0ca3be07 100644
--- a/tests/models/test_renderedtifields.py
+++ b/tests/models/test_renderedtifields.py
@@ -27,6 +27,8 @@ from unittest import mock
import pytest
from airflow import settings
+from airflow.configuration import conf
+from airflow.decorators import task as task_decorator
from airflow.models import Variable
from airflow.models.renderedtifields import RenderedTaskInstanceFields as RTIF
from airflow.operators.bash import BashOperator
@@ -62,6 +64,17 @@ class ClassWithCustomAttributes:
return not self.__eq__(other)
+class LargeStrObject:
+ def __init__(self):
+ self.a = "a" * 5000
+
+ def __str__(self):
+ return self.a
+
+
+max_length = conf.getint("core", "max_templated_field_length")
+
+
class TestRenderedTaskInstanceFields:
"""Unit tests for RenderedTaskInstanceFields."""
@@ -111,6 +124,14 @@ class TestRenderedTaskInstanceFields:
"{'att3': '{{ task.task_id }}', 'att4': '{{ task.task_id }}', 'template_fields': ['att3']}), "
"'template_fields': ['nested1']})",
),
+ (
+ "a" * 5000,
+ f"Truncated. You can change this behaviour in [core]max_templated_field_length. {('a'*5000)[:max_length-79]!r}... ",
+ ),
+ (
+ LargeStrObject(),
+ f"Truncated. You can change this behaviour in [core]max_templated_field_length. {str(LargeStrObject())[:max_length-79]!r}... ",
+ ),
],
)
def test_get_templated_fields(self, templated_field, expected_rendered_field, dag_maker):
@@ -148,6 +169,45 @@ class TestRenderedTaskInstanceFields:
# Fetching them will return None
assert RTIF.get_templated_fields(ti=ti2) is None
+ def test_secrets_are_masked_when_large_string(self, dag_maker):
+ """
+ Test that secrets are masked when the templated field is a large string
+ """
+ Variable.set(
+ key="api_key",
+ value="test api key are still masked" * 5000,
+ )
+ with dag_maker("test_serialized_rendered_fields"):
+ task = BashOperator(task_id="test", bash_command="echo {{ var.value.api_key }}")
+ dr = dag_maker.create_dagrun()
+ ti = dr.task_instances[0]
+ ti.task = task
+ rtif = RTIF(ti=ti)
+ assert "***" in rtif.rendered_fields.get("bash_command")
+
+ @mock.patch("airflow.models.BaseOperator.render_template")
+ def test_pandas_dataframes_works_with_the_string_compare(self, render_mock, dag_maker):
+ """Test that rendered dataframe gets passed through the serialized template fields."""
+ import pandas
+
+ render_mock.return_value = pandas.DataFrame({"a": [1, 2, 3]})
+ with dag_maker("test_serialized_rendered_fields"):
+
+ @task_decorator
+ def generate_pd():
+ return pandas.DataFrame({"a": [1, 2, 3]})
+
+ @task_decorator
+ def consume_pd(data):
+ return data
+
+ consume_pd(generate_pd())
+
+ dr = dag_maker.create_dagrun()
+ ti, ti2 = dr.task_instances
+ rtif = RTIF(ti=ti2)
+ rtif.write()
+
@pytest.mark.parametrize(
"rtif_num, num_to_keep, remaining_rtifs, expected_query_count",
[