add max_concurrent_job to async retriever

Author: pnilan

airbyte_cdk/sources/declarative/async_job/job_tracker.py

@@ -3,9 +3,11 @@
 import logging
 import threading
 import uuid
-from typing import Set
+from typing import Any, Mapping, Set, Union
 
+import json
 from airbyte_cdk.logger import lazy_log
+from airbyte_cdk.sources.declarative.interpolation import InterpolatedString
 
 LOGGER = logging.getLogger("airbyte")
 
@@ -15,7 +17,13 @@ class ConcurrentJobLimitReached(Exception):
 
 
 class JobTracker:
-    def __init__(self, limit: int):
+    def __init__(self, limit: Union[int, InterpolatedString], config: Mapping[str, Any]):
+        if isinstance(limit, InterpolatedString):
+            limit = int(limit.eval(config=config, json_loads=json.loads))
+
+        if limit < 1:
+            raise ValueError(f"Invalid max concurrent jobs limit: {limit}. Minimum value is 1.")
+
         self._jobs: Set[str] = set()
         self._limit = limit
         self._lock = threading.Lock()

airbyte_cdk/sources/declarative/declarative_component_schema.yaml

@@ -3376,12 +3376,16 @@ definitions:
           - "$ref": "#/definitions/IterableDecoder"
           - "$ref": "#/definitions/XmlDecoder"
           - "$ref": "#/definitions/ZipfileDecoder"
-      maximum_job_count:
-        title: Maximum Job Count
-        description: Maximum number of asynchronous jobs to run concurrently.
+      max_concurrent_jobs:
+        title: Maximum Conccurent Job Count
+        description: Maximum number of concurrent jobs to run.
         anyOf:
-          - type: number
+          - type: integer
+          - type: string
         default: 1
+        examples:
+          - 2
+          - "{{ config['max_concurrent_jobs'] }}"
       $parameters:
         type: object
         additionalProperties: true

airbyte_cdk/sources/declarative/models/declarative_component_schema.py

@@ -2368,10 +2368,11 @@ class AsyncRetriever(BaseModel):
         description="Component decoding the download response so records can be extracted.",
         title="Download Decoder",
     )
-    maximum_job_count: Optional[float] = Field(
+    max_concurrent_jobs: Optional[Union[int, str]] = Field(
         1,
-        description="Maximum number of asynchronous jobs to run concurrently.",
-        title="Maximum Job Count",
+        description="Maximum number of concurrent jobs to run.",
+        examples=[2, "{{ config['max_concurrent_jobs'] }}"],
+        title="Maximum Conccurent Job Count",
     )
     parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters")
 

airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py

@@ -2888,7 +2888,7 @@ def _get_download_retriever() -> SimpleRetrieverTestReadDecorator | SimpleRetrie
             job_orchestrator_factory=lambda stream_slices: AsyncJobOrchestrator(
                 job_repository,
                 stream_slices,
-                JobTracker(model.maximum_job_count),
+                JobTracker(model.max_concurrent_jobs, config),
                 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