CDK: Add schema normalization to declarative stream (#32786)

Co-authored-by: Eugene Kulak <kulak.eugene@gmail.com>
Co-authored-by: Yevhenii Kurochkin <ykurochkin@flyaps.com>
Co-authored-by: Alexandre Girard <alexandre@airbyte.io>

Author: 4 people

airbyte-cdk/python/airbyte_cdk/sources/declarative/declarative_component_schema.yaml

@@ -1203,12 +1203,10 @@ definitions:
       http_method:
         title: HTTP Method
         description: The HTTP method used to fetch data from the source (can be GET or POST).
-        anyOf:
-          - type: string
-          - type: string
-            enum:
-              - GET
-              - POST
+        type: string
+        enum:
+          - GET
+          - POST
         default: GET
         examples:
           - GET
@@ -1822,9 +1820,22 @@ definitions:
         title: Record Filter
         description: Responsible for filtering records to be emitted by the Source.
         "$ref": "#/definitions/RecordFilter"
+      schema_normalization:
+        "$ref": "#/definitions/SchemaNormalization"
+        default: None
       $parameters:
         type: object
         additionalProperties: true
+  SchemaNormalization:
+    title: Schema Normalization
+    description: Responsible for normalization according to the schema.
+    type: string
+    enum:
+      - None
+      - Default
+    examples:
+      - None
+      - Default
   RemoveFields:
     title: Remove Fields
     description: A transformation which removes fields from a record. The fields removed are designated using FieldPointers. During transformation, if a field or any of its parents does not exist in the record, no error is thrown.

airbyte-cdk/python/airbyte_cdk/sources/declarative/declarative_stream.py

@@ -101,7 +101,7 @@ def read_records(
         """
         :param: stream_state We knowingly avoid using stream_state as we want cursors to manage their own state.
         """
-        yield from self.retriever.read_records(stream_slice)
+        yield from self.retriever.read_records(self.get_json_schema(), stream_slice)
 
     def get_json_schema(self) -> Mapping[str, Any]:  # type: ignore
         """

airbyte-cdk/python/airbyte_cdk/sources/declarative/extractors/http_selector.py

@@ -22,13 +22,15 @@ def select_records(
         self,
         response: requests.Response,
         stream_state: StreamState,
+        records_schema: Mapping[str, Any],
         stream_slice: Optional[StreamSlice] = None,
         next_page_token: Optional[Mapping[str, Any]] = None,
     ) -> List[Record]:
         """
         Selects records from the response
         :param response: The response to select the records from
         :param stream_state: The stream state
+        :param records_schema: json schema of records to return
         :param stream_slice: The stream slice
         :param next_page_token: The paginator token
         :return: List of Records selected from the response

airbyte-cdk/python/airbyte_cdk/sources/declarative/extractors/record_selector.py

@@ -9,8 +9,15 @@
 from airbyte_cdk.sources.declarative.extractors.http_selector import HttpSelector
 from airbyte_cdk.sources.declarative.extractors.record_extractor import RecordExtractor
 from airbyte_cdk.sources.declarative.extractors.record_filter import RecordFilter
+from airbyte_cdk.sources.declarative.models import SchemaNormalization
 from airbyte_cdk.sources.declarative.transformations import RecordTransformation
 from airbyte_cdk.sources.declarative.types import Config, Record, StreamSlice, StreamState
+from airbyte_cdk.sources.utils.transform import TransformConfig, TypeTransformer
+
+SCHEMA_TRANSFORMER_TYPE_MAPPING = {
+    SchemaNormalization.None_: TransformConfig.NoTransform,
+    SchemaNormalization.Default: TransformConfig.DefaultSchemaNormalization,
+}
 
 
 @dataclass
@@ -21,13 +28,15 @@ class RecordSelector(HttpSelector):
 
     Attributes:
         extractor (RecordExtractor): The record extractor responsible for extracting records from a response
+        schema_normalization (TypeTransformer): The record normalizer responsible for casting record values to stream schema types
         record_filter (RecordFilter): The record filter responsible for filtering extracted records
         transformations (List[RecordTransformation]): The transformations to be done on the records
     """
 
     extractor: RecordExtractor
     config: Config
     parameters: InitVar[Mapping[str, Any]]
+    schema_normalization: TypeTransformer
     record_filter: Optional[RecordFilter] = None
     transformations: List[RecordTransformation] = field(default_factory=lambda: [])
 
@@ -38,14 +47,31 @@ def select_records(
         self,
         response: requests.Response,
         stream_state: StreamState,
+        records_schema: Mapping[str, Any],
         stream_slice: Optional[StreamSlice] = None,
         next_page_token: Optional[Mapping[str, Any]] = None,
     ) -> List[Record]:
+        """
+        Selects records from the response
+        :param response: The response to select the records from
+        :param stream_state: The stream state
+        :param records_schema: json schema of records to return
+        :param stream_slice: The stream slice
+        :param next_page_token: The paginator token
+        :return: List of Records selected from the response
+        """
         all_data = self.extractor.extract_records(response)
         filtered_data = self._filter(all_data, stream_state, stream_slice, next_page_token)
         self._transform(filtered_data, stream_state, stream_slice)
+        self._normalize_by_schema(filtered_data, schema=records_schema)
         return [Record(data, stream_slice) for data in filtered_data]
 
+    def _normalize_by_schema(self, records: List[Mapping[str, Any]], schema: Optional[Mapping[str, Any]]) -> List[Mapping[str, Any]]:
+        if schema:
+            # record has type Mapping[str, Any], but dict[str, Any] expected
+            return [self.schema_normalization.transform(record, schema) for record in records]  # type: ignore
+        return records
+
     def _filter(
         self,
         records: List[Mapping[str, Any]],
@@ -67,4 +93,5 @@ def _transform(
     ) -> None:
         for record in records:
             for transformation in self.transformations:
-                transformation.transform(record, config=self.config, stream_state=stream_state, stream_slice=stream_slice)
+                # record has type Mapping[str, Any], but Record expected
+                transformation.transform(record, config=self.config, stream_state=stream_state, stream_slice=stream_slice)  # type: ignore

airbyte-cdk/python/airbyte_cdk/sources/declarative/models/declarative_component_schema.py

@@ -340,7 +340,7 @@ class SessionTokenRequestBearerAuthenticator(BaseModel):
     type: Literal['Bearer']
 
 
-class HttpMethodEnum(Enum):
+class HttpMethod(Enum):
     GET = 'GET'
     POST = 'POST'
 
@@ -572,6 +572,11 @@ class RecordFilter(BaseModel):
     parameters: Optional[Dict[str, Any]] = Field(None, alias='$parameters')
 
 
+class SchemaNormalization(Enum):
+    None_ = 'None'
+    Default = 'Default'
+
+
 class RemoveFields(BaseModel):
     type: Literal['RemoveFields']
     field_pointers: List[List[str]] = Field(
@@ -1019,6 +1024,7 @@ class RecordSelector(BaseModel):
         description='Responsible for filtering records to be emitted by the Source.',
         title='Record Filter',
     )
+    schema_normalization: Optional[SchemaNormalization] = SchemaNormalization.None_
     parameters: Optional[Dict[str, Any]] = Field(None, alias='$parameters')
 
 
@@ -1232,8 +1238,8 @@ class HttpRequester(BaseModel):
         description='Error handler component that defines how to handle errors.',
         title='Error Handler',
     )
-    http_method: Optional[Union[str, HttpMethodEnum]] = Field(
-        'GET',
+    http_method: Optional[HttpMethod] = Field(
+        HttpMethod.GET,
         description='The HTTP method used to fetch data from the source (can be GET or POST).',
         examples=['GET', 'POST'],
         title='HTTP Method',

airbyte-cdk/python/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py

@@ -26,6 +26,7 @@
 from airbyte_cdk.sources.declarative.declarative_stream import DeclarativeStream
 from airbyte_cdk.sources.declarative.decoders import JsonDecoder
 from airbyte_cdk.sources.declarative.extractors import DpathExtractor, RecordFilter, RecordSelector
+from airbyte_cdk.sources.declarative.extractors.record_selector import SCHEMA_TRANSFORMER_TYPE_MAPPING
 from airbyte_cdk.sources.declarative.incremental import Cursor, CursorFactory, DatetimeBasedCursor, PerPartitionCursor
 from airbyte_cdk.sources.declarative.interpolation import InterpolatedString
 from airbyte_cdk.sources.declarative.interpolation.interpolated_mapping import InterpolatedMapping
@@ -107,6 +108,7 @@
 from airbyte_cdk.sources.declarative.requesters.request_option import RequestOptionType
 from airbyte_cdk.sources.declarative.requesters.request_options import InterpolatedRequestOptionsProvider
 from airbyte_cdk.sources.declarative.requesters.request_path import RequestPath
+from airbyte_cdk.sources.declarative.requesters.requester import HttpMethod
 from airbyte_cdk.sources.declarative.retrievers import SimpleRetriever, SimpleRetrieverTestReadDecorator
 from airbyte_cdk.sources.declarative.schema import DefaultSchemaLoader, InlineSchemaLoader, JsonFileSchemaLoader
 from airbyte_cdk.sources.declarative.spec import Spec
@@ -115,6 +117,7 @@
 from airbyte_cdk.sources.declarative.transformations.add_fields import AddedFieldDefinition
 from airbyte_cdk.sources.declarative.types import Config
 from airbyte_cdk.sources.message import InMemoryMessageRepository, LogAppenderMessageRepositoryDecorator, MessageRepository
+from airbyte_cdk.sources.utils.transform import TypeTransformer
 from isodate import parse_duration
 from pydantic import BaseModel
 
@@ -710,9 +713,8 @@ def create_http_requester(self, model: HttpRequesterModel, config: Config, *, na
             parameters=model.parameters or {},
         )
 
-        model_http_method = (
-            model.http_method if isinstance(model.http_method, str) else model.http_method.value if model.http_method is not None else "GET"
-        )
+        assert model.use_cache is not None  # for mypy
+        assert model.http_method is not None  # for mypy
 
         assert model.use_cache is not None  # for mypy
 
@@ -722,7 +724,7 @@ def create_http_requester(self, model: HttpRequesterModel, config: Config, *, na
             path=model.path,
             authenticator=authenticator,
             error_handler=error_handler,
-            http_method=model_http_method,
+            http_method=HttpMethod[model.http_method.value],
             request_options_provider=request_options_provider,
             config=config,
             disable_retries=self._disable_retries,
@@ -884,16 +886,24 @@ def create_request_option(model: RequestOptionModel, config: Config, **kwargs: A
         return RequestOption(field_name=model.field_name, inject_into=inject_into, parameters={})
 
     def create_record_selector(
-        self, model: RecordSelectorModel, config: Config, *, transformations: List[RecordTransformation], **kwargs: Any
+        self,
+        model: RecordSelectorModel,
+        config: Config,
+        *,
+        transformations: List[RecordTransformation],
+        **kwargs: Any,
     ) -> RecordSelector:
+        assert model.schema_normalization is not None  # for mypy
         extractor = self._create_component_from_model(model=model.extractor, config=config)
         record_filter = self._create_component_from_model(model.record_filter, config=config) if model.record_filter else None
+        schema_normalization = TypeTransformer(SCHEMA_TRANSFORMER_TYPE_MAPPING[model.schema_normalization])
 
         return RecordSelector(
             extractor=extractor,
             config=config,
             record_filter=record_filter,
             transformations=transformations,
+            schema_normalization=schema_normalization,
             parameters=model.parameters or {},
         )
 

airbyte-cdk/python/airbyte_cdk/sources/declarative/requesters/http_requester.py

@@ -59,7 +59,7 @@ class HttpRequester(Requester):
     config: Config
     parameters: InitVar[Mapping[str, Any]]
     authenticator: Optional[DeclarativeAuthenticator] = None
-    http_method: Union[str, HttpMethod] = HttpMethod.GET
+    http_method: HttpMethod = HttpMethod.GET
     request_options_provider: Optional[InterpolatedRequestOptionsProvider] = None
     error_handler: Optional[ErrorHandler] = None
     disable_retries: bool = False
@@ -80,7 +80,6 @@ def __post_init__(self, parameters: Mapping[str, Any]) -> None:
         else:
             self._request_options_provider = self.request_options_provider
         self._authenticator = self.authenticator or NoAuth(parameters=parameters)
-        self._http_method = HttpMethod[self.http_method] if isinstance(self.http_method, str) else self.http_method
         self.error_handler = self.error_handler
         self._parameters = parameters
         self.decoder = JsonDecoder(parameters={})
@@ -139,7 +138,7 @@ def get_path(
         return path.lstrip("/")
 
     def get_method(self) -> HttpMethod:
-        return self._http_method
+        return self.http_method
 
     def interpret_response_status(self, response: requests.Response) -> ResponseStatus:
         if self.error_handler is None:
@@ -420,7 +419,7 @@ def _create_prepared_request(
         data: Any = None,
     ) -> requests.PreparedRequest:
         url = urljoin(self.get_url_base(), path)
-        http_method = str(self._http_method.value)
+        http_method = str(self.http_method.value)
         query_params = self.deduplicate_query_params(url, params)
         args = {"method": http_method, "url": url, "headers": headers, "params": query_params}
         if http_method.upper() in BODY_REQUEST_METHODS:

airbyte-cdk/python/airbyte_cdk/sources/declarative/retrievers/retriever.py

@@ -4,7 +4,7 @@
 
 from abc import abstractmethod
 from dataclasses import dataclass
-from typing import Iterable, Optional
+from typing import Any, Iterable, Mapping, Optional
 
 from airbyte_cdk.sources.declarative.types import StreamSlice, StreamState
 from airbyte_cdk.sources.streams.core import StreamData
@@ -19,15 +19,14 @@ class Retriever:
     @abstractmethod
     def read_records(
         self,
+        records_schema: Mapping[str, Any],
         stream_slice: Optional[StreamSlice] = None,
     ) -> Iterable[StreamData]:
         """
         Fetch a stream's records from an HTTP API source
 
-        :param sync_mode: Unused but currently necessary for integrating with HttpStream
-        :param cursor_field: Unused but currently necessary for integrating with HttpStream
+        :param records_schema: json schema to describe record
         :param stream_slice: The stream slice to read data for
-        :param stream_state: The initial stream state
         :return: The records read from the API source
         """
 

airbyte-cdk/python/airbyte_cdk/sources/declarative/retrievers/simple_retriever.py

@@ -3,6 +3,7 @@
 #
 
 from dataclasses import InitVar, dataclass, field
+from functools import partial
 from itertools import islice
 from typing import Any, Callable, Iterable, List, Mapping, Optional, Set, Tuple, Union
 
@@ -215,6 +216,7 @@ def _parse_response(
         self,
         response: Optional[requests.Response],
         stream_state: StreamState,
+        records_schema: Mapping[str, Any],
         stream_slice: Optional[StreamSlice] = None,
         next_page_token: Optional[Mapping[str, Any]] = None,
     ) -> Iterable[Record]:
@@ -225,7 +227,11 @@ def _parse_response(
 
         self._last_response = response
         records = self.record_selector.select_records(
-            response=response, stream_state=stream_state, stream_slice=stream_slice, next_page_token=next_page_token
+            response=response,
+            stream_state=stream_state,
+            records_schema=records_schema,
+            stream_slice=stream_slice,
+            next_page_token=next_page_token,
         )
         self._records_from_last_response = records
         return records
@@ -271,16 +277,15 @@ def _fetch_next_page(
     # This logic is similar to _read_pages in the HttpStream class. When making changes here, consider making changes there as well.
     def _read_pages(
         self,
-        records_generator_fn: Callable[[Optional[requests.Response], Mapping[str, Any], Mapping[str, Any]], Iterable[StreamData]],
+        records_generator_fn: Callable[[Optional[requests.Response]], Iterable[StreamData]],
         stream_state: Mapping[str, Any],
         stream_slice: Mapping[str, Any],
     ) -> Iterable[StreamData]:
-        stream_state = stream_state or {}
         pagination_complete = False
         next_page_token = None
         while not pagination_complete:
             response = self._fetch_next_page(stream_state, stream_slice, next_page_token)
-            yield from records_generator_fn(response, stream_state, stream_slice)
+            yield from records_generator_fn(response)
 
             if not response:
                 pagination_complete = True
@@ -294,14 +299,28 @@ def _read_pages(
 
     def read_records(
         self,
+        records_schema: Mapping[str, Any],
         stream_slice: Optional[StreamSlice] = None,
     ) -> Iterable[StreamData]:
+        """
+        Fetch a stream's records from an HTTP API source
+
+        :param records_schema: json schema to describe record
+        :param stream_slice: The stream slice to read data for
+        :return: The records read from the API source
+        """
         stream_slice = stream_slice or {}  # None-check
         # Fixing paginator types has a long tail of dependencies
         self._paginator.reset()
 
         most_recent_record_from_slice = None
-        for stream_data in self._read_pages(self._parse_records, self.state, stream_slice):
+        record_generator = partial(
+            self._parse_records,
+            stream_state=self.state or {},
+            stream_slice=stream_slice,
+            records_schema=records_schema,
+        )
+        for stream_data in self._read_pages(record_generator, self.state, stream_slice):
             most_recent_record_from_slice = self._get_most_recent_record(most_recent_record_from_slice, stream_data, stream_slice)
             yield stream_data
 
@@ -361,9 +380,15 @@ def _parse_records(
         self,
         response: Optional[requests.Response],
         stream_state: Mapping[str, Any],
+        records_schema: Mapping[str, Any],
         stream_slice: Optional[Mapping[str, Any]],
     ) -> Iterable[StreamData]:
-        yield from self._parse_response(response, stream_slice=stream_slice, stream_state=stream_state)
+        yield from self._parse_response(
+            response,
+            stream_slice=stream_slice,
+            stream_state=stream_state,
+            records_schema=records_schema,
+        )
 
     def must_deduplicate_query_params(self) -> bool:
         return True

airbyte-cdk/python/bin/generate-component-manifest-files.sh

@@ -19,7 +19,8 @@ function main() {
       --input "/airbyte/$YAML_DIR/$filename_wo_ext.yaml" \
       --output "/airbyte/$OUTPUT_DIR/$filename_wo_ext.py" \
       --disable-timestamp \
-      --enum-field-as-literal one
+      --enum-field-as-literal one \
+      --set-default-enum-member
 
     # There is a limitation of Pydantic where a model's private fields starting with an underscore are inaccessible.
     # The Pydantic model generator replaces special characters like $ with the underscore which results in all

airbyte-cdk/python/unit_tests/sources/declarative/extractors/test_record_selector.py

@@ -13,6 +13,7 @@
 from airbyte_cdk.sources.declarative.extractors.record_selector import RecordSelector
 from airbyte_cdk.sources.declarative.transformations import RecordTransformation
 from airbyte_cdk.sources.declarative.types import Record
+from airbyte_cdk.sources.utils.transform import TransformConfig, TypeTransformer
 
 
 @pytest.mark.parametrize(
@@ -68,6 +69,7 @@ def test_record_filter(test_name, field_path, filter_template, body, expected_da
     stream_state = {"created_at": "06-06-21"}
     stream_slice = {"last_seen": "06-10-21"}
     next_page_token = {"last_seen_id": 14}
+    schema = create_schema()
     first_transformation = Mock(spec=RecordTransformation)
     second_transformation = Mock(spec=RecordTransformation)
     transformations = [first_transformation, second_transformation]
@@ -80,13 +82,19 @@ def test_record_filter(test_name, field_path, filter_template, body, expected_da
     else:
         record_filter = RecordFilter(config=config, condition=filter_template, parameters=parameters)
     record_selector = RecordSelector(
-        extractor=extractor, record_filter=record_filter, transformations=transformations, config=config, parameters=parameters
+        extractor=extractor,
+        record_filter=record_filter,
+        transformations=transformations,
+        config=config,
+        parameters=parameters,
+        schema_normalization=TypeTransformer(TransformConfig.NoTransform),
     )
 
     actual_records = record_selector.select_records(
-        response=response, stream_state=stream_state, stream_slice=stream_slice, next_page_token=next_page_token
+        response=response, records_schema=schema, stream_state=stream_state, stream_slice=stream_slice, next_page_token=next_page_token
     )
     assert actual_records == [Record(data, stream_slice) for data in expected_data]
+
     calls = []
     for record in expected_data:
         calls.append(call(record, config=config, stream_state=stream_state, stream_slice=stream_slice))
@@ -95,7 +103,77 @@ def test_record_filter(test_name, field_path, filter_template, body, expected_da
         transformation.transform.assert_has_calls(calls)
 
 
+@pytest.mark.parametrize(
+    "test_name, schema, schema_transformation, body, expected_data",
+    [
+        (
+            "test_with_empty_schema",
+            {},
+            TransformConfig.NoTransform,
+            {"data": [{"id": 1, "created_at": "06-06-21", "field_int": "100", "field_float": "123.3"}]},
+            [{"id": 1, "created_at": "06-06-21", "field_int": "100", "field_float": "123.3"}],
+        ),
+        (
+            "test_with_schema_none_normalizer",
+            {},
+            TransformConfig.NoTransform,
+            {"data": [{"id": 1, "created_at": "06-06-21", "field_int": "100", "field_float": "123.3"}]},
+            [{"id": 1, "created_at": "06-06-21", "field_int": "100", "field_float": "123.3"}],
+        ),
+        (
+            "test_with_schema_and_default_normalizer",
+            {},
+            TransformConfig.DefaultSchemaNormalization,
+            {"data": [{"id": 1, "created_at": "06-06-21", "field_int": "100", "field_float": "123.3"}]},
+            [{"id": "1", "created_at": "06-06-21", "field_int": 100, "field_float": 123.3}],
+        ),
+    ],
+)
+def test_schema_normalization(test_name, schema, schema_transformation, body, expected_data):
+    config = {"response_override": "stop_if_you_see_me"}
+    parameters = {"parameters_field": "data", "created_at": "06-07-21"}
+    stream_state = {"created_at": "06-06-21"}
+    stream_slice = {"last_seen": "06-10-21"}
+    next_page_token = {"last_seen_id": 14}
+
+    response = create_response(body)
+    schema = create_schema()
+    decoder = JsonDecoder(parameters={})
+    extractor = DpathExtractor(field_path=["data"], decoder=decoder, config=config, parameters=parameters)
+    record_selector = RecordSelector(
+        extractor=extractor,
+        record_filter=None,
+        transformations=[],
+        config=config,
+        parameters=parameters,
+        schema_normalization=TypeTransformer(schema_transformation),
+    )
+
+    actual_records = record_selector.select_records(
+        response=response,
+        stream_state=stream_state,
+        stream_slice=stream_slice,
+        next_page_token=next_page_token,
+        records_schema=schema,
+    )
+
+    assert actual_records == [Record(data, stream_slice) for data in expected_data]
+
+
 def create_response(body):
     response = requests.Response()
     response._content = json.dumps(body).encode("utf-8")
     return response
+
+
+def create_schema():
+    return {
+        "$schema": "http://json-schema.org/draft-07/schema#",
+        "type": "object",
+        "properties": {
+            "id": {"type": "string"},
+            "created_at": {"type": "string"},
+            "field_int": {"type": "integer"},
+            "field_float": {"type": "number"},
+        },
+    }

airbyte-cdk/python/unit_tests/sources/declarative/parsers/test_model_to_component_factory.py

@@ -245,7 +245,7 @@ def test_full_config_stream():
     assert stream.retriever.paginator.pagination_strategy.page_size == 10
 
     assert isinstance(stream.retriever.requester, HttpRequester)
-    assert stream.retriever.requester._http_method == HttpMethod.GET
+    assert stream.retriever.requester.http_method == HttpMethod.GET
     assert stream.retriever.requester.name == stream.name
     assert stream.retriever.requester._path.string == "{{ next_page_token['next_page_url'] }}"
     assert stream.retriever.requester._path.default == "{{ next_page_token['next_page_url'] }}"
@@ -829,7 +829,7 @@ def test_create_requester(test_name, error_handler, expected_backoff_strategy_ty
     )
 
     assert isinstance(selector, HttpRequester)
-    assert selector._http_method == HttpMethod.GET
+    assert selector.http_method == HttpMethod.GET
     assert selector.name == "name"
     assert selector._path.string == "/v3/marketing/lists"
     assert selector._url_base.string == "https://api.sendgrid.com"
@@ -1075,7 +1075,7 @@ def test_config_with_defaults():
     assert stream.schema_loader.file_path.default == "./source_sendgrid/schemas/{{ parameters.name }}.yaml"
 
     assert isinstance(stream.retriever.requester, HttpRequester)
-    assert stream.retriever.requester._http_method == HttpMethod.GET
+    assert stream.retriever.requester.http_method == HttpMethod.GET
 
     assert isinstance(stream.retriever.requester.authenticator, BearerAuthenticator)
     assert stream.retriever.requester.authenticator.token_provider.get_token() == "verysecrettoken"

airbyte-cdk/python/unit_tests/sources/declarative/requesters/test_http_requester.py

@@ -61,7 +61,7 @@ def factory(
 
 
 def test_http_requester():
-    http_method = "GET"
+    http_method = HttpMethod.GET
 
     request_options_provider = MagicMock()
     request_params = {"param": "value"}
@@ -106,7 +106,7 @@ def test_http_requester():
     assert requester.get_url_base() == "https://airbyte.io/"
     assert requester.get_path(stream_state={}, stream_slice=stream_slice, next_page_token={}) == "v1/1234"
     assert requester.get_authenticator() == authenticator
-    assert requester.get_method() == HttpMethod.GET
+    assert requester.get_method() == http_method
     assert requester.get_request_params(stream_state={}, stream_slice=None, next_page_token=None) == request_params
     assert requester.get_request_body_data(stream_state={}, stream_slice=None, next_page_token=None) == request_body_data
     assert requester.get_request_body_json(stream_state={}, stream_slice=None, next_page_token=None) == request_body_json

airbyte-cdk/python/unit_tests/sources/declarative/retrievers/test_simple_retriever.py

@@ -1,7 +1,6 @@
 #
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
-
 from unittest.mock import MagicMock, Mock, patch
 
 import pytest
@@ -94,7 +93,7 @@ def test_simple_retriever_full(mock_http_stream):
 
     assert retriever._last_response is None
     assert retriever._records_from_last_response == []
-    assert retriever._parse_response(response, stream_state={}) == records
+    assert retriever._parse_response(response, stream_state={}, records_schema={}) == records
     assert retriever._last_response == response
     assert retriever._records_from_last_response == records
 
@@ -170,7 +169,7 @@ def test_simple_retriever_with_request_response_log_last_records(mock_http_strea
 
     assert retriever._last_response is None
     assert retriever._records_from_last_response == []
-    assert retriever._parse_response(response, stream_state={}) == request_response_logs
+    assert retriever._parse_response(response, stream_state={}, records_schema={}) == request_response_logs
     assert retriever._last_response == response
     assert retriever._records_from_last_response == request_response_logs
 
@@ -396,13 +395,16 @@ def test_when_read_records_then_cursor_close_slice_with_greater_record(test_name
     )
     stream_slice = {"repository": "airbyte"}
 
+    def retriever_read_pages(_, __, ___):
+        return retriever._parse_records(response=MagicMock(), stream_state={}, stream_slice=stream_slice, records_schema={})
+
     with patch.object(
         SimpleRetriever,
         "_read_pages",
         return_value=iter([first_record, second_record]),
-        side_effect=lambda _, __, ___: retriever._parse_records(response=MagicMock(), stream_state=None, stream_slice=stream_slice),
+        side_effect=retriever_read_pages,
     ):
-        list(retriever.read_records(stream_slice=stream_slice))
+        list(retriever.read_records(stream_slice=stream_slice, records_schema={}))
         cursor.close_slice.assert_called_once_with(stream_slice, first_record if first_greater_than_second else second_record)
 
 
@@ -425,13 +427,16 @@ def test_given_stream_data_is_not_record_when_read_records_then_update_slice_wit
     )
     stream_slice = {"repository": "airbyte"}
 
+    def retriever_read_pages(_, __, ___):
+        return retriever._parse_records(response=MagicMock(), stream_state={}, stream_slice=stream_slice, records_schema={})
+
     with patch.object(
         SimpleRetriever,
         "_read_pages",
         return_value=iter(stream_data),
-        side_effect=lambda _, __, ___: retriever._parse_records(response=MagicMock(), stream_state=None, stream_slice=stream_slice),
+        side_effect=retriever_read_pages,
     ):
-        list(retriever.read_records(stream_slice=stream_slice))
+        list(retriever.read_records(stream_slice=stream_slice, records_schema={}))
         cursor.close_slice.assert_called_once_with(stream_slice, None)
 
 
@@ -440,7 +445,7 @@ def _generate_slices(number_of_slices):
 
 
 @patch.object(SimpleRetriever, "_read_pages", return_value=iter([]))
-def test_given_state_selector_when_read_records_use_stream_state(http_stream_read_pages):
+def test_given_state_selector_when_read_records_use_stream_state(http_stream_read_pages, mocker):
     requester = MagicMock()
     paginator = MagicMock()
     record_selector = MagicMock()
@@ -459,9 +464,10 @@ def test_given_state_selector_when_read_records_use_stream_state(http_stream_rea
         parameters={},
         config={},
     )
-    list(retriever.read_records(stream_slice=A_STREAM_SLICE))
 
-    http_stream_read_pages.assert_called_once_with(retriever._parse_records, A_STREAM_STATE, A_STREAM_SLICE)
+    list(retriever.read_records(stream_slice=A_STREAM_SLICE, records_schema={}))
+
+    http_stream_read_pages.assert_called_once_with(mocker.ANY, A_STREAM_STATE, A_STREAM_SLICE)
 
 
 def test_emit_log_request_response_messages(mocker):