Improve performance of interpolation in decalarative sources (#44027)

Co-authored-by: Natik Gadzhi <[email protected]>
Co-authored-by: Alexandre Girard <[email protected]>

Author: 3 people

airbyte-cdk/python/airbyte_cdk/sources/declarative/interpolation/interpolated_string.py

@@ -28,6 +28,9 @@ def __post_init__(self, parameters: Mapping[str, Any]) -> None:
         self.default = self.default or self.string
         self._interpolation = JinjaInterpolation()
         self._parameters = parameters
+        # indicates whether passed string is just a plain string, not Jinja template
+        # This allows for optimization, but we do not know it yet at this stage
+        self._is_plain_string = None
 
     def eval(self, config: Config, **kwargs: Any) -> Any:
         """
@@ -37,6 +40,14 @@ def eval(self, config: Config, **kwargs: Any) -> Any:
         :param kwargs: Optional parameters used for interpolation
         :return: The interpolated string
         """
+        if self._is_plain_string:
+            return self.string
+        if self._is_plain_string is None:
+            # Let's check whether output from evaluation is the same as input.
+            # This indicates occurence of a plain string, not a template and we can skip Jinja in subsequent runs.
+            evaluated = self._interpolation.eval(self.string, config, self.default, parameters=self._parameters, **kwargs)
+            self._is_plain_string = self.string == evaluated
+            return evaluated
         return self._interpolation.eval(self.string, config, self.default, parameters=self._parameters, **kwargs)
 
     def __eq__(self, other: Any) -> bool:

airbyte-cdk/python/airbyte_cdk/sources/declarative/interpolation/jinja.py

@@ -3,13 +3,15 @@
 #
 
 import ast
+from functools import cache
 from typing import Any, Mapping, Optional, Tuple, Type
 
 from airbyte_cdk.sources.declarative.interpolation.filters import filters
 from airbyte_cdk.sources.declarative.interpolation.interpolation import Interpolation
 from airbyte_cdk.sources.declarative.interpolation.macros import macros
 from airbyte_cdk.sources.types import Config
 from jinja2 import meta
+from jinja2.environment import Template
 from jinja2.exceptions import UndefinedError
 from jinja2.sandbox import SandboxedEnvironment
 
@@ -114,13 +116,27 @@ def _literal_eval(self, result: Optional[str], valid_types: Optional[Tuple[Type[
 
     def _eval(self, s: Optional[str], context: Mapping[str, Any]) -> Optional[str]:
         try:
-            ast = self._environment.parse(s)  # type: ignore # parse is able to handle None
-            undeclared = meta.find_undeclared_variables(ast)
+            undeclared = self._find_undeclared_variables(s)
             undeclared_not_in_context = {var for var in undeclared if var not in context}
             if undeclared_not_in_context:
                 raise ValueError(f"Jinja macro has undeclared variables: {undeclared_not_in_context}. Context: {context}")
-            return self._environment.from_string(s).render(context)  # type: ignore # from_string is able to handle None
+            return self._compile(s).render(context)  # type: ignore # from_string is able to handle None
         except TypeError:
             # The string is a static value, not a jinja template
             # It can be returned as is
             return s
+
+    @cache
+    def _find_undeclared_variables(self, s: Optional[str]) -> Template:
+        """
+        Find undeclared variables and cache them
+        """
+        ast = self._environment.parse(s)  # type: ignore # parse is able to handle None
+        return meta.find_undeclared_variables(ast)
+
+    @cache
+    def _compile(self, s: Optional[str]) -> Template:
+        """
+        We must cache the Jinja Template ourselves because we're using `from_string` instead of a template loader
+        """
+        return self._environment.from_string(s)

airbyte-cdk/python/unit_tests/sources/declarative/interpolation/test_jinja.py

@@ -20,6 +20,13 @@ def test_get_value_from_config():
     assert val == "2022-01-01"
 
 
+def test_get_missing_value_from_config():
+    s = "{{ config['date'] }}"
+    config = {}
+    val = interpolation.eval(s, config)
+    assert val is None
+
+
 @pytest.mark.parametrize(
     "valid_types, expected_value",
     [
@@ -42,6 +49,14 @@ def test_get_value_from_stream_slice():
     assert val == "2020-09-09"
 
 
+def test_get_missing_value_from_stream_slice():
+    s = "{{ stream_slice['date'] }}"
+    config = {"date": "2022-01-01"}
+    stream_slice = {}
+    val = interpolation.eval(s, config, **{"stream_slice": stream_slice})
+    assert val is None
+
+
 def test_get_value_from_a_list_of_mappings():
     s = "{{ records[0]['date'] }}"
     config = {"date": "2022-01-01"}
@@ -142,7 +157,7 @@ def test_negative_day_delta():
         ("test_false_to_string", False, "false"),
         ("test_array_to_string", ["hello", "world"], '["hello", "world"]'),
         ("test_object_to_array", {"hello": "world"}, '{"hello": "world"}'),
-    ]
+    ],
 )
 def test_to_string(test_name, input_value, expected_output):
     interpolation = JinjaInterpolation()
@@ -254,7 +269,7 @@ def test_undeclared_variables(template_string, expected_error, expected_value):
             id="test_now_utc_with_duration_and_format",
         ),
         pytest.param("{{ 1 | string }}", "1", id="test_int_to_string"),
-        pytest.param("{{ [\"hello\", \"world\"] | string }}", "[\"hello\", \"world\"]", id="test_array_to_string"),
+        pytest.param('{{ ["hello", "world"] | string }}', '["hello", "world"]', id="test_array_to_string"),
     ],
 )
 def test_macros_examples(template_string, expected_value):