[airflow] branch main updated: Make native environment Airflow-flavoured like sandbox (#20704)

[po...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

potiuk 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 3700b3a  Make native environment Airflow-flavoured like sandbox (#20704)
3700b3a is described below

commit 3700b3a918986c9caa2885f69f023d251b2881cd
Author: Malthe Borch <mb...@gmail.com>
AuthorDate: Thu Jan 6 11:24:45 2022 +0100

    Make native environment Airflow-flavoured like sandbox (#20704)
---
 airflow/models/dag.py |  3 +--
 airflow/templates.py  | 13 ++++++++++---
 2 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/airflow/models/dag.py b/airflow/models/dag.py
index 136ed04..9157bae 100644
--- a/airflow/models/dag.py
+++ b/airflow/models/dag.py
@@ -50,7 +50,6 @@ from typing import (
 import jinja2
 import pendulum
 from dateutil.relativedelta import relativedelta
-from jinja2.nativetypes import NativeEnvironment
 from pendulum.tz.timezone import Timezone
 from sqlalchemy import Boolean, Column, ForeignKey, Index, Integer, String, Text, func, or_
 from sqlalchemy.orm import backref, joinedload, relationship
@@ -1287,7 +1286,7 @@ class DAG(LoggingMixin):
         if self.jinja_environment_kwargs:
             jinja_env_options.update(self.jinja_environment_kwargs)
         if self.render_template_as_native_obj:
-            env_class: Any = NativeEnvironment
+            env_class = airflow.templates.NativeEnvironment
         else:
             env_class = airflow.templates.SandboxedEnvironment
         env: jinja2.Environment = env_class(**jinja_env_options)
diff --git a/airflow/templates.py b/airflow/templates.py
index 690951e..6ec010f 100644
--- a/airflow/templates.py
+++ b/airflow/templates.py
@@ -16,12 +16,11 @@
 # specific language governing permissions and limitations
 # under the License.
 
+import jinja2.nativetypes
 import jinja2.sandbox
 
 
-class SandboxedEnvironment(jinja2.sandbox.SandboxedEnvironment):
-    """SandboxedEnvironment for Airflow task templates."""
-
+class _AirflowEnvironmentMixin:
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
 
@@ -37,6 +36,14 @@ class SandboxedEnvironment(jinja2.sandbox.SandboxedEnvironment):
         return not jinja2.sandbox.is_internal_attribute(obj, attr)
 
 
+class NativeEnvironment(_AirflowEnvironmentMixin, jinja2.nativetypes.NativeEnvironment):
+    """NativeEnvironment for Airflow task templates."""
+
+
+class SandboxedEnvironment(_AirflowEnvironmentMixin, jinja2.sandbox.SandboxedEnvironment):
+    """SandboxedEnvironment for Airflow task templates."""
+
+
 def ds_filter(value):
     return value.strftime('%Y-%m-%d')