cache properties from endpoint

Author: maxime.c

airbyte_cdk/sources/declarative/requesters/query_properties/properties_from_endpoint.py

@@ -7,6 +7,7 @@
 
 from airbyte_cdk.sources.declarative.interpolation import InterpolatedString
 from airbyte_cdk.sources.declarative.retrievers import Retriever
+from airbyte_cdk.sources.streams.core import StreamData
 from airbyte_cdk.sources.types import Config, StreamSlice
 
 
@@ -22,19 +23,22 @@ class PropertiesFromEndpoint:
     config: Config
     parameters: InitVar[Mapping[str, Any]]
 
+    _cached_properties: Optional[List[str]] = None
+
     def __post_init__(self, parameters: Mapping[str, Any]) -> None:
         self._property_field_path = [
             InterpolatedString(string=property_field, parameters=parameters)
             for property_field in self.property_field_path
         ]
 
-    def get_properties_from_endpoint(self, stream_slice: Optional[StreamSlice]) -> Iterable[str]:
-        response_properties = self.retriever.read_records(
-            records_schema={}, stream_slice=stream_slice
-        )
-        for property_obj in response_properties:
-            path = [
-                node.eval(self.config) if not isinstance(node, str) else node
-                for node in self._property_field_path
-            ]
-            yield dpath.get(property_obj, path, default=[])  # type: ignore # extracted will be a MutableMapping, given input data structure
+    def get_properties_from_endpoint(self, stream_slice: Optional[StreamSlice]) -> List[str]:
+        if self._cached_properties is None:
+            self._cached_properties = list(map(self._get_property, self.retriever.read_records(records_schema={}, stream_slice=stream_slice)))  # type: ignore # extracted will be a MutableMapping, given input data structure
+        return self._cached_properties
+
+    def _get_property(self, property_obj: Mapping[str, Any]) -> str:
+        path = [
+            node.eval(self.config) if not isinstance(node, str) else node
+            for node in self._property_field_path
+        ]
+        return str(dpath.get(property_obj, path, default=[]))  # type: ignore # extracted will be a MutableMapping, given input data structure

unit_tests/sources/declarative/requesters/query_properties/test_properties_from_endpoint.py

@@ -44,11 +44,7 @@ def test_get_properties_from_endpoint():
         parameters={},
     )
 
-    properties = list(
-        properties_from_endpoint.get_properties_from_endpoint(
-            stream_slice=StreamSlice(cursor_slice={}, partition={})
-        )
-    )
+    properties = properties_from_endpoint.get_properties_from_endpoint(stream_slice=StreamSlice(cursor_slice={}, partition={}))
 
     assert len(properties) == 9
     assert properties == expected_properties
@@ -89,11 +85,7 @@ def test_get_properties_from_endpoint_with_multiple_field_paths():
         parameters={},
     )
 
-    properties = list(
-        properties_from_endpoint.get_properties_from_endpoint(
-            stream_slice=StreamSlice(cursor_slice={}, partition={})
-        )
-    )
+    properties = properties_from_endpoint.get_properties_from_endpoint(stream_slice=StreamSlice(cursor_slice={}, partition={}))
 
     assert len(properties) == 9
     assert properties == expected_properties
@@ -135,11 +127,51 @@ def test_get_properties_from_endpoint_with_interpolation():
         parameters={},
     )
 
-    properties = list(
-        properties_from_endpoint.get_properties_from_endpoint(
-            stream_slice=StreamSlice(cursor_slice={}, partition={})
-        )
-    )
+    properties = properties_from_endpoint.get_properties_from_endpoint(stream_slice=StreamSlice(cursor_slice={}, partition={}))
 
     assert len(properties) == 9
     assert properties == expected_properties
+
+
+def test_given_multiple_calls_when_get_properties_from_endpoint_then_only_call_retriever_once():
+    retriever = Mock(spec=SimpleRetriever)
+    retriever.read_records.return_value = iter(
+        [
+            Record(stream_name="players", data={"id": "ace", "value": 1}),
+        ]
+    )
+
+    properties_from_endpoint = PropertiesFromEndpoint(
+        retriever=retriever,
+        property_field_path=["value"],
+        config={},
+        parameters={},
+    )
+
+    properties_from_endpoint.get_properties_from_endpoint(stream_slice=StreamSlice(cursor_slice={}, partition={}))
+    properties_from_endpoint.get_properties_from_endpoint(stream_slice=StreamSlice(cursor_slice={}, partition={}))
+    properties_from_endpoint.get_properties_from_endpoint(stream_slice=StreamSlice(cursor_slice={}, partition={}))
+
+    assert retriever.read_records.call_count == 1
+
+
+def test_given_value_is_int_when_get_properties_from_endpoint_then_return_str():
+    expected_properties = ["1"]
+
+    retriever = Mock(spec=SimpleRetriever)
+    retriever.read_records.return_value = iter(
+        [
+            Record(stream_name="players", data={"id": "ace", "value": 1}),
+        ]
+    )
+
+    properties_from_endpoint = PropertiesFromEndpoint(
+        retriever=retriever,
+        property_field_path=["value"],
+        config={},
+        parameters={},
+    )
+
+    properties = properties_from_endpoint.get_properties_from_endpoint(stream_slice=StreamSlice(cursor_slice={}, partition={}))
+
+    assert properties == expected_properties