cognite-extractor-utils 7.6.0__py3-none-any.whl → 7.8.0__py3-none-any.whl

cognite/extractorutils/unstable/core/_dto.py

@@ -2,10 +2,16 @@
 Temporary holding place for DTOs against Extraction Pipelines 2.0 until it's in the SDK.
 """
 
-from typing import Any, Literal
+from enum import Enum
+from typing import Annotated, Any, Literal, Optional
 
+from annotated_types import Len
 from humps import camelize
-from pydantic import BaseModel, ConfigDict
+from pydantic import BaseModel, ConfigDict, StringConstraints
+from typing_extensions import TypeAliasType
+
+from cognite.extractorutils.unstable.core.errors import Error as InternalError
+from cognite.extractorutils.unstable.core.errors import ErrorLevel
 
 
 class CogniteModel(BaseModel):
@@ -17,30 +23,97 @@ class CogniteModel(BaseModel):
       * exclude Nones from serialized JSON instead of having nulls in the response text.
     """
 
-    def model_dump(self, *args: Any, **kwargs: Any) -> dict[str, Any]:
+    def model_dump(self, *args: Any, **kwargs: Any) -> dict[str, Any]:  # noqa: ANN401
         if kwargs:
             kwargs["exclude_none"] = True
         else:
             kwargs = {"exclude_none": True}
         return BaseModel.model_dump(self, *args, **kwargs)
 
-    def dict(self, *args: Any, **kwargs: Any) -> dict[str, Any]:
+    def dict(self, *args: Any, **kwargs: Any) -> dict[str, Any]:  # noqa: ANN401
         return self.model_dump(*args, **kwargs)
 
     model_config = ConfigDict(alias_generator=camelize, populate_by_name=True, extra="forbid")
 
 
+class WithExternalId(CogniteModel):
+    external_id: str
+
+
+MessageType = Annotated[str, StringConstraints(min_length=0, max_length=1000)]
+
+
 class TaskUpdate(CogniteModel):
     type: Literal["started"] | Literal["ended"]
     name: str
     timestamp: int
+    message: MessageType | None = None
 
 
-class Error(CogniteModel):
-    external_id: str
-    level: str
+class Error(WithExternalId):
+    level: ErrorLevel
     description: str
     details: str | None
     start_time: int
     end_time: int | None
     task: str | None
+
+    @classmethod
+    def from_internal(cls, error: InternalError) -> "Error":
+        """Convert the error into a DTO (Data Transfer Object) for reporting."""
+        return Error(
+            external_id=error.external_id,
+            level=error.level,
+            description=error.description,
+            details=error.details,
+            start_time=error.start_time,
+            end_time=error.end_time,
+            task=error._task_name,
+        )
+
+
+TaskUpdateList = Annotated[list[TaskUpdate], Len(min_length=1, max_length=1000)]
+ErrorList = Annotated[list[Error], Len(min_length=0, max_length=1000)]
+VersionType = Annotated[str, StringConstraints(min_length=1, max_length=32)]
+DescriptionType = Annotated[str, StringConstraints(min_length=0, max_length=500)]
+TaskList = Annotated[list["Task"], Len(min_length=1, max_length=1000)]
+JSONType = TypeAliasType(  # type: ignore[misc]
+    "JSONType",
+    bool | int | float | str | None | list[Optional["JSONType"]] | dict[str, Optional["JSONType"]],  # type: ignore[misc]
+)
+
+
+class WithVersion(CogniteModel):
+    version: VersionType | None = None
+
+
+class ExtractorInfo(WithExternalId, WithVersion):
+    pass
+
+
+class TaskType(Enum):
+    continuous = "continuous"
+    batch = "batch"
+
+
+class Task(CogniteModel):
+    type: TaskType
+    name: str
+    action: bool = False
+    description: DescriptionType | None = None
+
+
+class StartupRequest(WithExternalId):
+    extractor: ExtractorInfo
+    tasks: TaskList | None = None
+    active_config_revision: int | Literal["local"] | None = None
+    timestamp: int | None = None
+
+
+class CheckinRequest(WithExternalId):
+    task_events: TaskUpdateList | None = None
+    errors: ErrorList | None = None
+
+
+class CheckinResponse(WithExternalId):
+    last_config_revision: int | None = None

cognite/extractorutils/unstable/core/base.py

@@ -47,38 +47,56 @@ The subclass should also define several class attributes:
 import logging
 import time
 from concurrent.futures import ThreadPoolExecutor
+from datetime import datetime, timezone
 from functools import partial
-from logging.handlers import TimedRotatingFileHandler
 from multiprocessing import Queue
+from multiprocessing.synchronize import Event as MpEvent
 from threading import RLock, Thread
 from types import TracebackType
-from typing import Generic, Literal, TypeVar
+from typing import Generic, TypeVar
 
 from humps import pascalize
 from typing_extensions import Self, assert_never
 
 from cognite.extractorutils._inner_util import _resolve_log_level
+from cognite.extractorutils.statestore import (
+    AbstractStateStore,
+    LocalStateStore,
+    NoStateStore,
+)
 from cognite.extractorutils.threading import CancellationToken
 from cognite.extractorutils.unstable.configuration.models import (
+    ConfigRevision,
+    ConfigType,
     ConnectionConfig,
     ExtractorConfig,
     LogConsoleHandlerConfig,
     LogFileHandlerConfig,
 )
-from cognite.extractorutils.unstable.core._dto import Error as DtoError
-from cognite.extractorutils.unstable.core._dto import TaskUpdate
+from cognite.extractorutils.unstable.core._dto import (
+    CogniteModel,
+    ExtractorInfo,
+    StartupRequest,
+    TaskType,
+)
+from cognite.extractorutils.unstable.core._dto import (
+    Task as DtoTask,
+)
 from cognite.extractorutils.unstable.core._messaging import RuntimeMessage
+from cognite.extractorutils.unstable.core.checkin_worker import CheckinWorker
 from cognite.extractorutils.unstable.core.errors import Error, ErrorLevel
-from cognite.extractorutils.unstable.core.logger import CogniteLogger
+from cognite.extractorutils.unstable.core.logger import CogniteLogger, RobustFileHandler
 from cognite.extractorutils.unstable.core.restart_policy import WHEN_CONTINUOUS_TASKS_CRASHES, RestartPolicy
 from cognite.extractorutils.unstable.core.tasks import ContinuousTask, ScheduledTask, StartupTask, Task, TaskContext
 from cognite.extractorutils.unstable.scheduling import TaskScheduler
 from cognite.extractorutils.util import now
 
-__all__ = ["ConfigRevision", "ConfigType", "Extractor"]
-
-ConfigType = TypeVar("ConfigType", bound=ExtractorConfig)
-ConfigRevision = Literal["local"] | int
+__all__ = [
+    "CogniteModel",
+    "ConfigRevision",
+    "ConfigType",
+    "Extractor",
+]
 
 
 _T = TypeVar("_T", bound=ExtractorConfig)
@@ -97,10 +115,12 @@ class FullConfig(Generic[_T]):
         connection_config: ConnectionConfig,
         application_config: _T,
         current_config_revision: ConfigRevision,
+        log_level_override: str | None = None,
     ) -> None:
         self.connection_config = connection_config
         self.application_config = application_config
-        self.current_config_revision = current_config_revision
+        self.current_config_revision: ConfigRevision = current_config_revision
+        self.log_level_override = log_level_override
 
 
 class Extractor(Generic[ConfigType], CogniteLogger):
@@ -122,33 +142,60 @@ class Extractor(Generic[ConfigType], CogniteLogger):
     CONFIG_TYPE: type[ConfigType]
 
     RESTART_POLICY: RestartPolicy = WHEN_CONTINUOUS_TASKS_CRASHES
+    USE_DEFAULT_STATE_STORE: bool = True
+    _statestore_singleton: AbstractStateStore | None = None
+
+    cancellation_token: CancellationToken
 
-    def __init__(self, config: FullConfig[ConfigType]) -> None:
+    def __init__(self, config: FullConfig[ConfigType], checkin_worker: CheckinWorker) -> None:
         self._logger = logging.getLogger(f"{self.EXTERNAL_ID}.main")
+        self._checkin_worker = checkin_worker
 
         self.cancellation_token = CancellationToken()
         self.cancellation_token.cancel_on_interrupt()
 
         self.connection_config = config.connection_config
         self.application_config = config.application_config
-        self.current_config_revision = config.current_config_revision
+        self.current_config_revision: ConfigRevision = config.current_config_revision
+        self.log_level_override = config.log_level_override
 
         self.cognite_client = self.connection_config.get_cognite_client(f"{self.EXTERNAL_ID}-{self.VERSION}")
 
+        self.state_store: AbstractStateStore
+
         self._checkin_lock = RLock()
         self._runtime_messages: Queue[RuntimeMessage] | None = None
 
         self._scheduler = TaskScheduler(self.cancellation_token.create_child_token())
 
         self._tasks: list[Task] = []
-        self._task_updates: list[TaskUpdate] = []
-        self._errors: dict[str, Error] = {}
+        self._start_time: datetime
 
         self.__init_tasks__()
 
+    def _setup_cancellation_watcher(self, cancel_event: MpEvent) -> None:
+        """Starts a daemon thread to watch the inter-process event."""
+
+        def watcher() -> None:
+            """Blocks until the event is set, then cancels the local token."""
+            cancel_event.wait()
+            if not self.cancellation_token.is_cancelled:
+                self._logger.info("Cancellation signal received from runtime. Shutting down gracefully.")
+                self.cancellation_token.cancel()
+
+        thread = Thread(target=watcher, name="RuntimeCancellationWatcher", daemon=True)
+        thread.start()
+
     def _setup_logging(self) -> None:
-        min_level = min([_resolve_log_level(h.level.value) for h in self.application_config.log_handlers])
-        max_level = max([_resolve_log_level(h.level.value) for h in self.application_config.log_handlers])
+        if self.log_level_override:
+            level_to_set = _resolve_log_level(self.log_level_override)
+            # Use the override level if provided
+            min_level = level_to_set
+            max_level = level_to_set
+        else:
+            # Otherwise, use the levels from the config file
+            min_level = min([_resolve_log_level(h.level.value) for h in self.application_config.log_handlers])
+            max_level = max([_resolve_log_level(h.level.value) for h in self.application_config.log_handlers])
 
         root = logging.getLogger()
         root.setLevel(min_level)
@@ -173,21 +220,85 @@ class Extractor(Generic[ConfigType], CogniteLogger):
                 case LogConsoleHandlerConfig() as console_handler:
                     sh = logging.StreamHandler()
                     sh.setFormatter(fmt)
-                    sh.setLevel(_resolve_log_level(console_handler.level.value))
+                    level_for_handler = _resolve_log_level(
+                        self.log_level_override if self.log_level_override else console_handler.level.value
+                    )
+                    sh.setLevel(level_for_handler)
 
                     root.addHandler(sh)
 
                 case LogFileHandlerConfig() as file_handler:
-                    fh = TimedRotatingFileHandler(
-                        filename=file_handler.path,
-                        when="midnight",
-                        utc=True,
-                        backupCount=file_handler.retention,
-                    )
-                    fh.setLevel(_resolve_log_level(file_handler.level.value))
-                    fh.setFormatter(fmt)
+                    try:
+                        fh = RobustFileHandler(
+                            filename=file_handler.path,
+                            when="midnight",
+                            utc=True,
+                            backupCount=file_handler.retention,
+                            create_dirs=True,
+                        )
+                        level_for_handler = _resolve_log_level(
+                            self.log_level_override if self.log_level_override else file_handler.level.value
+                        )
+                        fh.setLevel(level_for_handler)
+                        fh.setFormatter(fmt)
+
+                        root.addHandler(fh)
+                    except (OSError, PermissionError) as e:
+                        self._logger.warning(
+                            f"Could not create or write to log file {file_handler.path}: {e}. "
+                            "Falling back to console logging."
+                        )
+                        if not any(isinstance(h, logging.StreamHandler) for h in root.handlers):
+                            sh = logging.StreamHandler()
+                            sh.setFormatter(fmt)
+                            sh.setLevel(level_for_handler)
+                            root.addHandler(sh)
+
+    def _load_state_store(self) -> None:
+        """
+        Searches through the config object for a StateStoreConfig.
+
+        If found, it will use that configuration to generate a state store, if no such config is found it will either
+        create a LocalStateStore or a NoStateStore depending on whether the ``use_default_state_store`` argument to the
+        constructor was true or false.
+
+        Either way, the state_store attribute is guaranteed to be set after calling this method.
+        """
+        state_store_config = self.application_config.state_store
+
+        if state_store_config:
+            self.state_store = state_store_config.create_state_store(
+                cdf_client=self.cognite_client,
+                default_to_local=self.USE_DEFAULT_STATE_STORE,
+                cancellation_token=self.cancellation_token,
+            )
+        elif self.USE_DEFAULT_STATE_STORE:
+            self.state_store = LocalStateStore("states.json", cancellation_token=self.cancellation_token)
+        else:
+            self.state_store = NoStateStore()
+
+        try:
+            self.state_store.initialize()
+        except ValueError:
+            self._logger.exception("Could not load state store, using an empty state store as default")
+            self.state_store = NoStateStore()
 
-                    root.addHandler(fh)
+        Extractor._statestore_singleton = self.state_store
+
+    @classmethod
+    def get_current_statestore(cls) -> AbstractStateStore:
+        """
+        Get the current state store singleton.
+
+        Returns:
+            The current state store singleton
+
+        Raises:
+            ValueError: If no state store singleton has been created, meaning no state store has been loaded.
+        """
+        if Extractor._statestore_singleton is None:
+            raise ValueError("No state store singleton created. Have a state store been loaded?")
+        return Extractor._statestore_singleton
 
     def __init_tasks__(self) -> None:
         """
@@ -202,55 +313,37 @@ class Extractor(Generic[ConfigType], CogniteLogger):
     def _set_runtime_message_queue(self, queue: Queue) -> None:
         self._runtime_messages = queue
 
-    def _checkin(self) -> None:
-        with self._checkin_lock:
-            task_updates = [t.model_dump() for t in self._task_updates]
-            self._task_updates.clear()
-
-            error_updates = [
-                DtoError(
-                    external_id=e.external_id,
-                    level=e.level.value,
-                    description=e.description,
-                    details=e.details,
-                    start_time=e.start_time,
-                    end_time=e.end_time,
-                    task=e._task_name if e._task_name is not None else None,
-                ).model_dump()
-                for e in self._errors.values()
+    def _attach_runtime_controls(self, *, cancel_event: MpEvent, message_queue: Queue) -> None:
+        self._set_runtime_message_queue(message_queue)
+        self._setup_cancellation_watcher(cancel_event)
+
+    def _get_startup_request(self) -> StartupRequest:
+        return StartupRequest(
+            external_id=self.connection_config.integration.external_id,
+            active_config_revision=self.current_config_revision,
+            extractor=ExtractorInfo(version=self.VERSION, external_id=self.EXTERNAL_ID),
+            tasks=[
+                DtoTask(
+                    type=TaskType.continuous if isinstance(t, ContinuousTask) else TaskType.batch,
+                    action=isinstance(t, ScheduledTask),
+                    description=t.description,
+                    name=t.name,
+                )
+                for t in self._tasks
             ]
-            self._errors.clear()
-
-        res = self.cognite_client.post(
-            f"/api/v1/projects/{self.cognite_client.config.project}/integrations/checkin",
-            json={
-                "externalId": self.connection_config.integration,
-                "taskEvents": task_updates,
-                "errors": error_updates,
-            },
-            headers={"cdf-version": "alpha"},
+            if len(self._tasks) > 0
+            else None,
+            timestamp=int(self._start_time.timestamp() * 1000),
         )
-        new_config_revision = res.json().get("lastConfigRevision")
-
-        if (
-            new_config_revision
-            and self.current_config_revision != "local"
-            and new_config_revision > self.current_config_revision
-        ):
-            self.restart()
 
     def _run_checkin(self) -> None:
-        while not self.cancellation_token.is_cancelled:
-            try:
-                self._logger.debug("Running checkin")
-                self._checkin()
-            except Exception:
-                self._logger.exception("Error during checkin")
-            self.cancellation_token.wait(10)
+        self._checkin_worker.run_periodic_checkin(self.cancellation_token, self._get_startup_request())
 
     def _report_error(self, error: Error) -> None:
-        with self._checkin_lock:
-            self._errors[error.external_id] = error
+        self._checkin_worker.report_error(error)
+
+    def _try_report_error(self, error: Error) -> None:
+        self._checkin_worker.try_report_error(error)
 
     def _new_error(
         self,
@@ -278,8 +371,8 @@ class Extractor(Generic[ConfigType], CogniteLogger):
         self.cancellation_token.cancel()
 
     @classmethod
-    def _init_from_runtime(cls, config: FullConfig[ConfigType]) -> Self:
-        return cls(config)
+    def _init_from_runtime(cls, config: FullConfig[ConfigType], checkin_worker: CheckinWorker) -> Self:
+        return cls(config, checkin_worker)
 
     def add_task(self, task: Task) -> None:
         """
@@ -298,10 +391,7 @@ class Extractor(Generic[ConfigType], CogniteLogger):
             A wrapped version of the task's target, with tracking and error handling.
             """
             # Record a task start
-            with self._checkin_lock:
-                self._task_updates.append(
-                    TaskUpdate(type="started", name=task.name, timestamp=now()),
-                )
+            self._checkin_worker.report_task_start(name=task.name, timestamp=now())
 
             try:
                 # Run task
@@ -320,10 +410,7 @@ class Extractor(Generic[ConfigType], CogniteLogger):
 
             finally:
                 # Record task end
-                with self._checkin_lock:
-                    self._task_updates.append(
-                        TaskUpdate(type="ended", name=task.name, timestamp=now()),
-                    )
+                self._checkin_worker.report_task_end(name=task.name, timestamp=now())
 
         task.target = run_task
         self._tasks.append(task)
@@ -341,29 +428,6 @@ class Extractor(Generic[ConfigType], CogniteLogger):
                     ),
                 )
 
-    def _report_extractor_info(self) -> None:
-        self.cognite_client.post(
-            f"/api/v1/projects/{self.cognite_client.config.project}/integrations/extractorinfo",
-            json={
-                "externalId": self.connection_config.integration,
-                "activeConfigRevision": self.current_config_revision,
-                "extractor": {
-                    "version": self.VERSION,
-                    "externalId": self.EXTERNAL_ID,
-                },
-                "tasks": [
-                    {
-                        "name": t.name,
-                        "type": "continuous" if isinstance(t, ContinuousTask) else "batch",
-                        "action": bool(isinstance(t, ScheduledTask)),
-                        "description": t.description,
-                    }
-                    for t in self._tasks
-                ],
-            },
-            headers={"cdf-version": "alpha"},
-        )
-
     def start(self) -> None:
         """
         Start the extractor.
@@ -372,7 +436,11 @@ class Extractor(Generic[ConfigType], CogniteLogger):
         ``with`` statement, which ensures proper cleanup on exit.
         """
         self._setup_logging()
-        self._report_extractor_info()
+        self._start_time = datetime.now(tz=timezone.utc)
+
+        self._load_state_store()
+        self.state_store.start()
+
         Thread(target=self._run_checkin, name="ExtractorCheckin", daemon=True).start()
 
     def stop(self) -> None:
@@ -401,10 +469,11 @@ class Extractor(Generic[ConfigType], CogniteLogger):
         Stop the extractor when exiting the context manager.
         """
         self.stop()
-        with self._checkin_lock:
-            self._checkin()
 
-        self._logger.info("Shutting down extractor")
+        if self.state_store:
+            self.state_store.synchronize()
+
+        self._checkin_worker.flush(self.cancellation_token)
         return exc_val is None
 
     def run(self) -> None: