airbyte-cdk 6.34.1.dev0__py3-none-any.whl → 6.34.1.dev1__py3-none-any.whl

airbyte_cdk/sources/declarative/async_job/job_orchestrator.py

@@ -437,10 +437,10 @@ class AsyncJobOrchestrator:
                 yield from self._process_running_partitions_and_yield_completed_ones()
                 self._wait_on_status_update()
             except Exception as exception:
+                LOGGER.warning(
+                    f"Caught exception that stops the processing of the jobs: {exception}. Traceback: {traceback.format_exc()}"
+                )
                 if self._is_breaking_exception(exception):
-                    LOGGER.warning(
-                        f"Caught exception that stops the processing of the jobs: {exception}"
-                    )
                     self._abort_all_running_jobs()
                     raise exception
 
@@ -482,16 +482,16 @@ class AsyncJobOrchestrator:
             and exception.failure_type == FailureType.config_error
         )
 
-    def fetch_records(self, partition: AsyncPartition) -> Iterable[Mapping[str, Any]]:
+    def fetch_records(self, async_jobs: Iterable[AsyncJob]) -> Iterable[Mapping[str, Any]]:
         """
-        Fetches records from the given partition's jobs.
+        Fetches records from the given jobs.
 
         Args:
-            partition (AsyncPartition): The partition containing the jobs.
+            async_jobs Iterable[AsyncJob]: The list of AsyncJobs.
 
         Yields:
             Iterable[Mapping[str, Any]]: The fetched records from the jobs.
         """
-        for job in partition.jobs:
+        for job in async_jobs:
             yield from self._job_repository.fetch_records(job)
             self._job_repository.delete(job)

airbyte_cdk/sources/declarative/auth/jwt.py

@@ -3,6 +3,7 @@
 #
 
 import base64
+import json
 from dataclasses import InitVar, dataclass
 from datetime import datetime
 from typing import Any, Mapping, Optional, Union
@@ -104,21 +105,21 @@ class JwtAuthenticator(DeclarativeAuthenticator):
         )
 
     def _get_jwt_headers(self) -> dict[str, Any]:
-        """ "
+        """
         Builds and returns the headers used when signing the JWT.
         """
-        headers = self._additional_jwt_headers.eval(self.config)
+        headers = self._additional_jwt_headers.eval(self.config, json_loads=json.loads)
         if any(prop in headers for prop in ["kid", "alg", "typ", "cty"]):
             raise ValueError(
                 "'kid', 'alg', 'typ', 'cty' are reserved headers and should not be set as part of 'additional_jwt_headers'"
             )
 
         if self._kid:
-            headers["kid"] = self._kid.eval(self.config)
+            headers["kid"] = self._kid.eval(self.config, json_loads=json.loads)
         if self._typ:
-            headers["typ"] = self._typ.eval(self.config)
+            headers["typ"] = self._typ.eval(self.config, json_loads=json.loads)
         if self._cty:
-            headers["cty"] = self._cty.eval(self.config)
+            headers["cty"] = self._cty.eval(self.config, json_loads=json.loads)
         headers["alg"] = self._algorithm
         return headers
 
@@ -130,18 +131,19 @@ class JwtAuthenticator(DeclarativeAuthenticator):
         exp = now + self._token_duration if isinstance(self._token_duration, int) else now
         nbf = now
 
-        payload = self._additional_jwt_payload.eval(self.config)
+        payload = self._additional_jwt_payload.eval(self.config, json_loads=json.loads)
         if any(prop in payload for prop in ["iss", "sub", "aud", "iat", "exp", "nbf"]):
             raise ValueError(
                 "'iss', 'sub', 'aud', 'iat', 'exp', 'nbf' are reserved properties and should not be set as part of 'additional_jwt_payload'"
             )
 
         if self._iss:
-            payload["iss"] = self._iss.eval(self.config)
+            payload["iss"] = self._iss.eval(self.config, json_loads=json.loads)
         if self._sub:
-            payload["sub"] = self._sub.eval(self.config)
+            payload["sub"] = self._sub.eval(self.config, json_loads=json.loads)
         if self._aud:
-            payload["aud"] = self._aud.eval(self.config)
+            payload["aud"] = self._aud.eval(self.config, json_loads=json.loads)
+
         payload["iat"] = now
         payload["exp"] = exp
         payload["nbf"] = nbf
@@ -151,7 +153,7 @@ class JwtAuthenticator(DeclarativeAuthenticator):
         """
         Returns the secret key used to sign the JWT.
         """
-        secret_key: str = self._secret_key.eval(self.config)
+        secret_key: str = self._secret_key.eval(self.config, json_loads=json.loads)
         return (
             base64.b64encode(secret_key.encode()).decode()
             if self._base64_encode_secret_key
@@ -176,7 +178,11 @@ class JwtAuthenticator(DeclarativeAuthenticator):
         """
         Returns the header prefix to be used when attaching the token to the request.
         """
-        return self._header_prefix.eval(self.config) if self._header_prefix else None
+        return (
+            self._header_prefix.eval(self.config, json_loads=json.loads)
+            if self._header_prefix
+            else None
+        )
 
     @property
     def auth_header(self) -> str:

airbyte_cdk/sources/declarative/auth/oauth.py

@@ -3,7 +3,7 @@
 #
 
 from dataclasses import InitVar, dataclass, field
-from datetime import timedelta
+from datetime import datetime, timedelta
 from typing import Any, List, Mapping, MutableMapping, Optional, Union
 
 from airbyte_cdk.sources.declarative.auth.declarative_authenticator import DeclarativeAuthenticator
@@ -232,8 +232,13 @@ class DeclarativeOauth2Authenticator(AbstractOauth2Authenticator, DeclarativeAut
         return self._refresh_request_headers.eval(self.config)
 
     def get_token_expiry_date(self) -> AirbyteDateTime:
+        if not self._has_access_token_been_initialized():
+            return AirbyteDateTime.from_datetime(datetime.min)
         return self._token_expiry_date  # type: ignore # _token_expiry_date is an AirbyteDateTime. It is never None despite what mypy thinks
 
+    def _has_access_token_been_initialized(self) -> bool:
+        return self._access_token is not None
+
     def set_token_expiry_date(self, value: Union[str, int]) -> None:
         self._token_expiry_date = self._parse_token_expiration_date(value)
 

airbyte_cdk/sources/declarative/auth/token.py

@@ -5,7 +5,7 @@
 import base64
 import logging
 from dataclasses import InitVar, dataclass
-from typing import Any, Mapping, Union
+from typing import Any, Mapping, MutableMapping, Union
 
 import requests
 from cachetools import TTLCache, cached
@@ -45,11 +45,6 @@ class ApiKeyAuthenticator(DeclarativeAuthenticator):
     config: Config
     parameters: InitVar[Mapping[str, Any]]
 
-    def __post_init__(self, parameters: Mapping[str, Any]) -> None:
-        self._field_name = InterpolatedString.create(
-            self.request_option.field_name, parameters=parameters
-        )
-
     @property
     def auth_header(self) -> str:
         options = self._get_request_options(RequestOptionType.header)
@@ -60,9 +55,9 @@ class ApiKeyAuthenticator(DeclarativeAuthenticator):
         return self.token_provider.get_token()
 
     def _get_request_options(self, option_type: RequestOptionType) -> Mapping[str, Any]:
-        options = {}
+        options: MutableMapping[str, Any] = {}
         if self.request_option.inject_into == option_type:
-            options[self._field_name.eval(self.config)] = self.token
+            self.request_option.inject_into_request(options, self.token, self.config)
         return options
 
     def get_request_params(self) -> Mapping[str, Any]:

airbyte_cdk/sources/declarative/concurrent_declarative_source.py

@@ -3,7 +3,7 @@
 #
 
 import logging
-from typing import Any, Generic, Iterator, List, Mapping, Optional, Tuple
+from typing import Any, Generic, Iterator, List, Mapping, MutableMapping, Optional, Tuple
 
 from airbyte_cdk.models import (
     AirbyteCatalog,
@@ -19,11 +19,11 @@ 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
 from airbyte_cdk.sources.declarative.incremental.datetime_based_cursor import DatetimeBasedCursor
 from airbyte_cdk.sources.declarative.incremental.per_partition_with_global import (
     PerPartitionWithGlobalCursor,
 )
-from airbyte_cdk.sources.declarative.interpolation import InterpolatedString
 from airbyte_cdk.sources.declarative.manifest_declarative_source import ManifestDeclarativeSource
 from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
     ConcurrencyLevel as ConcurrencyLevelModel,
@@ -35,13 +35,11 @@ from airbyte_cdk.sources.declarative.parsers.model_to_component_factory import (
     ModelToComponentFactory,
 )
 from airbyte_cdk.sources.declarative.partition_routers import AsyncJobPartitionRouter
-from airbyte_cdk.sources.declarative.requesters import HttpRequester
 from airbyte_cdk.sources.declarative.retrievers import AsyncRetriever, Retriever, SimpleRetriever
 from airbyte_cdk.sources.declarative.stream_slicers.declarative_partition_generator import (
     DeclarativePartitionFactory,
     StreamSlicerPartitionGenerator,
 )
-from airbyte_cdk.sources.declarative.transformations.add_fields import AddFields
 from airbyte_cdk.sources.declarative.types import ConnectionDefinition
 from airbyte_cdk.sources.source import TState
 from airbyte_cdk.sources.streams import Stream
@@ -230,6 +228,7 @@ class ConcurrentDeclarativeSource(ManifestDeclarativeSource, Generic[TState]):
                     stream_state = self._connector_state_manager.get_stream_state(
                         stream_name=declarative_stream.name, namespace=declarative_stream.namespace
                     )
+                    stream_state = self._migrate_state(declarative_stream, stream_state)
 
                     retriever = self._get_retriever(declarative_stream, stream_state)
 
@@ -238,7 +237,7 @@ class ConcurrentDeclarativeSource(ManifestDeclarativeSource, Generic[TState]):
                     ):
                         cursor = declarative_stream.retriever.stream_slicer.stream_slicer
 
-                        if not isinstance(cursor, ConcurrentCursor):
+                        if not isinstance(cursor, ConcurrentCursor | ConcurrentPerPartitionCursor):
                             # This should never happen since we instantiate ConcurrentCursor in
                             # model_to_component_factory.py
                             raise ValueError(
@@ -326,9 +325,6 @@ class ConcurrentDeclarativeSource(ManifestDeclarativeSource, Generic[TState]):
                     incremental_sync_component_definition
                     and incremental_sync_component_definition.get("type", "")
                     == DatetimeBasedCursorModel.__name__
-                    and self._stream_supports_concurrent_partition_processing(
-                        declarative_stream=declarative_stream
-                    )
                     and hasattr(declarative_stream.retriever, "stream_slicer")
                     and isinstance(
                         declarative_stream.retriever.stream_slicer, PerPartitionWithGlobalCursor
@@ -337,6 +333,8 @@ class ConcurrentDeclarativeSource(ManifestDeclarativeSource, Generic[TState]):
                     stream_state = self._connector_state_manager.get_stream_state(
                         stream_name=declarative_stream.name, namespace=declarative_stream.namespace
                     )
+                    stream_state = self._migrate_state(declarative_stream, stream_state)
+
                     partition_router = declarative_stream.retriever.stream_slicer._partition_router
 
                     perpartition_cursor = (
@@ -401,9 +399,6 @@ class ConcurrentDeclarativeSource(ManifestDeclarativeSource, Generic[TState]):
             and bool(incremental_sync_component_definition)
             and incremental_sync_component_definition.get("type", "")
             == DatetimeBasedCursorModel.__name__
-            and self._stream_supports_concurrent_partition_processing(
-                declarative_stream=declarative_stream
-            )
             and hasattr(declarative_stream.retriever, "stream_slicer")
             and (
                 isinstance(declarative_stream.retriever.stream_slicer, DatetimeBasedCursor)
@@ -411,72 +406,6 @@ class ConcurrentDeclarativeSource(ManifestDeclarativeSource, Generic[TState]):
             )
         )
 
-    def _stream_supports_concurrent_partition_processing(
-        self, declarative_stream: DeclarativeStream
-    ) -> bool:
-        """
-        Many connectors make use of stream_state during interpolation on a per-partition basis under the assumption that
-        state is updated sequentially. Because the concurrent CDK engine processes different partitions in parallel,
-        stream_state is no longer a thread-safe interpolation context. It would be a race condition because a cursor's
-        stream_state can be updated in any order depending on which stream partition's finish first.
-
-        We should start to move away from depending on the value of stream_state for low-code components that operate
-        per-partition, but we need to gate this otherwise some connectors will be blocked from publishing. See the
-        cdk-migrations.md for the full list of connectors.
-        """
-
-        if isinstance(declarative_stream.retriever, SimpleRetriever) and isinstance(
-            declarative_stream.retriever.requester, HttpRequester
-        ):
-            http_requester = declarative_stream.retriever.requester
-            if "stream_state" in http_requester._path.string:
-                self.logger.warning(
-                    f"Low-code stream '{declarative_stream.name}' uses interpolation of stream_state in the HttpRequester which is not thread-safe. Defaulting to synchronous processing"
-                )
-                return False
-
-            request_options_provider = http_requester._request_options_provider
-            if request_options_provider.request_options_contain_stream_state():
-                self.logger.warning(
-                    f"Low-code stream '{declarative_stream.name}' uses interpolation of stream_state in the HttpRequester which is not thread-safe. Defaulting to synchronous processing"
-                )
-                return False
-
-            record_selector = declarative_stream.retriever.record_selector
-            if isinstance(record_selector, RecordSelector):
-                if (
-                    record_selector.record_filter
-                    and not isinstance(
-                        record_selector.record_filter, ClientSideIncrementalRecordFilterDecorator
-                    )
-                    and "stream_state" in record_selector.record_filter.condition
-                ):
-                    self.logger.warning(
-                        f"Low-code stream '{declarative_stream.name}' uses interpolation of stream_state in the RecordFilter which is not thread-safe. Defaulting to synchronous processing"
-                    )
-                    return False
-
-                for add_fields in [
-                    transformation
-                    for transformation in record_selector.transformations
-                    if isinstance(transformation, AddFields)
-                ]:
-                    for field in add_fields.fields:
-                        if isinstance(field.value, str) and "stream_state" in field.value:
-                            self.logger.warning(
-                                f"Low-code stream '{declarative_stream.name}' uses interpolation of stream_state in the AddFields which is not thread-safe. Defaulting to synchronous processing"
-                            )
-                            return False
-                        if (
-                            isinstance(field.value, InterpolatedString)
-                            and "stream_state" in field.value.string
-                        ):
-                            self.logger.warning(
-                                f"Low-code stream '{declarative_stream.name}' uses interpolation of stream_state in the AddFields which is not thread-safe. Defaulting to synchronous processing"
-                            )
-                            return False
-        return True
-
     @staticmethod
     def _get_retriever(
         declarative_stream: DeclarativeStream, stream_state: Mapping[str, Any]
@@ -489,10 +418,21 @@ class ConcurrentDeclarativeSource(ManifestDeclarativeSource, Generic[TState]):
             # Also a temporary hack. In the legacy Stream implementation, as part of the read,
             # set_initial_state() is called to instantiate incoming state on the cursor. Although we no
             # longer rely on the legacy low-code cursor for concurrent checkpointing, low-code components
-            # like StopConditionPaginationStrategyDecorator and ClientSideIncrementalRecordFilterDecorator
-            # still rely on a DatetimeBasedCursor that is properly initialized with state.
+            # like StopConditionPaginationStrategyDecorator still rely on a DatetimeBasedCursor that is
+            # properly initialized with state.
             if retriever.cursor:
                 retriever.cursor.set_initial_state(stream_state=stream_state)
+
+            # Similar to above, the ClientSideIncrementalRecordFilterDecorator cursor is a separate instance
+            # from the one initialized on the SimpleRetriever, so it also must also have state initialized
+            # for semi-incremental streams using is_client_side_incremental to filter properly
+            if isinstance(retriever.record_selector, RecordSelector) and isinstance(
+                retriever.record_selector.record_filter, ClientSideIncrementalRecordFilterDecorator
+            ):
+                retriever.record_selector.record_filter._cursor.set_initial_state(
+                    stream_state=stream_state
+                )  # type: ignore  # After non-concurrent cursors are deprecated we can remove these cursor workarounds
+
             # We zero it out here, but since this is a cursor reference, the state is still properly
             # instantiated for the other components that reference it
             retriever.cursor = None
@@ -524,3 +464,14 @@ class ConcurrentDeclarativeSource(ManifestDeclarativeSource, Generic[TState]):
                 if stream.stream.name not in concurrent_stream_names
             ]
         )
+
+    @staticmethod
+    def _migrate_state(
+        declarative_stream: DeclarativeStream, stream_state: MutableMapping[str, Any]
+    ) -> MutableMapping[str, Any]:
+        for state_migration in declarative_stream.state_migrations:
+            if state_migration.should_migrate(stream_state):
+                # The state variable is expected to be mutable but the migrate method returns an immutable mapping.
+                stream_state = dict(state_migration.migrate(stream_state))
+
+        return stream_state