Merge branch 'main' into revert-705-revert_concurrent_changes

Author: brianjlai

airbyte_cdk/connector_builder/test_reader/reader.py

@@ -120,7 +120,11 @@ def run_test_read(
         deprecation_warnings: List[LogMessage] = source.deprecation_warnings()
 
         schema_inferrer = SchemaInferrer(
-            self._pk_to_nested_and_composite_field(stream.primary_key) if stream else None,
+            self._pk_to_nested_and_composite_field(
+                stream.primary_key if hasattr(stream, "primary_key") else stream._primary_key  # type: ignore  # We are accessing the private property here as the primary key is not exposed. We should either expose it or use `as_airbyte_stream` to retrieve it as this is the "official" way where it is exposed in the Airbyte protocol
+            )
+            if stream
+            else None,
             self._cursor_field_to_nested_and_composite_field(stream.cursor_field)
             if stream
             else None,

airbyte_cdk/sources/declarative/concurrent_declarative_source.py

@@ -5,7 +5,18 @@
 import logging
 from dataclasses import dataclass, field
 from queue import Queue
-from typing import Any, ClassVar, Generic, Iterator, List, Mapping, MutableMapping, Optional, Tuple
+from typing import (
+    Any,
+    ClassVar,
+    Generic,
+    Iterator,
+    List,
+    Mapping,
+    MutableMapping,
+    Optional,
+    Tuple,
+    Union,
+)
 
 from airbyte_protocol_dataclasses.models import Level
 
@@ -19,10 +30,6 @@
 from airbyte_cdk.sources.connector_state_manager import ConnectorStateManager
 from airbyte_cdk.sources.declarative.concurrency_level import ConcurrencyLevel
 from airbyte_cdk.sources.declarative.declarative_stream import DeclarativeStream
-from airbyte_cdk.sources.declarative.extractors import RecordSelector
-from airbyte_cdk.sources.declarative.extractors.record_filter import (
-    ClientSideIncrementalRecordFilterDecorator,
-)
 from airbyte_cdk.sources.declarative.incremental import (
     ConcurrentPerPartitionCursor,
     GlobalSubstreamCursor,
@@ -32,7 +39,6 @@
     PerPartitionWithGlobalCursor,
 )
 from airbyte_cdk.sources.declarative.manifest_declarative_source import ManifestDeclarativeSource
-from airbyte_cdk.sources.declarative.models import FileUploader
 from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
     ConcurrencyLevel as ConcurrencyLevelModel,
 )
@@ -117,7 +123,6 @@ def __init__(
         # incremental streams running in full refresh.
         component_factory = ModelToComponentFactory(
             emit_connector_builder_messages=emit_connector_builder_messages,
-            disable_resumable_full_refresh=True,
             message_repository=ConcurrentMessageRepository(queue, message_repository),
             connector_state_manager=self._connector_state_manager,
             max_concurrent_async_job_count=source_config.get("max_concurrent_async_job_count"),
@@ -223,7 +228,7 @@ def discover(self, logger: logging.Logger, config: Mapping[str, Any]) -> Airbyte
             ]
         )
 
-    def streams(self, config: Mapping[str, Any]) -> List[Stream]:
+    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
         """
         The `streams` method is used as part of the AbstractSource in the following cases:
         * ConcurrentDeclarativeSource.check -> ManifestDeclarativeSource.check -> AbstractSource.check -> DeclarativeSource.check_connection -> CheckStream.check_connection -> streams
@@ -253,6 +258,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]
@@ -322,7 +331,7 @@ def _group_streams(
                         partition_generator = StreamSlicerPartitionGenerator(
                             partition_factory=DeclarativePartitionFactory(
                                 stream_name=declarative_stream.name,
-                                json_schema=declarative_stream.get_json_schema(),
+                                schema_loader=declarative_stream._schema_loader,  # type: ignore  # We are accessing the private property but the public one is optional and we will remove this code soonish
                                 retriever=retriever,
                                 message_repository=self.message_repository,
                                 max_records_limit=self._limits.max_records
@@ -359,7 +368,7 @@ def _group_streams(
                         partition_generator = StreamSlicerPartitionGenerator(
                             partition_factory=DeclarativePartitionFactory(
                                 stream_name=declarative_stream.name,
-                                json_schema=declarative_stream.get_json_schema(),
+                                schema_loader=declarative_stream._schema_loader,  # type: ignore  # We are accessing the private property but the public one is optional and we will remove this code soonish
                                 retriever=retriever,
                                 message_repository=self.message_repository,
                                 max_records_limit=self._limits.max_records
@@ -393,7 +402,7 @@ def _group_streams(
                     partition_generator = StreamSlicerPartitionGenerator(
                         DeclarativePartitionFactory(
                             stream_name=declarative_stream.name,
-                            json_schema=declarative_stream.get_json_schema(),
+                            schema_loader=declarative_stream._schema_loader,  # type: ignore  # We are accessing the private property but the public one is optional and we will remove this code soonish
                             retriever=declarative_stream.retriever,
                             message_repository=self.message_repository,
                             max_records_limit=self._limits.max_records if self._limits else None,
@@ -457,7 +466,7 @@ def _group_streams(
                     partition_generator = StreamSlicerPartitionGenerator(
                         DeclarativePartitionFactory(
                             stream_name=declarative_stream.name,
-                            json_schema=declarative_stream.get_json_schema(),
+                            schema_loader=declarative_stream._schema_loader,  # type: ignore  # We are accessing the private property but the public one is optional and we will remove this code soonish
                             retriever=retriever,
                             message_repository=self.message_repository,
                             max_records_limit=self._limits.max_records if self._limits else None,

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]]:  # 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
+        fully decouple this from the AbstractSource.
+        """
         if self._spec_component:
             self._spec_component.validate_config(config)