have bland stream be instantiated as DefaultStream

Author: maxime.c

airbyte_cdk/sources/declarative/concurrent_declarative_source.py

@@ -209,6 +209,10 @@ def _group_streams(
             # these legacy Python streams the way we do low-code streams to determine if they are concurrent compatible,
             # so we need to treat them as synchronous
 
+            if isinstance(declarative_stream, AbstractStream):
+                concurrent_streams.append(declarative_stream)
+                continue
+
             supports_file_transfer = (
                 isinstance(declarative_stream, DeclarativeStream)
                 and "file_uploader" in name_to_stream_mapping[declarative_stream.name]
@@ -278,7 +282,7 @@ def _group_streams(
                         partition_generator = StreamSlicerPartitionGenerator(
                             partition_factory=DeclarativePartitionFactory(
                                 declarative_stream.name,
-                                declarative_stream.get_json_schema(),
+                                declarative_stream.schema_loader,
                                 retriever,
                                 self.message_repository,
                             ),
@@ -309,7 +313,7 @@ def _group_streams(
                         partition_generator = StreamSlicerPartitionGenerator(
                             partition_factory=DeclarativePartitionFactory(
                                 declarative_stream.name,
-                                declarative_stream.get_json_schema(),
+                                declarative_stream.schema_loader,
                                 retriever,
                                 self.message_repository,
                             ),
@@ -339,7 +343,7 @@ def _group_streams(
                     partition_generator = StreamSlicerPartitionGenerator(
                         DeclarativePartitionFactory(
                             declarative_stream.name,
-                            declarative_stream.get_json_schema(),
+                            declarative_stream.schema_loader,
                             declarative_stream.retriever,
                             self.message_repository,
                         ),
@@ -399,7 +403,7 @@ def _group_streams(
                     partition_generator = StreamSlicerPartitionGenerator(
                         DeclarativePartitionFactory(
                             declarative_stream.name,
-                            declarative_stream.get_json_schema(),
+                            declarative_stream.schema_loader,
                             retriever,
                             self.message_repository,
                         ),

airbyte_cdk/sources/declarative/manifest_declarative_source.py

@@ -8,7 +8,7 @@
 from copy import deepcopy
 from importlib import metadata
 from types import ModuleType
-from typing import Any, Dict, Iterator, List, Mapping, Optional, Set
+from typing import Any, Dict, Iterator, List, Mapping, Optional, Set, Union
 
 import orjson
 import yaml
@@ -66,6 +66,7 @@
 from airbyte_cdk.sources.declarative.resolvers import COMPONENTS_RESOLVER_TYPE_MAPPING
 from airbyte_cdk.sources.declarative.spec.spec import Spec
 from airbyte_cdk.sources.message import MessageRepository
+from airbyte_cdk.sources.streams.concurrent.abstract_stream import AbstractStream
 from airbyte_cdk.sources.streams.core import Stream
 from airbyte_cdk.sources.types import Config, ConnectionDefinition
 from airbyte_cdk.sources.utils.slice_logger import (
@@ -297,7 +298,12 @@ 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[Stream]:
+    def streams(self, config: Mapping[str, Any]) -> List[Union[Stream, AbstractStream]]:
+        """
+        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
+        fully decouple this from the AbstractSource.
+        """
         if self._spec_component:
             self._spec_component.validate_config(config)
 

airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py

@@ -598,6 +598,7 @@
     Rate,
     UnlimitedCallRatePolicy,
 )
+from airbyte_cdk.sources.streams.concurrent.abstract_stream import AbstractStream
 from airbyte_cdk.sources.streams.concurrent.clamping import (
     ClampingEndProvider,
     ClampingStrategy,
@@ -1920,8 +1921,8 @@ def create_datetime_based_cursor(
         )
 
     def create_declarative_stream(
-        self, model: DeclarativeStreamModel, config: Config, **kwargs: Any
-    ) -> DeclarativeStream:
+        self, model: DeclarativeStreamModel, config: Config, is_parent=False, **kwargs: Any
+    ) -> Union[DeclarativeStream, AbstractStream]:
         # When constructing a declarative stream, we assemble the incremental_sync component and retriever's partition_router field
         # components if they exist into a single CartesianProductStreamSlicer. This is then passed back as an argument when constructing the
         # Retriever. This is done in the declarative stream not the retriever to support custom retrievers. The custom create methods in
@@ -2065,8 +2066,38 @@ def create_declarative_stream(
                 options["name"] = model.name
             schema_loader = DefaultSchemaLoader(config=config, parameters=options)
 
-        cursor_field = model.incremental_sync.cursor_field if model.incremental_sync else None
+        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
+            # * Streams without partition router but with cursor
+            # * Streams with both partition router and cursor
+            # We specifically exclude parent streams here because SubstreamPartitionRouter has not been updated yet
+            # We specifically exclude Connector Builder stuff for now as Brian is working on this anyway
+            stream_name = model.name or ""
+            partition_generator = StreamSlicerPartitionGenerator(
+                DeclarativePartitionFactory(
+                    stream_name,
+                    schema_loader,
+                    retriever,
+                    self._message_repository,
+                ),
+                combined_slicers,
+            )
+            FinalStateCursor(stream_name, None, self._message_repository)
+            return DefaultStream(
+                partition_generator=partition_generator,
+                name=stream_name,
+                json_schema=schema_loader.get_json_schema,
+                primary_key=get_primary_key_from_stream(primary_key),
+                cursor_field=None,
+                # FIXME we should have the cursor field has part of the interface of cursor
+                logger=logging.getLogger(f"airbyte.{stream_name}"),
+                # FIXME this is a breaking change compared to the old implementation,
+                cursor=FinalStateCursor(stream_name, None, self._message_repository),
+            )
 
+        cursor_field = model.incremental_sync.cursor_field if model.incremental_sync else None
         if model.state_migrations:
             state_transformations = [
                 self._create_component_from_model(state_migration, config, declarative_stream=model)
@@ -2094,7 +2125,7 @@ def _build_stream_slicer_from_partition_router(
         ],
         config: Config,
         stream_name: Optional[str] = None,
-    ) -> Optional[PartitionRouter]:
+    ) -> PartitionRouter:
         if (
             hasattr(model, "partition_router")
             and isinstance(model, SimpleRetrieverModel | AsyncRetrieverModel)
@@ -2115,15 +2146,15 @@ def _build_stream_slicer_from_partition_router(
                 return self._create_component_from_model(  # type: ignore[no-any-return] # Will be created PartitionRouter as stream_slicer_model is model.partition_router
                     model=stream_slicer_model, config=config, stream_name=stream_name or ""
                 )
-        return None
+        return SinglePartitionRouter(parameters={})
 
     def _build_incremental_cursor(
         self,
         model: DeclarativeStreamModel,
         stream_slicer: Optional[PartitionRouter],
         config: Config,
     ) -> Optional[StreamSlicer]:
-        if model.incremental_sync and stream_slicer:
+        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
@@ -2871,7 +2902,7 @@ def create_parent_stream_config(
         self, model: ParentStreamConfigModel, config: Config, **kwargs: Any
     ) -> ParentStreamConfig:
         declarative_stream = self._create_component_from_model(
-            model.stream, config=config, **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

@@ -3,6 +3,7 @@
 from typing import Any, Iterable, Mapping, Optional
 
 from airbyte_cdk.sources.declarative.retrievers import Retriever
+from airbyte_cdk.sources.declarative.schema import SchemaLoader
 from airbyte_cdk.sources.message import MessageRepository
 from airbyte_cdk.sources.streams.concurrent.partitions.partition import Partition
 from airbyte_cdk.sources.streams.concurrent.partitions.partition_generator import PartitionGenerator
@@ -11,11 +12,23 @@
 from airbyte_cdk.utils.slice_hasher import SliceHasher
 
 
+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()
+        return self._loaded_schema
+
+
 class DeclarativePartitionFactory:
     def __init__(
         self,
         stream_name: str,
-        json_schema: Mapping[str, Any],
+        schema_loader: SchemaLoader,
         retriever: Retriever,
         message_repository: MessageRepository,
     ) -> None:
@@ -25,14 +38,14 @@ def __init__(
         In order to avoid these problems, we will create one retriever per thread which should make the processing thread-safe.
         """
         self._stream_name = stream_name
-        self._json_schema = json_schema
+        self._schema_loader = SchemaLoaderCachingDecorator(schema_loader)
         self._retriever = retriever
         self._message_repository = message_repository
 
     def create(self, stream_slice: StreamSlice) -> Partition:
         return DeclarativePartition(
             self._stream_name,
-            self._json_schema,
+            self._schema_loader,
             self._retriever,
             self._message_repository,
             stream_slice,
@@ -43,20 +56,20 @@ class DeclarativePartition(Partition):
     def __init__(
         self,
         stream_name: str,
-        json_schema: Mapping[str, Any],
+        schema_loader: SchemaLoader,
         retriever: Retriever,
         message_repository: MessageRepository,
         stream_slice: StreamSlice,
     ):
         self._stream_name = stream_name
-        self._json_schema = json_schema
+        self._schema_loader = schema_loader
         self._retriever = retriever
         self._message_repository = message_repository
         self._stream_slice = stream_slice
         self._hash = SliceHasher.hash(self._stream_name, self._stream_slice)
 
     def read(self) -> Iterable[Record]:
-        for stream_data in self._retriever.read_records(self._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,8 +5,7 @@
 import copy
 import json
 import logging
-from functools import lru_cache
-from typing import Any, Iterable, List, Mapping, MutableMapping, Optional, Tuple, Union
+from typing import Any, Iterable, List, Mapping, MutableMapping, Optional, Union
 
 from typing_extensions import deprecated
 
@@ -196,7 +195,6 @@ 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
 
-    @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

@@ -2,9 +2,8 @@
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
 
-from functools import lru_cache
 from logging import Logger
-from typing import Any, Iterable, List, Mapping, Optional
+from typing import Any, Iterable, List, Mapping, Optional, Union, Callable
 
 from airbyte_cdk.models import AirbyteStream, SyncMode
 from airbyte_cdk.sources.streams.concurrent.abstract_stream import AbstractStream
@@ -20,7 +19,7 @@ def __init__(
         self,
         partition_generator: PartitionGenerator,
         name: str,
-        json_schema: Mapping[str, Any],
+        json_schema: Union[Mapping[str, Any], Callable[[], Mapping[str, Any]]],
         primary_key: List[str],
         cursor_field: Optional[str],
         logger: Logger,
@@ -53,14 +52,15 @@ def namespace(self) -> Optional[str]:
     def cursor_field(self) -> Optional[str]:
         return self._cursor_field
 
-    @lru_cache(maxsize=None)
     def get_json_schema(self) -> Mapping[str, Any]:
+        if isinstance(self._json_schema, Callable):
+            return self._json_schema()
         return self._json_schema
 
     def as_airbyte_stream(self) -> AirbyteStream:
         stream = AirbyteStream(
             name=self.name,
-            json_schema=dict(self._json_schema),
+            json_schema=dict(self.get_json_schema()),
             supported_sync_modes=[SyncMode.full_refresh],
             is_resumable=False,
             is_file_based=self._supports_file_transfer,

unit_tests/connector_builder/test_connector_builder_handler.py

@@ -780,7 +780,7 @@ def test_config_update() -> None:
         "client_secret": "a client secret",
         "refresh_token": "a refresh token",
     }
-    source = ManifestDeclarativeSource(source_config=manifest)
+    source = ManifestDeclarativeSource(source_config=manifest, emit_connector_builder_messages=True)
 
     refresh_request_response = {
         "access_token": "an updated access token",