feat(airbyte-cdk) - Async job salesforce (airbytehq#45673)

Author: maxi297

airbyte-cdk/python/airbyte_cdk/sources/declarative/async_job/job.py

@@ -47,4 +47,4 @@ def update_status(self, status: AsyncJobStatus) -> None:
         self._status = status
 
     def __repr__(self) -> str:
-        return f"AsyncJob(data={self.api_job_id()}, job_parameters={self.job_parameters()}, status={self.status()})"
+        return f"AsyncJob(api_job_id={self.api_job_id()}, job_parameters={self.job_parameters()}, status={self.status()})"

airbyte-cdk/python/airbyte_cdk/sources/declarative/async_job/job_orchestrator.py

@@ -0,0 +1,57 @@
+# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+
+import logging
+import threading
+import uuid
+from typing import Set
+
+from airbyte_cdk.logger import lazy_log
+
+LOGGER = logging.getLogger("airbyte")
+
+
+class ConcurrentJobLimitReached(Exception):
+    pass
+
+
+class JobTracker:
+    def __init__(self, limit: int):
+        self._jobs: Set[str] = set()
+        self._limit = limit
+        self._lock = threading.Lock()
+
+    def try_to_get_intent(self) -> str:
+        lazy_log(LOGGER, logging.DEBUG, lambda: f"JobTracker - Trying to acquire lock by thread {threading.get_native_id()}...")
+        with self._lock:
+            if self._has_reached_limit():
+                raise ConcurrentJobLimitReached("Can't allocate more jobs right now: limit already reached")
+            intent = f"intent_{str(uuid.uuid4())}"
+            lazy_log(LOGGER, logging.DEBUG, lambda: f"JobTracker - Thread {threading.get_native_id()} has acquired {intent}!")
+            self._jobs.add(intent)
+            return intent
+
+    def add_job(self, intent_or_job_id: str, job_id: str) -> None:
+        if intent_or_job_id not in self._jobs:
+            raise ValueError(f"Can't add job: Unknown intent or job id, known values are {self._jobs}")
+
+        if intent_or_job_id == job_id:
+            # Nothing to do here as the ID to replace is the same
+            return
+
+        lazy_log(
+            LOGGER, logging.DEBUG, lambda: f"JobTracker - Thread {threading.get_native_id()} replacing job {intent_or_job_id} by {job_id}!"
+        )
+        with self._lock:
+            self._jobs.add(job_id)
+            self._jobs.remove(intent_or_job_id)
+
+    def remove_job(self, job_id: str) -> None:
+        """
+        If the job is not allocated as a running job, this method does nothing and it won't raise.
+        """
+        lazy_log(LOGGER, logging.DEBUG, lambda: f"JobTracker - Thread {threading.get_native_id()} removing job {job_id}")
+        with self._lock:
+            self._jobs.discard(job_id)
+
+    def _has_reached_limit(self) -> bool:
+        return len(self._jobs) >= self._limit

airbyte-cdk/python/airbyte_cdk/sources/declarative/async_job/job_tracker.py

@@ -19,3 +19,15 @@ def update_jobs_status(self, jobs: Set[AsyncJob]) -> None:
     @abstractmethod
     def fetch_records(self, job: AsyncJob) -> Iterable[Mapping[str, Any]]:
         pass
+
+    @abstractmethod
+    def abort(self, job: AsyncJob) -> None:
+        """
+        Called when we need to stop on the API side. This method can raise NotImplementedError as not all the APIs will support aborting
+        jobs.
+        """
+        raise NotImplementedError("Either the API or the AsyncJobRepository implementation do not support aborting jobs")
+
+    @abstractmethod
+    def delete(self, job: AsyncJob) -> None:
+        pass

airbyte-cdk/python/airbyte_cdk/sources/declarative/async_job/repository.py

@@ -2435,6 +2435,21 @@ definitions:
         anyOf:
           - "$ref": "#/definitions/CustomRequester"
           - "$ref": "#/definitions/HttpRequester"
+      download_paginator:
+        description: Paginator component that describes how to navigate through the API's pages during download.
+        anyOf:
+          - "$ref": "#/definitions/DefaultPaginator"
+          - "$ref": "#/definitions/NoPagination"
+      abort_requester:
+        description: Requester component that describes how to prepare HTTP requests to send to the source API to abort a job once it is timed out from the source's perspective.
+        anyOf:
+          - "$ref": "#/definitions/CustomRequester"
+          - "$ref": "#/definitions/HttpRequester"
+      delete_requester:
+        description: Requester component that describes how to prepare HTTP requests to send to the source API to delete a job once the records are extracted.
+        anyOf:
+          - "$ref": "#/definitions/CustomRequester"
+          - "$ref": "#/definitions/HttpRequester"
       partition_router:
         title: Partition Router
         description: PartitionRouter component that describes how to partition the stream, enabling incremental syncs and checkpointing.

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

@@ -4,5 +4,7 @@
 
 from airbyte_cdk.sources.declarative.decoders.decoder import Decoder
 from airbyte_cdk.sources.declarative.decoders.json_decoder import JsonDecoder, JsonlDecoder, IterableDecoder
+from airbyte_cdk.sources.declarative.decoders.noop_decoder import NoopDecoder
 
-__all__ = ["Decoder", "JsonDecoder", "JsonlDecoder", "IterableDecoder"]
+
+__all__ = ["Decoder", "JsonDecoder", "JsonlDecoder", "IterableDecoder", "NoopDecoder"]

airbyte-cdk/python/airbyte_cdk/sources/declarative/decoders/__init__.py

@@ -0,0 +1,17 @@
+# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+
+import logging
+from typing import Any, Generator, Mapping
+
+import requests
+from airbyte_cdk.sources.declarative.decoders.decoder import Decoder
+
+logger = logging.getLogger("airbyte")
+
+
+class NoopDecoder(Decoder):
+    def is_stream_response(self) -> bool:
+        return False
+
+    def decode(self, response: requests.Response) -> Generator[Mapping[str, Any], None, None]:
+        yield from [{}]

airbyte-cdk/python/airbyte_cdk/sources/declarative/decoders/noop_decoder.py

@@ -1653,6 +1653,18 @@ class AsyncRetriever(BaseModel):
         ...,
         description='Requester component that describes how to prepare HTTP requests to send to the source API to download the data provided by the completed async job.',
     )
+    download_paginator: Optional[Union[DefaultPaginator, NoPagination]] = Field(
+        None,
+        description="Paginator component that describes how to navigate through the API's pages during download.",
+    )
+    abort_requester: Optional[Union[CustomRequester, HttpRequester]] = Field(
+        None,
+        description="Requester component that describes how to prepare HTTP requests to send to the source API to abort a job once it is timed out from the source's perspective.",
+    )
+    delete_requester: Optional[Union[CustomRequester, HttpRequester]] = Field(
+        None,
+        description='Requester component that describes how to prepare HTTP requests to send to the source API to delete a job once the records are extracted.',
+    )
     partition_router: Optional[
         Union[
             CustomPartitionRouter,

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

@@ -12,6 +12,7 @@
 
 from airbyte_cdk.models import FailureType, Level
 from airbyte_cdk.sources.declarative.async_job.job_orchestrator import AsyncJobOrchestrator
+from airbyte_cdk.sources.declarative.async_job.job_tracker import JobTracker
 from airbyte_cdk.sources.declarative.async_job.repository import AsyncJobRepository
 from airbyte_cdk.sources.declarative.async_job.status import AsyncJobStatus
 from airbyte_cdk.sources.declarative.auth import DeclarativeOauth2Authenticator, JwtAuthenticator
@@ -30,7 +31,7 @@
 from airbyte_cdk.sources.declarative.datetime import MinMaxDatetime
 from airbyte_cdk.sources.declarative.declarative_stream import DeclarativeStream
 from airbyte_cdk.sources.declarative.decoders import Decoder, IterableDecoder, JsonDecoder, JsonlDecoder
-from airbyte_cdk.sources.declarative.extractors import DpathExtractor, RecordFilter, RecordSelector
+from airbyte_cdk.sources.declarative.extractors import DpathExtractor, RecordFilter, RecordSelector, ResponseToFileExtractor
 from airbyte_cdk.sources.declarative.extractors.record_filter import ClientSideIncrementalRecordFilterDecorator
 from airbyte_cdk.sources.declarative.extractors.record_selector import SCHEMA_TRANSFORMER_TYPE_MAPPING
 from airbyte_cdk.sources.declarative.incremental import (
@@ -158,7 +159,7 @@
 from airbyte_cdk.sources.message import InMemoryMessageRepository, LogAppenderMessageRepositoryDecorator, MessageRepository
 from airbyte_cdk.sources.streams.http.error_handlers.response_models import ResponseAction
 from airbyte_cdk.sources.types import Config
-from airbyte_cdk.sources.utils.transform import TypeTransformer
+from airbyte_cdk.sources.utils.transform import TransformConfig, TypeTransformer
 from isodate import parse_duration
 from pydantic.v1 import BaseModel
 
@@ -1298,22 +1299,59 @@ def create_async_retriever(
         polling_requester = self._create_component_from_model(
             model=model.polling_requester, decoder=decoder, config=config, name=f"job polling - {name}"
         )
+        job_download_components_name = f"job download - {name}"
         download_requester = self._create_component_from_model(
-            model=model.download_requester, decoder=decoder, config=config, name=f"job download - {name}"
+            model=model.download_requester, decoder=decoder, config=config, name=job_download_components_name
+        )
+        download_retriever = SimpleRetriever(
+            requester=download_requester,
+            record_selector=RecordSelector(
+                extractor=ResponseToFileExtractor(),
+                record_filter=None,
+                transformations=[],
+                schema_normalization=TypeTransformer(TransformConfig.NoTransform),
+                config=config,
+                parameters={},
+            ),
+            primary_key=None,
+            name=job_download_components_name,
+            paginator=self._create_component_from_model(model=model.download_paginator, decoder=decoder, config=config, url_base="")
+            if model.download_paginator
+            else NoPagination(parameters={}),
+            config=config,
+            parameters={},
+        )
+        abort_requester = (
+            self._create_component_from_model(model=model.abort_requester, decoder=decoder, config=config, name=f"job abort - {name}")
+            if model.abort_requester
+            else None
+        )
+        delete_requester = (
+            self._create_component_from_model(model=model.delete_requester, decoder=decoder, config=config, name=f"job delete - {name}")
+            if model.delete_requester
+            else None
         )
         status_extractor = self._create_component_from_model(model=model.status_extractor, decoder=decoder, config=config, name=name)
         urls_extractor = self._create_component_from_model(model=model.urls_extractor, decoder=decoder, config=config, name=name)
         job_repository: AsyncJobRepository = AsyncHttpJobRepository(
             creation_requester=creation_requester,
             polling_requester=polling_requester,
-            download_requester=download_requester,
+            download_retriever=download_retriever,
+            abort_requester=abort_requester,
+            delete_requester=delete_requester,
             status_extractor=status_extractor,
             status_mapping=self._create_async_job_status_mapping(model.status_mapping, config),
             urls_extractor=urls_extractor,
         )
 
         return AsyncRetriever(
-            job_orchestrator_factory=lambda stream_slices: AsyncJobOrchestrator(job_repository, stream_slices),
+            job_orchestrator_factory=lambda stream_slices: AsyncJobOrchestrator(
+                job_repository,
+                stream_slices,
+                JobTracker(1),  # FIXME eventually make the number of concurrent jobs in the API configurable. Until then, we limit to 1
+                self._message_repository,
+                has_bulk_parent=False,  # FIXME work would need to be done here in order to detect if a stream as a parent stream that is bulk
+            ),
             record_selector=record_selector,
             stream_slicer=stream_slicer,
             config=config,

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

@@ -2,18 +2,21 @@
 import logging
 import uuid
 from dataclasses import dataclass, field
+from datetime import timedelta
 from typing import Any, Dict, Iterable, Mapping, Optional
 
 import requests
+from airbyte_cdk import AirbyteMessage
 from airbyte_cdk.logger import lazy_log
-from airbyte_cdk.models import FailureType
+from airbyte_cdk.models import FailureType, Type
 from airbyte_cdk.sources.declarative.async_job.job import AsyncJob
 from airbyte_cdk.sources.declarative.async_job.repository import AsyncJobRepository
 from airbyte_cdk.sources.declarative.async_job.status import AsyncJobStatus
 from airbyte_cdk.sources.declarative.extractors.dpath_extractor import DpathExtractor, RecordExtractor
 from airbyte_cdk.sources.declarative.extractors.response_to_file_extractor import ResponseToFileExtractor
 from airbyte_cdk.sources.declarative.requesters.requester import Requester
-from airbyte_cdk.sources.types import StreamSlice
+from airbyte_cdk.sources.declarative.retrievers.simple_retriever import SimpleRetriever
+from airbyte_cdk.sources.types import Record, StreamSlice
 from airbyte_cdk.utils import AirbyteTracedException
 from requests import Response
 
@@ -24,11 +27,14 @@
 class AsyncHttpJobRepository(AsyncJobRepository):
     creation_requester: Requester
     polling_requester: Requester
-    download_requester: Requester
+    download_retriever: SimpleRetriever
+    abort_requester: Optional[Requester]
+    delete_requester: Optional[Requester]
     status_extractor: DpathExtractor
     status_mapping: Mapping[str, AsyncJobStatus]
     urls_extractor: DpathExtractor
 
+    job_timeout: Optional[timedelta] = None
     record_extractor: RecordExtractor = field(init=False, repr=False, default_factory=lambda: ResponseToFileExtractor())
 
     def __post_init__(self) -> None:
@@ -118,7 +124,7 @@ def start(self, stream_slice: StreamSlice) -> AsyncJob:
         job_id: str = str(uuid.uuid4())
         self._create_job_response_by_id[job_id] = response
 
-        return AsyncJob(api_job_id=job_id, job_parameters=stream_slice)
+        return AsyncJob(api_job_id=job_id, job_parameters=stream_slice, timeout=self.job_timeout)
 
     def update_jobs_status(self, jobs: Iterable[AsyncJob]) -> None:
         """
@@ -135,15 +141,14 @@ def update_jobs_status(self, jobs: Iterable[AsyncJob]) -> None:
             None
         """
         for job in jobs:
-            stream_slice = StreamSlice(
-                partition={"create_job_response": self._create_job_response_by_id[job.api_job_id()]},
-                cursor_slice={},
-            )
+            stream_slice = self._get_create_job_stream_slice(job)
             polling_response: requests.Response = self._get_validated_polling_response(stream_slice)
             job_status: AsyncJobStatus = self._get_validated_job_status(polling_response)
 
             if job_status != job.status():
                 lazy_log(LOGGER, logging.DEBUG, lambda: f"Status of job {job.api_job_id()} changed from {job.status()} to {job_status}")
+            else:
+                lazy_log(LOGGER, logging.DEBUG, lambda: f"Status of job {job.api_job_id()} is still {job.status()}")
 
             job.update_status(job_status)
             if job_status == AsyncJobStatus.COMPLETED:
@@ -163,15 +168,39 @@ def fetch_records(self, job: AsyncJob) -> Iterable[Mapping[str, Any]]:
 
         for url in self.urls_extractor.extract_records(self._polling_job_response_by_id[job.api_job_id()]):
             stream_slice: StreamSlice = StreamSlice(partition={"url": url}, cursor_slice={})
-            # FIXME salesforce will require pagination here
-            response = self.download_requester.send_request(stream_slice=stream_slice)
-            if response:
-                yield from self.record_extractor.extract_records(response)
+            for message in self.download_retriever.read_records({}, stream_slice):
+                if isinstance(message, Record):
+                    yield message.data
+                elif isinstance(message, AirbyteMessage):
+                    if message.type == Type.RECORD:
+                        yield message.record.data  # type: ignore  # message.record won't be None here as the message is a record
+                elif isinstance(message, (dict, Mapping)):
+                    yield message
+                else:
+                    raise TypeError(f"Unknown type `{type(message)}` for message")
 
         yield from []
 
+    def abort(self, job: AsyncJob) -> None:
+        if not self.abort_requester:
+            return
+
+        self.abort_requester.send_request(stream_slice=self._get_create_job_stream_slice(job))
+
+    def delete(self, job: AsyncJob) -> None:
+        if not self.delete_requester:
+            return
+
+        self.delete_requester.send_request(stream_slice=self._get_create_job_stream_slice(job))
         self._clean_up_job(job.api_job_id())
 
     def _clean_up_job(self, job_id: str) -> None:
         del self._create_job_response_by_id[job_id]
         del self._polling_job_response_by_id[job_id]
+
+    def _get_create_job_stream_slice(self, job: AsyncJob) -> StreamSlice:
+        stream_slice = StreamSlice(
+            partition={"create_job_response": self._create_job_response_by_id[job.api_job_id()]},
+            cursor_slice={},
+        )
+        return stream_slice

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

@@ -17,7 +17,9 @@ class HttpMethod(Enum):
     Http Method to use when submitting an outgoing HTTP request
     """
 
+    DELETE = "DELETE"
     GET = "GET"
+    PATCH = "PATCH"
     POST = "POST"