airbyte_cdk.sources.declarative.async_job.job_tracker

 1# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
 2
 3import logging
 4import threading
 5import uuid
 6from dataclasses import dataclass, field
 7from typing import Any, Mapping, Set, Union
 8
 9from airbyte_cdk.logger import lazy_log
10from airbyte_cdk.sources.declarative.interpolation import InterpolatedString
11
12LOGGER = logging.getLogger("airbyte")
13
14
15class ConcurrentJobLimitReached(Exception):
16    pass
17
18
19@dataclass
20class JobTracker:
21    limit: Union[int, str]
22    config: Mapping[str, Any] = field(default_factory=dict)
23
24    def __post_init__(self) -> None:
25        self._jobs: Set[str] = set()
26        self._lock = threading.Lock()
27        if isinstance(self.limit, str):
28            try:
29                self.limit = int(
30                    InterpolatedString(self.limit, parameters={}).eval(config=self.config)
31                )
32            except Exception as e:
33                LOGGER.warning(
34                    f"Error interpolating max job count: {self.limit}. Setting to 1. {e}"
35                )
36                self.limit = 1
37        if self.limit < 1:
38            LOGGER.warning(
39                f"The `max_concurrent_async_job_count` property is less than 1: {self.limit}. Setting to 1. Please update the source manifest to set a valid value."
40            )
41        self._limit = self.limit if self.limit >= 1 else 1
42
43    def try_to_get_intent(self) -> str:
44        lazy_log(
45            LOGGER,
46            logging.DEBUG,
47            lambda: f"JobTracker - Trying to acquire lock by thread {threading.get_native_id()}...",
48        )
49        with self._lock:
50            if self._has_reached_limit():
51                raise ConcurrentJobLimitReached(
52                    "Can't allocate more jobs right now: limit already reached"
53                )
54            intent = f"intent_{str(uuid.uuid4())}"
55            lazy_log(
56                LOGGER,
57                logging.DEBUG,
58                lambda: f"JobTracker - Thread {threading.get_native_id()} has acquired {intent}!",
59            )
60            self._jobs.add(intent)
61            return intent
62
63    def add_job(self, intent_or_job_id: str, job_id: str) -> None:
64        if intent_or_job_id not in self._jobs:
65            raise ValueError(
66                f"Can't add job: Unknown intent or job id, known values are {self._jobs}"
67            )
68
69        if intent_or_job_id == job_id:
70            # Nothing to do here as the ID to replace is the same
71            return
72
73        lazy_log(
74            LOGGER,
75            logging.DEBUG,
76            lambda: f"JobTracker - Thread {threading.get_native_id()} replacing job {intent_or_job_id} by {job_id}!",
77        )
78        with self._lock:
79            self._jobs.add(job_id)
80            self._jobs.remove(intent_or_job_id)
81
82    def remove_job(self, job_id: str) -> None:
83        """
84        If the job is not allocated as a running job, this method does nothing and it won't raise.
85        """
86        lazy_log(
87            LOGGER,
88            logging.DEBUG,
89            lambda: f"JobTracker - Thread {threading.get_native_id()} removing job {job_id}",
90        )
91        with self._lock:
92            self._jobs.discard(job_id)
93
94    def _has_reached_limit(self) -> bool:
95        return len(self._jobs) >= self._limit