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

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 typing import TYPE_CHECKING, Any, Collection, Iterable, Sequence 

21 

22from airflow.utils.context import Context 

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 

26from airflow.utils.session import NEW_SESSION, provide_session 

27 

28if TYPE_CHECKING: 

29 import jinja2 

30 from sqlalchemy.orm import Session 

31 

32 from airflow import DAG 

33 

34 

35class Templater(LoggingMixin): 

36 """ 

37 This renders the template fields of object. 

38 

39 :meta private: 

40 """ 

41 

42 # For derived classes to define which fields will get jinjaified. 

43 template_fields: Collection[str] 

44 # Defines which files extensions to look for in the templated fields. 

45 template_ext: Sequence[str] 

46 

47 def get_template_env(self, dag: DAG | None = None) -> jinja2.Environment: 

48 """Fetch a Jinja template environment from the DAG or instantiate empty environment if no DAG.""" 

49 # This is imported locally since Jinja2 is heavy and we don't need it 

50 # for most of the functionalities. It is imported by get_template_env() 

51 # though, so we don't need to put this after the 'if dag' check. 

52 from airflow.templates import SandboxedEnvironment 

53 

54 if dag: 

55 return dag.get_template_env(force_sandboxed=False) 

56 return SandboxedEnvironment(cache_size=0) 

57 

58 def prepare_template(self) -> None: 

59 """Hook triggered after the templated fields get replaced by their content. 

60 

61 If you need your object to alter the content of the file before the 

62 template is rendered, it should override this method to do so. 

63 """ 

64 

65 def resolve_template_files(self) -> None: 

66 """Getting the content of files for template_field / template_ext.""" 

67 if self.template_ext: 

68 for field in self.template_fields: 

69 content = getattr(self, field, None) 

70 if content is None: 

71 continue 

72 elif isinstance(content, str) and any(content.endswith(ext) for ext in self.template_ext): 

73 env = self.get_template_env() 

74 try: 

75 setattr(self, field, env.loader.get_source(env, content)[0]) # type: ignore 

76 except Exception: 

77 self.log.exception("Failed to resolve template field %r", field) 

78 elif isinstance(content, list): 

79 env = self.get_template_env() 

80 for i, item in enumerate(content): 

81 if isinstance(item, str) and any(item.endswith(ext) for ext in self.template_ext): 

82 try: 

83 content[i] = env.loader.get_source(env, item)[0] # type: ignore 

84 except Exception: 

85 self.log.exception("Failed to get source %s", item) 

86 self.prepare_template() 

87 

88 @provide_session 

89 def _do_render_template_fields( 

90 self, 

91 parent: Any, 

92 template_fields: Iterable[str], 

93 context: Context, 

94 jinja_env: jinja2.Environment, 

95 seen_oids: set[int], 

96 *, 

97 session: Session = NEW_SESSION, 

98 ) -> None: 

99 for attr_name in template_fields: 

100 value = getattr(parent, attr_name) 

101 rendered_content = self.render_template( 

102 value, 

103 context, 

104 jinja_env, 

105 seen_oids, 

106 ) 

107 if rendered_content: 

108 setattr(parent, attr_name, rendered_content) 

109 

110 def _render(self, template, context, dag: DAG | None = None) -> Any: 

111 if dag and dag.render_template_as_native_obj: 

112 return render_template_as_native(template, context) 

113 return render_template_to_string(template, context) 

114 

115 def render_template( 

116 self, 

117 content: Any, 

118 context: Context, 

119 jinja_env: jinja2.Environment | None = None, 

120 seen_oids: set[int] | None = None, 

121 ) -> Any: 

122 """Render a templated string. 

123 

124 If *content* is a collection holding multiple templated strings, strings 

125 in the collection will be templated recursively. 

126 

127 :param content: Content to template. Only strings can be templated (may 

128 be inside a collection). 

129 :param context: Dict with values to apply on templated content 

130 :param jinja_env: Jinja environment. Can be provided to avoid 

131 re-creating Jinja environments during recursion. 

132 :param seen_oids: template fields already rendered (to avoid 

133 *RecursionError* on circular dependencies) 

134 :return: Templated content 

135 """ 

136 # "content" is a bad name, but we're stuck to it being public API. 

137 value = content 

138 del content 

139 

140 if seen_oids is not None: 

141 oids = seen_oids 

142 else: 

143 oids = set() 

144 

145 if id(value) in oids: 

146 return value 

147 

148 if not jinja_env: 

149 jinja_env = self.get_template_env() 

150 

151 if isinstance(value, str): 

152 if any(value.endswith(ext) for ext in self.template_ext): # A filepath. 

153 template = jinja_env.get_template(value) 

154 else: 

155 template = jinja_env.from_string(value) 

156 return self._render(template, context) 

157 if isinstance(value, ResolveMixin): 

158 return value.resolve(context) 

159 

160 # Fast path for common built-in collections. 

161 if value.__class__ is tuple: 

162 return tuple(self.render_template(element, context, jinja_env, oids) for element in value) 

163 elif isinstance(value, tuple): # Special case for named tuples. 

164 return value.__class__(*(self.render_template(el, context, jinja_env, oids) for el in value)) 

165 elif isinstance(value, list): 

166 return [self.render_template(element, context, jinja_env, oids) for element in value] 

167 elif isinstance(value, dict): 

168 return {k: self.render_template(v, context, jinja_env, oids) for k, v in value.items()} 

169 elif isinstance(value, set): 

170 return {self.render_template(element, context, jinja_env, oids) for element in value} 

171 

172 # More complex collections. 

173 self._render_nested_template_fields(value, context, jinja_env, oids) 

174 return value 

175 

176 def _render_nested_template_fields( 

177 self, 

178 value: Any, 

179 context: Context, 

180 jinja_env: jinja2.Environment, 

181 seen_oids: set[int], 

182 ) -> None: 

183 if id(value) in seen_oids: 

184 return 

185 seen_oids.add(id(value)) 

186 try: 

187 nested_template_fields = value.template_fields 

188 except AttributeError: 

189 # content has no inner template fields 

190 return 

191 self._do_render_template_fields(value, nested_template_fields, context, jinja_env, seen_oids)