Coverage for /pythoncovmergedfiles/medio/medio/src/airflow/airflow/template/templater.py: 23%

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)