From 057258e3ec3f585eb7afb6dd2857e9bf5e262103 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?G=C3=A9raud?= Date: Wed, 23 Jan 2019 08:02:04 -0500 Subject: [PATCH] [AIRFLOW-2508] Handle non string types in Operators templatized fields (#4292) --- airflow/models.py | 11 +---- tests/models.py | 117 ++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 119 insertions(+), 9 deletions(-) diff --git a/airflow/models.py b/airflow/models.py index e7d39335f191e..b837438148701 100755 --- a/airflow/models.py +++ b/airflow/models.py @@ -39,7 +39,6 @@ import jinja2 import json import logging -import numbers import os import pendulum import pickle @@ -2879,25 +2878,19 @@ def render_template_from_field(self, attr, content, context, jinja_env): Renders a template from a field. If the field is a string, it will simply render the string and return the result. If it is a collection or nested set of collections, it will traverse the structure and render - all strings in it. + all elements in it. If the field has another type, it will return it as it is. """ rt = self.render_template if isinstance(content, six.string_types): result = jinja_env.from_string(content).render(**context) elif isinstance(content, (list, tuple)): result = [rt(attr, e, context) for e in content] - elif isinstance(content, numbers.Number): - result = content elif isinstance(content, dict): result = { k: rt("{}[{}]".format(attr, k), v, context) for k, v in list(content.items())} else: - param_type = type(content) - msg = ( - "Type '{param_type}' used for parameter '{attr}' is " - "not supported for templating").format(**locals()) - raise AirflowException(msg) + result = content return result def render_template(self, attr, content, context): diff --git a/tests/models.py b/tests/models.py index d45c0b870e780..6c60ef956d0fd 100644 --- a/tests/models.py +++ b/tests/models.py @@ -31,6 +31,7 @@ import time import unittest import urllib +import uuid from tempfile import NamedTemporaryFile, mkdtemp import pendulum @@ -395,6 +396,122 @@ def test_render_template_field(self): result = task.render_template('', '{{ foo }}', dict(foo='bar')) self.assertEqual(result, 'bar') + def test_render_template_list_field(self): + """Tests if render_template from a list field works""" + + dag = DAG('test-dag', + start_date=DEFAULT_DATE) + + with dag: + task = DummyOperator(task_id='op1') + + self.assertListEqual( + task.render_template('', ['{{ foo }}_1', '{{ foo }}_2'], {'foo': 'bar'}), + ['bar_1', 'bar_2'] + ) + + def test_render_template_tuple_field(self): + """Tests if render_template from a tuple field works""" + + dag = DAG('test-dag', + start_date=DEFAULT_DATE) + + with dag: + task = DummyOperator(task_id='op1') + + # tuple is replaced by a list + self.assertListEqual( + task.render_template('', ('{{ foo }}_1', '{{ foo }}_2'), {'foo': 'bar'}), + ['bar_1', 'bar_2'] + ) + + def test_render_template_dict_field(self): + """Tests if render_template from a dict field works""" + + dag = DAG('test-dag', + start_date=DEFAULT_DATE) + + with dag: + task = DummyOperator(task_id='op1') + + self.assertDictEqual( + task.render_template('', {'key1': '{{ foo }}_1', 'key2': '{{ foo }}_2'}, {'foo': 'bar'}), + {'key1': 'bar_1', 'key2': 'bar_2'} + ) + + def test_render_template_dict_field_with_templated_keys(self): + """Tests if render_template from a dict field works as expected: + dictionary keys are not templated""" + + dag = DAG('test-dag', + start_date=DEFAULT_DATE) + + with dag: + task = DummyOperator(task_id='op1') + + self.assertDictEqual( + task.render_template('', {'key_{{ foo }}_1': 1, 'key_2': '{{ foo }}_2'}, {'foo': 'bar'}), + {'key_{{ foo }}_1': 1, 'key_2': 'bar_2'} + ) + + def test_render_template_date_field(self): + """Tests if render_template from a date field works""" + + dag = DAG('test-dag', + start_date=DEFAULT_DATE) + + with dag: + task = DummyOperator(task_id='op1') + + self.assertEqual( + task.render_template('', datetime.date(2018, 12, 6), {'foo': 'bar'}), + datetime.date(2018, 12, 6) + ) + + def test_render_template_datetime_field(self): + """Tests if render_template from a datetime field works""" + + dag = DAG('test-dag', + start_date=DEFAULT_DATE) + + with dag: + task = DummyOperator(task_id='op1') + + self.assertEqual( + task.render_template('', datetime.datetime(2018, 12, 6, 10, 55), {'foo': 'bar'}), + datetime.datetime(2018, 12, 6, 10, 55) + ) + + def test_render_template_UUID_field(self): + """Tests if render_template from a UUID field works""" + + dag = DAG('test-dag', + start_date=DEFAULT_DATE) + + with dag: + task = DummyOperator(task_id='op1') + + random_uuid = uuid.uuid4() + self.assertIs( + task.render_template('', random_uuid, {'foo': 'bar'}), + random_uuid + ) + + def test_render_template_object_field(self): + """Tests if render_template from an object field works""" + + dag = DAG('test-dag', + start_date=DEFAULT_DATE) + + with dag: + task = DummyOperator(task_id='op1') + + test_object = object() + self.assertIs( + task.render_template('', test_object, {'foo': 'bar'}), + test_object + ) + def test_render_template_field_macro(self): """ Tests if render_template from a field works, if a custom filter was defined"""