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

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

2 

3from typing import TYPE_CHECKING, Any, Optional 1 ctx1a

4 

5from jinja2 import ChainableUndefined, nodes 1 ctx1a

6from jinja2.sandbox import ImmutableSandboxedEnvironment 1 ctx1a

7 

8from prefect.logging import get_logger 1 ctx1a

9 

10if TYPE_CHECKING: 10 ↛ 11line 10 didn't jump to line 11 because the condition on line 10 was never true1 ctx1a

11 import logging 

12 

13logger: "logging.Logger" = get_logger(__name__) 1 ctx1a

14 

15 

16MAX_TEMPLATE_RANGE = 100 1 ctx1a

17MAX_LOOP_COUNT = 10 1 ctx1a

18MAX_NESTED_LOOP_DEPTH = 2 1 ctx1a

19 

20 

21def _check_template_range(*args: int) -> range: 1 ctx1a

22 rng = range(*args) 

23 

24 if len(rng) > MAX_TEMPLATE_RANGE: 

25 raise OverflowError( 

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

27 f" {MAX_TEMPLATE_RANGE=}." 

28 ) 

29 return rng 

30 

31 

32class UserTemplateEnvironment(ImmutableSandboxedEnvironment): 1 ctx1a

33 def __init__(self, *args: Any, **kwargs: Any) -> None: 1 ctx1a

34 super().__init__(*args, **kwargs) 1 ctx1a

35 # Override the range function to limit its size 

36 self.globals["range"] = _check_template_range # type: ignore 1 ctx1a

37 

38 

39_template_environment = UserTemplateEnvironment( 1 ctx1a

40 undefined=ChainableUndefined, 

41 enable_async=True, 

42 extensions=[ 

43 # Supports human-friendly rendering of dates and times 

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

45 "jinja2_humanize_extension.HumanizeExtension", 

46 ], 

47) 

48 

49_sync_template_environment = UserTemplateEnvironment( 1 ctx1a

50 undefined=ChainableUndefined, 

51 enable_async=False, 

52 extensions=[ 

53 # Supports human-friendly rendering of dates and times 

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

55 "jinja2_humanize_extension.HumanizeExtension", 

56 ], 

57) 

58 

59 

60class TemplateSecurityError(Exception): 1 ctx1a

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

62 

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

64 self.lineno = line_number 

65 self.message = message 

66 super().__init__(message) 

67 

68 

69def register_user_template_filters(filters: dict[str, Any]) -> None: 1 ctx1a

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

71 _template_environment.filters.update(filters) 

72 

73 

74def validate_user_template(template: str) -> None: 1 ctx1a

75 root_node = _template_environment.parse(template) 

76 _validate_loop_constraints(root_node) 

77 

78 

79def _validate_loop_constraints(root_node: nodes.Template): 1 ctx1a

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

81 

82 if not for_nodes: 

83 return 

84 

85 if len(for_nodes) > MAX_LOOP_COUNT: 

86 raise TemplateSecurityError( 

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

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

89 ) 

90 

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

92 if max_nested_depth > MAX_NESTED_LOOP_DEPTH: 

93 raise TemplateSecurityError( 

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

95 "Templates can nest for loops no more than " 

96 f"{MAX_NESTED_LOOP_DEPTH} loops deep." 

97 ) 

98 

99 

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

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

102 

103 if isinstance(node, nodes.For): 

104 depth += 1 

105 

106 if not children: 

107 return depth 

108 

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

110 

111 

112def matching_types_in_templates(templates: list[str], types: set[str]) -> list[str]: 1 ctx1a

113 found: set[str] = set() 

114 

115 for template in templates: 

116 root_node = _template_environment.parse(template) 

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

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

119 found.add(node.name) 

120 

121 return list(found) 

122 

123 

124def maybe_template(possible: str) -> bool: 1 ctx1a

125 return "{{" in possible or "{%" in possible 

126 

127 

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

129 if not maybe_template(template): 

130 return template 

131 

132 try: 

133 loaded = _template_environment.from_string(template) 

134 return await loaded.render_async(context) 

135 except Exception as e: 

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

137 return ( 

138 f"Failed to render template due to the following error: {e!r}\n" 

139 "Template source:\n" 

140 ) + template 

141 

142 

143def render_user_template_sync(template: str, context: dict[str, Any]) -> str: 1 ctx1a

144 if not maybe_template(template): 

145 return template 

146 

147 try: 

148 loaded = _sync_template_environment.from_string(template) 

149 return loaded.render(context) 

150 except Exception as e: 

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

152 return ( 

153 f"Failed to render template due to the following error: {e!r}\n" 

154 "Template source:\n" 

155 ) + template