Coverage for /usr/local/lib/python3.12/site-packages/prefect/server/utilities/user_templates.py: 36%

1"""Utilities to support safely rendering user-supplied templates""" 

2 

3from collections.abc import Mapping 

4from typing import TYPE_CHECKING, Any, Optional 

5 

6from jinja2 import ChainableUndefined, nodes, pass_eval_context 

7from jinja2.nodes import EvalContext 

8from jinja2.sandbox import ImmutableSandboxedEnvironment 

9from jinja2.utils import htmlsafe_json_dumps 

10from markupsafe import Markup 

11from pydantic import BaseModel 

12 

13from prefect.logging import get_logger 

14 

15if TYPE_CHECKING: 15 ↛ 16line 15 didn't jump to line 16 because the condition on line 15 was never true

16 import logging 

17 

18logger: "logging.Logger" = get_logger(__name__) 

19 

20 

21MAX_TEMPLATE_RANGE = 100 

22MAX_LOOP_COUNT = 10 

23MAX_NESTED_LOOP_DEPTH = 2 

24 

25 

26def _check_template_range(*args: int) -> range: 

27 rng = range(*args) 

28 

29 if len(rng) > MAX_TEMPLATE_RANGE: 

30 raise OverflowError( 

31 "Range too big. The sandbox blocks ranges larger than" 

32 f" {MAX_TEMPLATE_RANGE=}." 

33 ) 

34 return rng 

35 

36 

37def _prepare_value_for_json(value: object) -> object: 

38 if isinstance(value, BaseModel): 

39 return _prepare_value_for_json(value.model_dump(mode="json")) 

40 elif isinstance(value, Mapping): 

41 return { 

42 key: _prepare_value_for_json(nested_value) 

43 for key, nested_value in value.items() 

44 } 

45 elif isinstance(value, list): 

46 return [_prepare_value_for_json(nested_value) for nested_value in value] 

47 elif isinstance(value, tuple): 

48 return tuple(_prepare_value_for_json(nested_value) for nested_value in value) 

49 else: 

50 return value 

51 

52 

53@pass_eval_context 

54def _tojson( 

55 eval_ctx: EvalContext, value: object, indent: Optional[int] = None 

56) -> Markup: 

57 policies = eval_ctx.environment.policies 

58 dumps = policies["json.dumps_function"] 

59 kwargs = policies["json.dumps_kwargs"] 

60 

61 if indent is not None: 

62 kwargs = kwargs.copy() 

63 kwargs["indent"] = indent 

64 

65 return htmlsafe_json_dumps(_prepare_value_for_json(value), dumps=dumps, **kwargs) 

66 

67 

68class UserTemplateEnvironment(ImmutableSandboxedEnvironment): 

69 def __init__(self, *args: Any, **kwargs: Any) -> None: 

70 super().__init__(*args, **kwargs) 

71 # Override the range function to limit its size 

72 self.globals["range"] = _check_template_range # type: ignore 

73 self.filters["tojson"] = _tojson 

74 

75 

76_template_environment = UserTemplateEnvironment( 

77 undefined=ChainableUndefined, 

78 enable_async=True, 

79 extensions=[ 

80 # Supports human-friendly rendering of dates and times 

81 # https://pypi.org/project/jinja2-humanize-extension/ 

82 "jinja2_humanize_extension.HumanizeExtension", 

83 ], 

84) 

85 

86_sync_template_environment = UserTemplateEnvironment( 

87 undefined=ChainableUndefined, 

88 enable_async=False, 

89 extensions=[ 

90 # Supports human-friendly rendering of dates and times 

91 # https://pypi.org/project/jinja2-humanize-extension/ 

92 "jinja2_humanize_extension.HumanizeExtension", 

93 ], 

94) 

95 

96 

97class TemplateSecurityError(Exception): 

98 """Raised when extended validation of a template fails.""" 

99 

100 def __init__(self, message: Optional[str] = None, line_number: int = 0) -> None: 

101 self.lineno = line_number 

102 self.message = message 

103 super().__init__(message) 

104 

105 

106class TemplateRenderError(Exception): 

107 """Raised when a user-supplied template fails to render.""" 

108 

109 def __init__(self, error: Exception, template: str) -> None: 

110 self.error = error 

111 self.template = template 

112 super().__init__( 

113 f"Failed to render template due to the following error: {error!r}" 

114 ) 

115 

116 

117def register_user_template_filters(filters: dict[str, Any]) -> None: 

118 """Register additional filters that will be available to user templates""" 

119 _template_environment.filters.update(filters) 

120 _sync_template_environment.filters.update(filters) 

121 

122 

123def validate_user_template(template: str) -> None: 

124 root_node = _template_environment.parse(template) 

125 _validate_loop_constraints(root_node) 

126 

127 

128def _validate_loop_constraints(root_node: nodes.Template): 

129 for_nodes = [node for node in root_node.find_all(nodes.For)] 

130 

131 if not for_nodes: 131 ↛ 134line 131 didn't jump to line 134 because the condition on line 131 was always true

132 return 

133 

134 if len(for_nodes) > MAX_LOOP_COUNT: 

135 raise TemplateSecurityError( 

136 f"Contains {len(for_nodes)} for loops. Templates can contain no " 

137 f"more than {MAX_LOOP_COUNT} for loops." 

138 ) 

139 

140 max_nested_depth = max(_nested_loop_depth(for_node) for for_node in for_nodes) 

141 if max_nested_depth > MAX_NESTED_LOOP_DEPTH: 

142 raise TemplateSecurityError( 

143 f"Contains nested for loops at a depth of {max_nested_depth}. " 

144 "Templates can nest for loops no more than " 

145 f"{MAX_NESTED_LOOP_DEPTH} loops deep." 

146 ) 

147 

148 

149def _nested_loop_depth(node: nodes.Node, depth: int = 0) -> int: 

150 children = [child for child in node.iter_child_nodes()] 

151 

152 if isinstance(node, nodes.For): 

153 depth += 1 

154 

155 if not children: 

156 return depth 

157 

158 return max(_nested_loop_depth(child, depth) for child in children) 

159 

160 

161def matching_types_in_templates(templates: list[str], types: set[str]) -> list[str]: 

162 found: set[str] = set() 

163 

164 for template in templates: 164 ↛ 165line 164 didn't jump to line 165 because the loop on line 164 never started

165 root_node = _template_environment.parse(template) 

166 for node in root_node.find_all(nodes.Name): 

167 if node.ctx == "load" and node.name in types: 

168 found.add(node.name) 

169 

170 return list(found) 

171 

172 

173def maybe_template(possible: str) -> bool: 

174 return "{{" in possible or "{%" in possible 

175 

176 

177async def render_user_template(template: str, context: dict[str, Any]) -> str: 

178 if not maybe_template(template): 

179 return template 

180 

181 try: 

182 loaded = _template_environment.from_string(template) 

183 return await loaded.render_async(context) 

184 except Exception as e: 

185 logger.warning("Unhandled exception rendering template", exc_info=True) 

186 raise TemplateRenderError(e, template) from e 

187 

188 

189def render_user_template_sync(template: str, context: dict[str, Any]) -> str: 

190 if not maybe_template(template): 

191 return template 

192 

193 try: 

194 loaded = _sync_template_environment.from_string(template) 

195 return loaded.render(context) 

196 except Exception as e: 

197 logger.warning("Unhandled exception rendering template", exc_info=True) 

198 raise TemplateRenderError(e, template) from e