Repository: apache/airflow

Issue:
Create UndefinedJinjaVariablesRule to ease upgrade to 2.0
This issue is part of #8765

## Rule
Create `UndefinedJinjaVariablesRule` which corresponds to
> Prior to Airflow 2.0 Jinja Templates would permit the use of undefined variables. They would render as an
empty string, with no indication to the user an undefined variable was used. With this release, any template
rendering involving undefined variables will fail the task, as well as displaying an error in the UI when
rendering.

entry in UPDATING.md. Introduced in #11016 
This rule should check if any of the jinja templates use undefined variables. This might not be always possible, as this might involve the dynamic generation of some context variables, but likely we can at least print some warning or "double-check this" statements.

## How to guide
To implement a new rule, create a class that inherits from `airflow.upgrade.rules.base_rule.BaseRule`.
It will be auto-registered and used by `airflow upgrade-check` command. The custom rule class has to have `title`,
`description` properties and should implement `check` method which returns a list of error messages in case of
incompatibility.

For example:
https://github.com/apache/airflow/blob/ea36166961ca35fc51ddc262ec245590c3e236fb/airflow/upgrade/rules/conn_type_is_not_nullable.py#L25-L42

**Remember** to open the PR against `v1-10-test` branch.



Discussion:
Hi, I am new to both Airflow and contributing to FOSS.
Saw this issue unpicked and I would like to pick this up.

I had one quick query - 
The base branch from which I create a new feature branch in my fork would be `v1-10-test`?

Thanks! Looking forward to contribute and learn more! 
Yep.
I had a few doubts asking them here - 

Currently I was writing the following logic for `check` - 
For a given dag - 
1. It will gather all the tasks - `dag.tasks`
2. Iterate over the task list and for every task
        - Get Jinja2 environment from `task.get_template_env()`
        - Find Jinja2 variables for all `task.template_fields`
        - Verify that all the above gathered variables are present in `task.params`
3. Return a list of messages for every missing variable in `task.params`

Below is the script version

```
            for task in dag.tasks:
                env = task.get_template_env()
                params = set(task.params.keys())
                for attr_name in task.template_fields:
                    content = getattr(task, attr_name)
                    parsed = env.parse(content)
                    template_varaibles = meta.find_undeclared_variables(parsed) # Returns a set of variables present in the template
                    undefined_variables = template_varaibles-params # Set difference
```

Does this logic look like something we could follow?
And if yes then how can I write this logic to check for every DAG? 


I did have a look at `DagBag` but I do not know how should I write a test for it.
Sorry for the newbie questions - would be really helpful if you could guide on these points.
@ashmeet13 That check looks like a good start yes, but not quite there. I don't think your `params` is right. `TaskInstance.get_template_context()` is the parameters that templates are executed in, and for that you will need to create a TaskInstance (it can be in memory, no need to write it to the DB)
Also check our BaseOperator.render_template for all the possible locations/field types we render templates on. You'll need to replicate much of that logic to handle all the types (inside lists, dicts, etc. etc.)
Got it - thank you so much. Will update on the progress or in case of any further doubts!
Some progress that I have made. Thanks @ashb  for guiding towards `TaskInstance` for getting the complete context.
Also towards `BaseOperator.render_template`

With this I was able to improve the logic to extract undefined variables in a template to this - 

```
def check_rendered(rendered):
    if isinstance(rendered, six.string_types):
        return set(re.findall(r"{{(.*?)}}", rendered))

    elif isinstance(rendered, (tuple,list,set)):
        parsed_templates = set()
        for element in rendered:
            parsed_templates.union(check_rendered(element))
        return parsed_templates

    elif isinstance(rendered, dict):
        parsed_templates = set()
        for key, value in rendered.items():
            parsed_templates.union(check_rendered(value))
        return parsed_templates 



def check(dag):
    dag.template_undefined = jinja2.DebugUndefined
    for task in dag.tasks:
        jinja_env = task.get_template_env()
        task_instance = TaskInstance(task=task, execution_date=timezone.utcnow())
        template_context = task_instance.get_template_context()

        for attr_name in task.template_fields:
            content = getattr(task, attr_name)
            if content:
                rendered = task.render_template(content, template_context)
                undefined = check_rendered(rendered)
                for element in undefined:
                    print(element) 
```

Take into the consideration the following DAG -
```
dag = DAG(
    'tutorial',
    default_args=default_args,
    description='A simple tutorial DAG',
    schedule_interval=timedelta(days=1),
)

templated_command_string = """
{% for i in range(5) %}
    echo "{{ params.defined }}"
    echo "{{ execution_date.month }}"
    echo "{{ params.undefined }}"
    echo "{{ foo }}"
{% endfor %}
"""

t1 = BashOperator(
    task_id='templated_string',
    depends_on_past=False,
    bash_command=templated_command_string,
    params={'defined': 'test_val'},
    dag=dag,
)
```

The above check would print/return this -  
```
foo
no such element: dict object['undefined'] 
```


One blocker still that I have is how exactly will I be able to get the dags when the user runs `airflow upgrade-check`
Thanks again for the help!


Edited: Adding the [Stack Overflow](https://stackoverflow.com/questions/46619830/how-to-get-all-undefined-variables-from-a-jinja2-template) source I used

Resolving patch:
diff --git a/airflow/models/dag.py b/airflow/models/dag.py
--- a/airflow/models/dag.py
+++ b/airflow/models/dag.py
@@ -224,7 +224,7 @@ def __init__(
         end_date=None,  # type: Optional[datetime]
         full_filepath=None,  # type: Optional[str]
         template_searchpath=None,  # type: Optional[Union[str, Iterable[str]]]
-        template_undefined=jinja2.Undefined,  # type: Type[jinja2.Undefined]
+        template_undefined=None,  # type: Optional[Type[jinja2.Undefined]]
         user_defined_macros=None,  # type: Optional[Dict]
         user_defined_filters=None,  # type: Optional[Dict]
         default_args=None,  # type: Optional[Dict]
@@ -807,7 +807,7 @@ def get_template_env(self):  # type: () -> jinja2.Environment
         # Default values (for backward compatibility)
         jinja_env_options = {
             'loader': jinja2.FileSystemLoader(searchpath),
-            'undefined': self.template_undefined,
+            'undefined': self.template_undefined or jinja2.Undefined,
             'extensions': ["jinja2.ext.do"],
             'cache_size': 0
         }
diff --git a/airflow/upgrade/rules/undefined_jinja_varaibles.py b/airflow/upgrade/rules/undefined_jinja_varaibles.py
new file mode 100644
--- /dev/null
+++ b/airflow/upgrade/rules/undefined_jinja_varaibles.py
@@ -0,0 +1,153 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import absolute_import
+
+import re
+
+import jinja2
+import six
+
+from airflow import conf
+from airflow.models import DagBag, TaskInstance
+from airflow.upgrade.rules.base_rule import BaseRule
+from airflow.utils import timezone
+
+
+class UndefinedJinjaVariablesRule(BaseRule):
+
+    title = "Jinja Template Variables cannot be undefined"
+
+    description = """\
+The default behavior for DAG's Jinja templates has changed. Now, more restrictive validation
+of non-existent variables is applied - `jinja2.StrictUndefined`.
+
+The user should do either of the following to fix this -
+1. Fix the Jinja Templates by defining every variable or providing default values
+2. Explicitly declare `template_undefined=jinja2.Undefined` while defining the DAG
+"""
+
+    def _check_rendered_content(self, rendered_content, seen_oids=None):
+        """Replicates the logic in BaseOperator.render_template() to
+        cover all the cases needed to be checked.
+        """
+        if isinstance(rendered_content, six.string_types):
+            return set(re.findall(r"{{(.*?)}}", rendered_content))
+
+        elif isinstance(rendered_content, (int, float, bool)):
+            return set()
+
+        elif isinstance(rendered_content, (tuple, list, set)):
+            debug_error_messages = set()
+            for element in rendered_content:
+                debug_error_messages.update(self._check_rendered_content(element))
+            return debug_error_messages
+
+        elif isinstance(rendered_content, dict):
+            debug_error_messages = set()
+            for key, value in rendered_content.items():
+                debug_error_messages.update(self._check_rendered_content(value))
+            return debug_error_messages
+
+        else:
+            if seen_oids is None:
+                seen_oids = set()
+            return self._nested_check_rendered(rendered_content, seen_oids)
+
+    def _nested_check_rendered(self, rendered_content, seen_oids):
+        debug_error_messages = set()
+        if id(rendered_content) not in seen_oids:
+            seen_oids.add(id(rendered_content))
+            nested_template_fields = rendered_content.template_fields
+            for attr_name in nested_template_fields:
+                nested_rendered_content = getattr(rendered_content, attr_name)
+
+                if nested_rendered_content:
+                    errors = list(
+                        self._check_rendered_content(nested_rendered_content, seen_oids)
+                    )
+                    for i in range(len(errors)):
+                        errors[i].strip()
+                        errors[i] += " NestedTemplateField={}".format(attr_name)
+                    debug_error_messages.update(errors)
+        return debug_error_messages
+
+    def _render_task_content(self, task, content, context):
+        completed_rendering = False
+        errors_while_rendering = []
+        while not completed_rendering:
+            # Catch errors such as {{ object.element }} where
+            # object is not defined
+            try:
+                renderend_content = task.render_template(content, context)
+                completed_rendering = True
+            except Exception as e:
+                undefined_variable = re.sub(" is undefined", "", str(e))
+                undefined_variable = re.sub("'", "", undefined_variable)
+                context[undefined_variable] = dict()
+                message = "Could not find the object '{}'".format(undefined_variable)
+                errors_while_rendering.append(message)
+        return renderend_content, errors_while_rendering
+
+    def iterate_over_template_fields(self, task):
+        messages = {}
+        task_instance = TaskInstance(task=task, execution_date=timezone.utcnow())
+        context = task_instance.get_template_context()
+        for attr_name in task.template_fields:
+            content = getattr(task, attr_name)
+            if content:
+                rendered_content, errors_while_rendering = self._render_task_content(
+                    task, content, context
+                )
+                debug_error_messages = list(
+                    self._check_rendered_content(rendered_content, set())
+                )
+                messages[attr_name] = errors_while_rendering + debug_error_messages
+
+        return messages
+
+    def iterate_over_dag_tasks(self, dag):
+        dag.template_undefined = jinja2.DebugUndefined
+        tasks = dag.tasks
+        messages = {}
+        for task in tasks:
+            error_messages = self.iterate_over_template_fields(task)
+            messages[task.task_id] = error_messages
+        return messages
+
+    def check(self, dagbag=None):
+        if not dagbag:
+            dag_folder = conf.get("core", "dags_folder")
+            dagbag = DagBag(dag_folder)
+        dags = dagbag.dags
+        messages = []
+        for dag_id, dag in dags.items():
+            if dag.template_undefined:
+                continue
+            dag_messages = self.iterate_over_dag_tasks(dag)
+
+            for task_id, task_messages in dag_messages.items():
+                for attr_name, error_messages in task_messages.items():
+                    for error_message in error_messages:
+                        message = (
+                            "Possible UndefinedJinjaVariable -> DAG: {}, Task: {}, "
+                            "Attribute: {}, Error: {}".format(
+                                dag_id, task_id, attr_name, error_message.strip()
+                            )
+                        )
+                        messages.append(message)
+        return messages


Test patch:
diff --git a/tests/upgrade/rules/test_undefined_jinja_varaibles.py b/tests/upgrade/rules/test_undefined_jinja_varaibles.py
new file mode 100644
--- /dev/null
+++ b/tests/upgrade/rules/test_undefined_jinja_varaibles.py
@@ -0,0 +1,192 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from tempfile import mkdtemp
+from unittest import TestCase
+
+import jinja2
+
+from airflow import DAG
+from airflow.models import DagBag
+from airflow.operators.bash_operator import BashOperator
+from airflow.upgrade.rules.undefined_jinja_varaibles import UndefinedJinjaVariablesRule
+from tests.models import DEFAULT_DATE
+
+
+class ClassWithCustomAttributes:
+    """Class for testing purpose: allows to create objects with custom attributes in one single statement."""
+
+    def __init__(self, **kwargs):
+        for key, value in kwargs.items():
+            setattr(self, key, value)
+
+    def __str__(self):
+        return "{}({})".format(ClassWithCustomAttributes.__name__, str(self.__dict__))
+
+    def __repr__(self):
+        return self.__str__()
+
+    def __eq__(self, other):
+        return self.__dict__ == other.__dict__
+
+    def __ne__(self, other):
+        return not self.__eq__(other)
+
+
+class TestUndefinedJinjaVariablesRule(TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.empty_dir = mkdtemp()
+
+    def setUpValidDag(self):
+        self.valid_dag = DAG(
+            dag_id="test-defined-jinja-variables", start_date=DEFAULT_DATE
+        )
+
+        BashOperator(
+            task_id="templated_string",
+            depends_on_past=False,
+            bash_command="echo",
+            env={
+                "integer": "{{ params.integer }}",
+                "float": "{{ params.float }}",
+                "string": "{{ params.string }}",
+                "boolean": "{{ params.boolean }}",
+            },
+            params={
+                "integer": 1,
+                "float": 1.0,
+                "string": "test_string",
+                "boolean": True,
+            },
+            dag=self.valid_dag,
+        )
+
+    def setUpDagToSkip(self):
+        self.skip_dag = DAG(
+            dag_id="test-defined-jinja-variables",
+            start_date=DEFAULT_DATE,
+            template_undefined=jinja2.Undefined,
+        )
+
+        BashOperator(
+            task_id="templated_string",
+            depends_on_past=False,
+            bash_command="{{ undefined }}",
+            dag=self.skip_dag,
+        )
+
+    def setUpInvalidDag(self):
+        self.invalid_dag = DAG(
+            dag_id="test-undefined-jinja-variables", start_date=DEFAULT_DATE
+        )
+
+        invalid_template_command = """
+            {% for i in range(5) %}
+                echo "{{ params.defined_variable }}"
+                echo "{{ execution_date.today }}"
+                echo "{{ execution_date.invalid_element }}"
+                echo "{{ params.undefined_variable }}"
+                echo "{{ foo }}"
+            {% endfor %}
+            """
+
+        nested_validation = ClassWithCustomAttributes(
+            nested1=ClassWithCustomAttributes(
+                att1="{{ nested.undefined }}", template_fields=["att1"]
+            ),
+            nested2=ClassWithCustomAttributes(
+                att2="{{ bar }}", template_fields=["att2"]
+            ),
+            template_fields=["nested1", "nested2"],
+        )
+
+        BashOperator(
+            task_id="templated_string",
+            depends_on_past=False,
+            bash_command=invalid_template_command,
+            env={
+                "undefined_object": "{{ undefined_object.element }}",
+                "nested_object": nested_validation,
+            },
+            params={"defined_variable": "defined_value"},
+            dag=self.invalid_dag,
+        )
+
+    def setUp(self):
+        self.setUpValidDag()
+        self.setUpDagToSkip()
+        self.setUpInvalidDag()
+
+    def test_description_and_title_is_defined(self):
+        rule = UndefinedJinjaVariablesRule()
+        assert isinstance(rule.description, str)
+        assert isinstance(rule.title, str)
+
+    def test_valid_check(self):
+        dagbag = DagBag(dag_folder=self.empty_dir, include_examples=False)
+        dagbag.dags[self.valid_dag.dag_id] = self.valid_dag
+        rule = UndefinedJinjaVariablesRule()
+
+        messages = rule.check(dagbag)
+
+        assert len(messages) == 0
+
+    def test_skipping_dag_check(self):
+        dagbag = DagBag(dag_folder=self.empty_dir, include_examples=False)
+        dagbag.dags[self.skip_dag.dag_id] = self.skip_dag
+        rule = UndefinedJinjaVariablesRule()
+
+        messages = rule.check(dagbag)
+
+        assert len(messages) == 0
+
+    def test_invalid_check(self):
+        dagbag = DagBag(dag_folder=self.empty_dir, include_examples=False)
+        dagbag.dags[self.invalid_dag.dag_id] = self.invalid_dag
+        rule = UndefinedJinjaVariablesRule()
+
+        messages = rule.check(dagbag)
+
+        expected_messages = [
+            "Possible UndefinedJinjaVariable -> DAG: test-undefined-jinja-variables, "
+            "Task: templated_string, Attribute: bash_command, Error: no such element: "
+            "dict object['undefined_variable']",
+            "Possible UndefinedJinjaVariable -> DAG: test-undefined-jinja-variables, "
+            "Task: templated_string, Attribute: bash_command, Error: no such element: "
+            "pendulum.pendulum.Pendulum object['invalid_element']",
+            "Possible UndefinedJinjaVariable -> DAG: test-undefined-jinja-variables, "
+            "Task: templated_string, Attribute: bash_command, Error: foo",
+            "Possible UndefinedJinjaVariable -> DAG: test-undefined-jinja-variables, "
+            "Task: templated_string, Attribute: env, Error: Could not find the "
+            "object 'undefined_object",
+            "Possible UndefinedJinjaVariable -> DAG: test-undefined-jinja-variables, "
+            "Task: templated_string, Attribute: env, Error: Could not find the object 'nested'",
+            "Possible UndefinedJinjaVariable -> DAG: test-undefined-jinja-variables, "
+            "Task: templated_string, Attribute: env, Error: bar  NestedTemplateField=att2 "
+            "NestedTemplateField=nested2",
+            "Possible UndefinedJinjaVariable -> DAG: test-undefined-jinja-variables, "
+            "Task: templated_string, Attribute: env, Error: no such element: "
+            "dict object['undefined']  NestedTemplateField=att1 NestedTemplateField=nested1",
+            "Possible UndefinedJinjaVariable -> DAG: test-undefined-jinja-variables, "
+            "Task: templated_string, Attribute: env, Error: no such element: dict object['element']",
+        ]
+
+        assert len(messages) == len(expected_messages)
+        assert [m for m in messages if m in expected_messages], len(messages) == len(
+            expected_messages
+        )