1# Licensed to the Apache Software Foundation (ASF) under one
2# or more contributor license agreements. See the NOTICE file
3# distributed with this work for additional information
4# regarding copyright ownership. The ASF licenses this file
5# to you under the Apache License, Version 2.0 (the
6# "License"); you may not use this file except in compliance
7# with the License. You may obtain a copy of the License at
8#
9# http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing,
12# software distributed under the License is distributed on an
13# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14# KIND, either express or implied. See the License for the
15# specific language governing permissions and limitations
16# under the License.
17
18from __future__ import annotations
19
20from dataclasses import dataclass
21from typing import TYPE_CHECKING, Any, Collection, Iterable, Sequence
22
23from airflow.utils.helpers import render_template_as_native, render_template_to_string
24from airflow.utils.log.logging_mixin import LoggingMixin
25from airflow.utils.mixins import ResolveMixin
26
27if TYPE_CHECKING:
28 import jinja2
29
30 from airflow import DAG
31 from airflow.models.operator import Operator
32 from airflow.utils.context import Context
33
34
35@dataclass(frozen=True)
36class LiteralValue(ResolveMixin):
37 """
38 A wrapper for a value that should be rendered as-is, without applying jinja templating to its contents.
39
40 :param value: The value to be rendered without templating
41 """
42
43 value: Any
44
45 def iter_references(self) -> Iterable[tuple[Operator, str]]:
46 return ()
47
48 def resolve(self, context: Context) -> Any:
49 return self.value
50
51
52class Templater(LoggingMixin):
53 """
54 This renders the template fields of object.
55
56 :meta private:
57 """
58
59 # For derived classes to define which fields will get jinjaified.
60 template_fields: Collection[str]
61 # Defines which files extensions to look for in the templated fields.
62 template_ext: Sequence[str]
63
64 def get_template_env(self, dag: DAG | None = None) -> jinja2.Environment:
65 """Fetch a Jinja template environment from the DAG or instantiate empty environment if no DAG."""
66 # This is imported locally since Jinja2 is heavy and we don't need it
67 # for most of the functionalities. It is imported by get_template_env()
68 # though, so we don't need to put this after the 'if dag' check.
69 from airflow.templates import SandboxedEnvironment
70
71 if dag:
72 return dag.get_template_env(force_sandboxed=False)
73 return SandboxedEnvironment(cache_size=0)
74
75 def prepare_template(self) -> None:
76 """
77 Execute after the templated fields get replaced by their content.
78
79 If you need your object to alter the content of the file before the
80 template is rendered, it should override this method to do so.
81 """
82
83 def resolve_template_files(self) -> None:
84 """Get the content of files for template_field / template_ext."""
85 if self.template_ext:
86 for field in self.template_fields:
87 content = getattr(self, field, None)
88 if isinstance(content, str) and content.endswith(tuple(self.template_ext)):
89 env = self.get_template_env()
90 try:
91 setattr(self, field, env.loader.get_source(env, content)[0]) # type: ignore
92 except Exception:
93 self.log.exception("Failed to resolve template field %r", field)
94 elif isinstance(content, list):
95 env = self.get_template_env()
96 for i, item in enumerate(content):
97 if isinstance(item, str) and item.endswith(tuple(self.template_ext)):
98 try:
99 content[i] = env.loader.get_source(env, item)[0] # type: ignore
100 except Exception:
101 self.log.exception("Failed to get source %s", item)
102 self.prepare_template()
103
104 def _do_render_template_fields(
105 self,
106 parent: Any,
107 template_fields: Iterable[str],
108 context: Context,
109 jinja_env: jinja2.Environment,
110 seen_oids: set[int],
111 ) -> None:
112 for attr_name in template_fields:
113 value = getattr(parent, attr_name)
114 rendered_content = self.render_template(
115 value,
116 context,
117 jinja_env,
118 seen_oids,
119 )
120 if rendered_content:
121 setattr(parent, attr_name, rendered_content)
122
123 def _render(self, template, context, dag: DAG | None = None) -> Any:
124 if dag and dag.render_template_as_native_obj:
125 return render_template_as_native(template, context)
126 return render_template_to_string(template, context)
127
128 def render_template(
129 self,
130 content: Any,
131 context: Context,
132 jinja_env: jinja2.Environment | None = None,
133 seen_oids: set[int] | None = None,
134 ) -> Any:
135 """Render a templated string.
136
137 If *content* is a collection holding multiple templated strings, strings
138 in the collection will be templated recursively.
139
140 :param content: Content to template. Only strings can be templated (may
141 be inside a collection).
142 :param context: Dict with values to apply on templated content
143 :param jinja_env: Jinja environment. Can be provided to avoid
144 re-creating Jinja environments during recursion.
145 :param seen_oids: template fields already rendered (to avoid
146 *RecursionError* on circular dependencies)
147 :return: Templated content
148 """
149 # "content" is a bad name, but we're stuck to it being public API.
150 value = content
151 del content
152
153 if seen_oids is not None:
154 oids = seen_oids
155 else:
156 oids = set()
157
158 if id(value) in oids:
159 return value
160
161 if not jinja_env:
162 jinja_env = self.get_template_env()
163
164 if isinstance(value, str):
165 if value.endswith(tuple(self.template_ext)): # A filepath.
166 template = jinja_env.get_template(value)
167 else:
168 template = jinja_env.from_string(value)
169 return self._render(template, context)
170 if isinstance(value, ResolveMixin):
171 return value.resolve(context)
172
173 # Fast path for common built-in collections.
174 if value.__class__ is tuple:
175 return tuple(self.render_template(element, context, jinja_env, oids) for element in value)
176 elif isinstance(value, tuple): # Special case for named tuples.
177 return value.__class__(*(self.render_template(el, context, jinja_env, oids) for el in value))
178 elif isinstance(value, list):
179 return [self.render_template(element, context, jinja_env, oids) for element in value]
180 elif isinstance(value, dict):
181 return {k: self.render_template(v, context, jinja_env, oids) for k, v in value.items()}
182 elif isinstance(value, set):
183 return {self.render_template(element, context, jinja_env, oids) for element in value}
184
185 # More complex collections.
186 self._render_nested_template_fields(value, context, jinja_env, oids)
187 return value
188
189 def _render_nested_template_fields(
190 self,
191 value: Any,
192 context: Context,
193 jinja_env: jinja2.Environment,
194 seen_oids: set[int],
195 ) -> None:
196 if id(value) in seen_oids:
197 return
198 seen_oids.add(id(value))
199 try:
200 nested_template_fields = value.template_fields
201 except AttributeError:
202 # content has no inner template fields
203 return
204 self._do_render_template_fields(value, nested_template_fields, context, jinja_env, seen_oids)