fix test, format, lint and a bit of mypy

Author: maxime.c

airbyte_cdk/sources/declarative/manifest_declarative_source.py

@@ -298,7 +298,7 @@ def connection_checker(self) -> ConnectionChecker:
                 f"Expected to generate a ConnectionChecker component, but received {check_stream.__class__}"
             )
 
-    def streams(self, config: Mapping[str, Any]) -> List[Union[Stream, AbstractStream]]:
+    def streams(self, config: Mapping[str, Any]) -> List[Union[Stream, AbstractStream]]:  # type: ignore  # we are migrating away from the AbstractSource and are expecting that this will only be called by ConcurrentDeclarativeSource or the Connector Builder
         """
         As a migration step, this method will return both legacy stream (Stream) and concurrent stream (AbstractStream).
         Once the migration is done, we can probably have this method throw "not implemented" as we figure out how to

airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py

@@ -545,8 +545,10 @@
     StreamSlicer,
     StreamSlicerTestReadDecorator,
 )
-from airbyte_cdk.sources.declarative.stream_slicers.declarative_partition_generator import \
-    StreamSlicerPartitionGenerator, DeclarativePartitionFactory
+from airbyte_cdk.sources.declarative.stream_slicers.declarative_partition_generator import (
+    DeclarativePartitionFactory,
+    StreamSlicerPartitionGenerator,
+)
 from airbyte_cdk.sources.declarative.transformations import (
     AddFields,
     RecordTransformation,
@@ -608,7 +610,12 @@
     WeekClampingStrategy,
     Weekday,
 )
-from airbyte_cdk.sources.streams.concurrent.cursor import ConcurrentCursor, CursorField, Cursor, FinalStateCursor
+from airbyte_cdk.sources.streams.concurrent.cursor import (
+    ConcurrentCursor,
+    Cursor,
+    CursorField,
+    FinalStateCursor,
+)
 from airbyte_cdk.sources.streams.concurrent.default_stream import DefaultStream
 from airbyte_cdk.sources.streams.concurrent.helpers import get_primary_key_from_stream
 from airbyte_cdk.sources.streams.concurrent.state_converters.datetime_stream_state_converter import (
@@ -2066,7 +2073,11 @@ def create_declarative_stream(
                 options["name"] = model.name
             schema_loader = DefaultSchemaLoader(config=config, parameters=options)
 
-        if isinstance(combined_slicers, PartitionRouter) and not is_parent and not self._emit_connector_builder_messages:
+        if (
+            isinstance(combined_slicers, PartitionRouter)
+            and not is_parent
+            and not self._emit_connector_builder_messages
+        ):
             # We are starting to migrate streams to instantiate directly the DefaultStream instead of instantiating the
             # DeclarativeStream and assembling the DefaultStream from that. The plan is the following:
             # * Streams without partition router nor cursors and streams with only partition router. This is the `isinstance(combined_slicers, PartitionRouter)` condition as the first kind with have a SinglePartitionRouter
@@ -2154,7 +2165,9 @@ def _build_incremental_cursor(
         stream_slicer: Optional[PartitionRouter],
         config: Config,
     ) -> Optional[StreamSlicer]:
-        if model.incremental_sync and (stream_slicer and not isinstance(stream_slicer, SinglePartitionRouter)):
+        if model.incremental_sync and (
+            stream_slicer and not isinstance(stream_slicer, SinglePartitionRouter)
+        ):
             if model.retriever.type == "AsyncRetriever":
                 stream_name = model.name or ""
                 stream_namespace = None
@@ -2902,7 +2915,10 @@ def create_parent_stream_config(
         self, model: ParentStreamConfigModel, config: Config, **kwargs: Any
     ) -> ParentStreamConfig:
         declarative_stream = self._create_component_from_model(
-            model.stream, config=config, is_parent=True, **kwargs,
+            model.stream,
+            config=config,
+            is_parent=True,
+            **kwargs,
         )
         request_option = (
             self._create_component_from_model(model.request_option, config=config)

airbyte_cdk/sources/declarative/stream_slicers/declarative_partition_generator.py

@@ -13,14 +13,16 @@
 
 
 class SchemaLoaderCachingDecorator(SchemaLoader):
-
     def __init__(self, schema_loader: SchemaLoader):
         self._decorated = schema_loader
         self._loaded_schema = None
 
     def get_json_schema(self) -> Mapping[str, Any]:
         if self._loaded_schema is None:
             self._loaded_schema = self._decorated.get_json_schema()
+
+        if self._loaded_schema is None:
+            raise ValueError("Could not load schema")
         return self._loaded_schema
 
 
@@ -69,7 +71,9 @@ def __init__(
         self._hash = SliceHasher.hash(self._stream_name, self._stream_slice)
 
     def read(self) -> Iterable[Record]:
-        for stream_data in self._retriever.read_records(self._schema_loader.get_json_schema(), self._stream_slice):
+        for stream_data in self._retriever.read_records(
+            self._schema_loader.get_json_schema(), self._stream_slice
+        ):
             if isinstance(stream_data, Mapping):
                 record = (
                     stream_data

airbyte_cdk/sources/streams/concurrent/adapters.py

@@ -5,6 +5,7 @@
 import copy
 import json
 import logging
+from functools import lru_cache
 from typing import Any, Iterable, List, Mapping, MutableMapping, Optional, Union
 
 from typing_extensions import deprecated
@@ -195,6 +196,8 @@ def cursor_field(self) -> Union[str, List[str]]:
     def cursor(self) -> Optional[Cursor]:  # type: ignore[override] # StreamFaced expects to use only airbyte_cdk.sources.streams.concurrent.cursor.Cursor
         return self._cursor
 
+    # FIXME the lru_cache seems to be mostly there because of typing issue
+    @lru_cache(maxsize=None)
     def get_json_schema(self) -> Mapping[str, Any]:
         return self._abstract_stream.get_json_schema()
 

airbyte_cdk/sources/streams/concurrent/default_stream.py

@@ -3,7 +3,7 @@
 #
 
 from logging import Logger
-from typing import Any, Iterable, List, Mapping, Optional, Union, Callable
+from typing import Any, Callable, Iterable, List, Mapping, Optional, Union
 
 from airbyte_cdk.models import AirbyteStream, SyncMode
 from airbyte_cdk.sources.streams.concurrent.abstract_stream import AbstractStream
@@ -109,7 +109,9 @@ def check_availability(self) -> StreamAvailability:
                 f"Cannot attempt to connect to stream {self.name} - no stream slices were found"
             )
         except AirbyteTracedException as error:
-            return StreamAvailability.unavailable(error.message or error.internal_message or "<no error message>")
+            return StreamAvailability.unavailable(
+                error.message or error.internal_message or "<no error message>"
+            )
 
         try:
             next(iter(partition.read()))
@@ -118,4 +120,6 @@ def check_availability(self) -> StreamAvailability:
             self._logger.info(f"Successfully connected to stream {self.name}, but got 0 records.")
             return StreamAvailability.available()
         except AirbyteTracedException as error:
-            return StreamAvailability.unavailable(error.message or error.internal_message or "<no error message>")
+            return StreamAvailability.unavailable(
+                error.message or error.internal_message or "<no error message>"
+            )

unit_tests/sources/declarative/decoders/test_decoders_memory_usage.py

@@ -93,9 +93,8 @@ def get_body():
     requests_mock.get("https://for-all-mankind.nasa.com/api/v1/users/users3", body=get_body())
     requests_mock.get("https://for-all-mankind.nasa.com/api/v1/users/users4", body=get_body())
 
-    stream_slices = list(stream.stream_slices(sync_mode=SyncMode.full_refresh))
-    for stream_slice in stream_slices:
-        for _ in stream.retriever.read_records(records_schema={}, stream_slice=stream_slice):
+    for partition in stream.generate_partitions():
+        for _ in partition.read():
             counter += 1
 
-    assert counter == lines_in_response * len(stream_slices)
+    assert counter == lines_in_response * 4  # 4 partitions

unit_tests/sources/declarative/parsers/test_model_to_component_factory.py

@@ -157,7 +157,9 @@
 from airbyte_cdk.sources.declarative.schema.schema_loader import SchemaLoader
 from airbyte_cdk.sources.declarative.spec import Spec
 from airbyte_cdk.sources.declarative.stream_slicers import StreamSlicerTestReadDecorator
-from airbyte_cdk.sources.declarative.stream_slicers.declarative_partition_generator import SchemaLoaderCachingDecorator
+from airbyte_cdk.sources.declarative.stream_slicers.declarative_partition_generator import (
+    SchemaLoaderCachingDecorator,
+)
 from airbyte_cdk.sources.declarative.transformations import AddFields, RemoveFields
 from airbyte_cdk.sources.declarative.transformations.add_fields import AddedFieldDefinition
 from airbyte_cdk.sources.declarative.yaml_declarative_source import YamlDeclarativeSource
@@ -1768,14 +1770,8 @@ def test_config_with_defaults():
 
     schema_loader = get_schema_loader(stream)
     assert isinstance(schema_loader, JsonFileSchemaLoader)
-    assert (
-        schema_loader.file_path.string
-        == "./source_sendgrid/schemas/{{ parameters.name }}.yaml"
-    )
-    assert (
-        schema_loader.file_path.default
-        == "./source_sendgrid/schemas/{{ parameters.name }}.yaml"
-    )
+    assert schema_loader.file_path.string == "./source_sendgrid/schemas/{{ parameters.name }}.yaml"
+    assert schema_loader.file_path.default == "./source_sendgrid/schemas/{{ parameters.name }}.yaml"
 
     assert isinstance(retriever.requester, HttpRequester)
     assert retriever.requester.http_method == HttpMethod.GET
@@ -1785,9 +1781,9 @@ def test_config_with_defaults():
 
     assert isinstance(retriever.record_selector, RecordSelector)
     assert isinstance(retriever.record_selector.extractor, DpathExtractor)
-    assert [
-        fp.eval(input_config) for fp in retriever.record_selector.extractor._field_path
-    ] == ["result"]
+    assert [fp.eval(input_config) for fp in retriever.record_selector.extractor._field_path] == [
+        "result"
+    ]
 
     assert isinstance(retriever.paginator, DefaultPaginator)
     assert retriever.paginator.url_base.string == "https://api.sendgrid.com"
@@ -2508,7 +2504,9 @@ def test_default_schema_loader(self):
         ),
     ],
 )
-def test_merge_incremental_and_partition_router(incremental, partition_router, expected_router_type, expected_stream_type):
+def test_merge_incremental_and_partition_router(
+    incremental, partition_router, expected_router_type, expected_stream_type
+):
     stream_model = {
         "type": "DeclarativeStream",
         "retriever": {
@@ -2542,7 +2540,11 @@ def test_merge_incremental_and_partition_router(incremental, partition_router, e
     assert isinstance(stream, expected_stream_type)
     retriever = get_retriever(stream)
     assert isinstance(retriever, SimpleRetriever)
-    stream_slicer = retriever.stream_slicer if expected_stream_type == DeclarativeStream else stream._stream_partition_generator._stream_slicer
+    stream_slicer = (
+        retriever.stream_slicer
+        if expected_stream_type == DeclarativeStream
+        else stream._stream_partition_generator._stream_slicer
+    )
     assert isinstance(stream_slicer, expected_router_type)
 
     if incremental and partition_router:
@@ -2722,7 +2724,9 @@ def test_create_custom_retriever():
     )
 
     assert isinstance(stream, DefaultStream)
-    assert isinstance(stream._stream_partition_generator._partition_factory._retriever, MyCustomRetriever)
+    assert isinstance(
+        stream._stream_partition_generator._partition_factory._retriever, MyCustomRetriever
+    )
 
 
 @freezegun.freeze_time("2021-01-01 00:00:00")
@@ -4667,9 +4671,16 @@ def test_create_stream_with_multiple_schema_loaders():
 
 
 def get_schema_loader(stream: DefaultStream):
-    assert isinstance(stream._stream_partition_generator._partition_factory._schema_loader, SchemaLoaderCachingDecorator)
+    assert isinstance(
+        stream._stream_partition_generator._partition_factory._schema_loader,
+        SchemaLoaderCachingDecorator,
+    )
     return stream._stream_partition_generator._partition_factory._schema_loader._decorated
 
 
 def get_retriever(stream: Union[DeclarativeStream, DefaultStream]):
-    return stream.retriever if isinstance(stream, DeclarativeStream) else stream._stream_partition_generator._partition_factory._retriever
+    return (
+        stream.retriever
+        if isinstance(stream, DeclarativeStream)
+        else stream._stream_partition_generator._partition_factory._retriever
+    )

unit_tests/sources/declarative/resolvers/test_config_components_resolver.py

@@ -383,5 +383,6 @@ def test_component_mapping_conditions(manifest, config, expected_conditional_par
     for stream in source.streams(config):
         if stream.name in expected_conditional_params:
             assert (
-                stream._stream_partition_generator._partition_factory._retriever.requester._parameters == expected_conditional_params[stream.name]
+                stream._stream_partition_generator._partition_factory._retriever.requester._parameters
+                == expected_conditional_params[stream.name]
             )

unit_tests/sources/declarative/test_manifest_declarative_source.py

@@ -28,7 +28,9 @@
     SyncMode,
     Type,
 )
-from airbyte_cdk.sources.declarative.concurrent_declarative_source import ConcurrentDeclarativeSource
+from airbyte_cdk.sources.declarative.concurrent_declarative_source import (
+    ConcurrentDeclarativeSource,
+)
 from airbyte_cdk.sources.declarative.declarative_stream import DeclarativeStream
 from airbyte_cdk.sources.declarative.manifest_declarative_source import ManifestDeclarativeSource
 from airbyte_cdk.sources.declarative.parsers.model_to_component_factory import (
@@ -2191,14 +2193,10 @@ def test_only_parent_streams_use_cache():
     assert not stream_1_retriever.requester.use_cache
 
     # Parent stream created for substream
-    assert (
-            stream_1_retriever.stream_slicer.parent_stream_configs[0].stream.name
-            == "applications"
-    )
-    assert (
-        stream_1_retriever.stream_slicer.parent_stream_configs[0]
-        .stream.retriever.requester.use_cache
-    )
+    assert stream_1_retriever.stream_slicer.parent_stream_configs[0].stream.name == "applications"
+    assert stream_1_retriever.stream_slicer.parent_stream_configs[
+        0
+    ].stream.retriever.requester.use_cache
 
     # Main stream without caching
     assert streams[2].name == "jobs"