44
55import ast
66from functools import cache
7- from typing import Any , Mapping , Optional , Tuple , Type
7+ from typing import Any , Mapping , Optional , Set , Tuple , Type
88
99from jinja2 import meta
1010from jinja2 .environment import Template
@@ -27,7 +27,35 @@ class StreamPartitionAccessEnvironment(SandboxedEnvironment):
2727 def is_safe_attribute (self , obj : Any , attr : str , value : Any ) -> bool :
2828 if attr in ["_partition" ]:
2929 return True
30- return super ().is_safe_attribute (obj , attr , value )
30+ return super ().is_safe_attribute (obj , attr , value ) # type: ignore # for some reason, mypy says 'Returning Any from function declared to return "bool"'
31+
32+
33+ # These aliases are used to deprecate existing keywords without breaking all existing connectors.
34+ _ALIASES = {
35+ "stream_interval" : "stream_slice" , # Use stream_interval to access incremental_sync values
36+ "stream_partition" : "stream_slice" , # Use stream_partition to access partition router's values
37+ }
38+
39+ # These extensions are not installed so they're not currently a problem,
40+ # but we're still explicitly removing them from the jinja context.
41+ # At worst, this is documentation that we do NOT want to include these extensions because of the potential security risks
42+ _RESTRICTED_EXTENSIONS = ["jinja2.ext.loopcontrols" ] # Adds support for break continue in loops
43+
44+ # By default, these Python builtin functions are available in the Jinja context.
45+ # We explicitly remove them because of the potential security risk.
46+ # Please add a unit test to test_jinja.py when adding a restriction.
47+ _RESTRICTED_BUILTIN_FUNCTIONS = [
48+ "range"
49+ ] # The range function can cause very expensive computations
50+
51+ _ENVIRONMENT = StreamPartitionAccessEnvironment ()
52+ _ENVIRONMENT .filters .update (** filters )
53+ _ENVIRONMENT .globals .update (** macros )
54+
55+ for extension in _RESTRICTED_EXTENSIONS :
56+ _ENVIRONMENT .extensions .pop (extension , None )
57+ for builtin in _RESTRICTED_BUILTIN_FUNCTIONS :
58+ _ENVIRONMENT .globals .pop (builtin , None )
3159
3260
3361class JinjaInterpolation (Interpolation ):
@@ -48,34 +76,6 @@ class JinjaInterpolation(Interpolation):
4876 Additional information on jinja templating can be found at https://jinja.palletsprojects.com/en/3.1.x/templates/#
4977 """
5078
51- # These aliases are used to deprecate existing keywords without breaking all existing connectors.
52- ALIASES = {
53- "stream_interval" : "stream_slice" , # Use stream_interval to access incremental_sync values
54- "stream_partition" : "stream_slice" , # Use stream_partition to access partition router's values
55- }
56-
57- # These extensions are not installed so they're not currently a problem,
58- # but we're still explicitely removing them from the jinja context.
59- # At worst, this is documentation that we do NOT want to include these extensions because of the potential security risks
60- RESTRICTED_EXTENSIONS = ["jinja2.ext.loopcontrols" ] # Adds support for break continue in loops
61-
62- # By default, these Python builtin functions are available in the Jinja context.
63- # We explicitely remove them because of the potential security risk.
64- # Please add a unit test to test_jinja.py when adding a restriction.
65- RESTRICTED_BUILTIN_FUNCTIONS = [
66- "range"
67- ] # The range function can cause very expensive computations
68-
69- def __init__ (self ) -> None :
70- self ._environment = StreamPartitionAccessEnvironment ()
71- self ._environment .filters .update (** filters )
72- self ._environment .globals .update (** macros )
73-
74- for extension in self .RESTRICTED_EXTENSIONS :
75- self ._environment .extensions .pop (extension , None )
76- for builtin in self .RESTRICTED_BUILTIN_FUNCTIONS :
77- self ._environment .globals .pop (builtin , None )
78-
7979 def eval (
8080 self ,
8181 input_str : str ,
@@ -86,7 +86,7 @@ def eval(
8686 ) -> Any :
8787 context = {"config" : config , ** additional_parameters }
8888
89- for alias , equivalent in self . ALIASES .items ():
89+ for alias , equivalent in _ALIASES .items ():
9090 if alias in context :
9191 # This is unexpected. We could ignore or log a warning, but failing loudly should result in fewer surprises
9292 raise ValueError (
@@ -105,6 +105,7 @@ def eval(
105105 raise Exception (f"Expected a string, got { input_str } " )
106106 except UndefinedError :
107107 pass
108+
108109 # If result is empty or resulted in an undefined error, evaluate and return the default string
109110 return self ._literal_eval (self ._eval (default , context ), valid_types )
110111
@@ -132,16 +133,16 @@ def _eval(self, s: Optional[str], context: Mapping[str, Any]) -> Optional[str]:
132133 return s
133134
134135 @cache
135- def _find_undeclared_variables (self , s : Optional [str ]) -> set [str ]:
136+ def _find_undeclared_variables (self , s : Optional [str ]) -> Set [str ]:
136137 """
137138 Find undeclared variables and cache them
138139 """
139- ast = self . _environment .parse (s ) # type: ignore # parse is able to handle None
140+ ast = _ENVIRONMENT .parse (s ) # type: ignore # parse is able to handle None
140141 return meta .find_undeclared_variables (ast )
141142
142143 @cache
143- def _compile (self , s : Optional [ str ] ) -> Template :
144+ def _compile (self , s : str ) -> Template :
144145 """
145146 We must cache the Jinja Template ourselves because we're using `from_string` instead of a template loader
146147 """
147- return self . _environment . from_string (s ) # type: ignore [arg-type] # Expected `str | Template` but passed `str | None`
148+ return _ENVIRONMENT . from_string (s )
0 commit comments