cognite-extractor-utils 7.5.7__py3-none-any.whl → 7.5.9__py3-none-any.whl

cognite/extractorutils/unstable/core/base.py

@@ -2,13 +2,12 @@ import logging
 import logging.config
 import time
 from concurrent.futures import ThreadPoolExecutor
-from contextvars import ContextVar, Token
+from functools import partial
 from logging.handlers import TimedRotatingFileHandler
 from multiprocessing import Queue
 from threading import RLock, Thread
-from traceback import format_exception
 from types import TracebackType
-from typing import Generic, Literal, Type, TypeVar
+from typing import Generic, Literal, TypeVar
 
 from humps import pascalize
 from typing_extensions import Self, assert_never
@@ -25,8 +24,9 @@ from cognite.extractorutils.unstable.core._dto import Error as DtoError
 from cognite.extractorutils.unstable.core._dto import TaskUpdate
 from cognite.extractorutils.unstable.core._messaging import RuntimeMessage
 from cognite.extractorutils.unstable.core.errors import Error, ErrorLevel
+from cognite.extractorutils.unstable.core.logger import CogniteLogger
 from cognite.extractorutils.unstable.core.restart_policy import WHEN_CONTINUOUS_TASKS_CRASHES, RestartPolicy
-from cognite.extractorutils.unstable.core.tasks import ContinuousTask, ScheduledTask, StartupTask, Task
+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
 
@@ -45,32 +45,31 @@ class FullConfig(Generic[_T]):
         connection_config: ConnectionConfig,
         application_config: _T,
         current_config_revision: ConfigRevision,
-        newest_config_revision: ConfigRevision,
     ) -> None:
         self.connection_config = connection_config
         self.application_config = application_config
         self.current_config_revision = current_config_revision
-        self.newest_config_revision = newest_config_revision
 
 
-class Extractor(Generic[ConfigType]):
+class Extractor(Generic[ConfigType], CogniteLogger):
     NAME: str
     EXTERNAL_ID: str
     DESCRIPTION: str
     VERSION: str
 
-    CONFIG_TYPE: Type[ConfigType]
+    CONFIG_TYPE: type[ConfigType]
 
     RESTART_POLICY: RestartPolicy = WHEN_CONTINUOUS_TASKS_CRASHES
 
     def __init__(self, config: FullConfig[ConfigType]) -> None:
+        self._logger = logging.getLogger(f"{self.EXTERNAL_ID}.main")
+
         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.newest_config_revision = config.newest_config_revision
 
         self.cognite_client = self.connection_config.get_cognite_client(f"{self.EXTERNAL_ID}-{self.VERSION}")
 
@@ -83,10 +82,6 @@ class Extractor(Generic[ConfigType]):
         self._task_updates: list[TaskUpdate] = []
         self._errors: dict[str, Error] = {}
 
-        self.logger = logging.getLogger(f"{self.EXTERNAL_ID}.main")
-
-        self._current_task: ContextVar[str | None] = ContextVar("current_task", default=None)
-
         self.__init_tasks__()
 
     def _setup_logging(self) -> None:
@@ -158,7 +153,7 @@ class Extractor(Generic[ConfigType]):
             self._errors.clear()
 
         res = self.cognite_client.post(
-            f"/api/v1/projects/{self.cognite_client.config.project}/odin/checkin",
+            f"/api/v1/projects/{self.cognite_client.config.project}/integrations/checkin",
             json={
                 "externalId": self.connection_config.integration,
                 "taskEvents": task_updates,
@@ -171,43 +166,41 @@ class Extractor(Generic[ConfigType]):
         if (
             new_config_revision
             and self.current_config_revision != "local"
-            and new_config_revision > self.newest_config_revision
+            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._logger.debug("Running checkin")
                 self._checkin()
             except Exception:
-                self.logger.exception("Error during checkin")
+                self._logger.exception("Error during checkin")
             self.cancellation_token.wait(10)
 
     def _report_error(self, error: Error) -> None:
         with self._checkin_lock:
             self._errors[error.external_id] = error
 
-    def error(
+    def _new_error(
         self,
         level: ErrorLevel,
         description: str,
-        details: str | None = None,
         *,
-        force_global: bool = False,
+        details: str | None = None,
+        task_name: str | None = None,
     ) -> Error:
-        task_name = self._current_task.get()
-
         return Error(
             level=level,
             description=description,
             details=details,
             extractor=self,
-            task_name=None if force_global else task_name,
+            task_name=task_name,
         )
 
     def restart(self) -> None:
-        self.logger.info("Restarting extractor")
+        self._logger.info("Restarting extractor")
         if self._runtime_messages:
             self._runtime_messages.put(RuntimeMessage.RESTART)
         self.cancellation_token.cancel()
@@ -220,7 +213,7 @@ class Extractor(Generic[ConfigType]):
         # Store this for later, since we'll override it with the wrapped version
         target = task.target
 
-        def run_task() -> None:
+        def run_task(task_context: TaskContext) -> None:
             """
             A wrapped version of the task's target, with tracking and error handling
             """
@@ -230,33 +223,22 @@ class Extractor(Generic[ConfigType]):
                     TaskUpdate(type="started", name=task.name, timestamp=now()),
                 )
 
-            context_token: Token[str | None] | None = None
-
             try:
-                # Set the current task context var, used to track that we're in a task for error reporting
-                context_token = self._current_task.set(task.name)
-
                 # Run task
-                target()
+                target(task_context)
 
             except Exception as e:
-                self.logger.exception(f"Unexpected error in {task.name}")
-
                 # Task crashed, record it as a fatal error
-                self.error(
-                    ErrorLevel.fatal,
-                    description="Task crashed unexpectedly",
-                    details="".join(format_exception(e)),
-                ).instant()
+                task_context.exception(
+                    f"Task {task.name} crashed unexpectedly",
+                    e,
+                    level=ErrorLevel.fatal,
+                )
 
                 if self.__class__.RESTART_POLICY(task, e):
                     self.restart()
 
             finally:
-                # Unset the current task
-                if context_token is not None:
-                    self._current_task.reset(context_token)
-
                 # Record task end
                 with self._checkin_lock:
                     self._task_updates.append(
@@ -268,11 +250,20 @@ class Extractor(Generic[ConfigType]):
 
         match task:
             case ScheduledTask() as t:
-                self._scheduler.schedule_task(name=t.name, schedule=t.schedule, task=t.target)
+                self._scheduler.schedule_task(
+                    name=t.name,
+                    schedule=t.schedule,
+                    task=lambda: t.target(
+                        TaskContext(
+                            task=task,
+                            extractor=self,
+                        )
+                    ),
+                )
 
     def _report_extractor_info(self) -> None:
         self.cognite_client.post(
-            f"/api/v1/projects/{self.cognite_client.config.project}/odin/extractorinfo",
+            f"/api/v1/projects/{self.cognite_client.config.project}/integrations/extractorinfo",
             json={
                 "externalId": self.connection_config.integration,
                 "activeConfigRevision": self.current_config_revision,
@@ -284,6 +275,8 @@ class Extractor(Generic[ConfigType]):
                     {
                         "name": t.name,
                         "type": "continuous" if isinstance(t, ContinuousTask) else "batch",
+                        "action": True if isinstance(t, ScheduledTask) else False,
+                        "description": t.description,
                     }
                     for t in self._tasks
                 ],
@@ -305,7 +298,7 @@ class Extractor(Generic[ConfigType]):
 
     def __exit__(
         self,
-        exc_type: Type[BaseException] | None,
+        exc_type: type[BaseException] | None,
         exc_val: BaseException | None,
         exc_tb: TracebackType | None,
     ) -> bool:
@@ -313,7 +306,7 @@ class Extractor(Generic[ConfigType]):
         with self._checkin_lock:
             self._checkin()
 
-        self.logger.info("Shutting down extractor")
+        self._logger.info("Shutting down extractor")
         return exc_val is None
 
     def run(self) -> None:
@@ -336,15 +329,27 @@ class Extractor(Generic[ConfigType]):
                 case _:
                     assert_never(task)
 
-        self.logger.info("Starting extractor")
+        self._logger.info("Starting extractor")
         if startup:
             with ThreadPoolExecutor() as pool:
                 for task in startup:
-                    pool.submit(task.target)
-        self.logger.info("Startup done")
+                    pool.submit(
+                        partial(
+                            task.target,
+                            TaskContext(
+                                task=task,
+                                extractor=self,
+                            ),
+                        )
+                    )
+        self._logger.info("Startup done")
 
         for task in continuous:
-            Thread(name=pascalize(task.name), target=task.target).start()
+            Thread(
+                name=pascalize(task.name),
+                target=task.target,
+                args=(TaskContext(task=task, extractor=self),),
+            ).start()
 
         if has_scheduled:
             self._scheduler.run()

cognite/extractorutils/unstable/core/errors.py

@@ -1,11 +1,14 @@
-import typing
+import logging
 from enum import Enum
 from types import TracebackType
+from typing import TYPE_CHECKING
 from uuid import uuid4
 
+from typing_extensions import assert_never
+
 from cognite.extractorutils.util import now
 
-if typing.TYPE_CHECKING:
+if TYPE_CHECKING:
     from .base import Extractor
 
 __all__ = ["Error", "ErrorLevel"]
@@ -16,6 +19,18 @@ class ErrorLevel(Enum):
     error = "error"
     fatal = "fatal"
 
+    @property
+    def log_level(self) -> int:
+        match self:
+            case ErrorLevel.warning:
+                return logging.WARNING
+            case ErrorLevel.error:
+                return logging.ERROR
+            case ErrorLevel.fatal:
+                return logging.CRITICAL
+            case _:
+                assert_never(self)
+
 
 class Error:
     def __init__(
@@ -64,7 +79,7 @@ class Error:
 
     def __exit__(
         self,
-        exc_type: typing.Type[BaseException] | None,
+        exc_type: type[BaseException] | None,
         exc_val: BaseException | None,
         exc_tb: TracebackType | None,
     ) -> bool:

cognite/extractorutils/unstable/core/logger.py

@@ -0,0 +1,149 @@
+from abc import ABC, abstractmethod
+from logging import Logger, getLogger
+from traceback import format_exception
+from typing import Literal
+
+from typing_extensions import assert_never
+
+from cognite.extractorutils.unstable.core.errors import Error, ErrorLevel
+
+
+class CogniteLogger(ABC):
+    def __init__(self) -> None:
+        self._logger: Logger = getLogger()
+
+    @abstractmethod
+    def _new_error(
+        self,
+        level: ErrorLevel,
+        description: str,
+        *,
+        details: str | None = None,
+        task_name: str | None = None,
+    ) -> Error:
+        pass
+
+    def debug(self, message: str) -> None:
+        self._logger.debug(message)
+
+    def info(self, message: str) -> None:
+        self._logger.info(message)
+
+    def begin_warning(
+        self,
+        message: str,
+        *,
+        details: str | None = None,
+        auto_log: bool = True,
+    ) -> Error:
+        if auto_log:
+            self._logger.warning(message)
+        return self._new_error(
+            level=ErrorLevel.warning,
+            description=message,
+            details=details,
+        )
+
+    def begin_error(
+        self,
+        message: str,
+        *,
+        details: str | None = None,
+        auto_log: bool = True,
+    ) -> Error:
+        if auto_log:
+            self._logger.error(message)
+        return self._new_error(
+            level=ErrorLevel.error,
+            description=message,
+            details=details,
+        )
+
+    def begin_fatal(
+        self,
+        message: str,
+        *,
+        details: str | None = None,
+        auto_log: bool = True,
+    ) -> Error:
+        if auto_log:
+            self._logger.critical(message)
+        return self._new_error(
+            level=ErrorLevel.fatal,
+            description=message,
+            details=details,
+        )
+
+    def warning(
+        self,
+        message: str,
+        *,
+        details: str | None = None,
+        auto_log: bool = True,
+    ) -> None:
+        if auto_log:
+            self._logger.warning(message)
+        self._new_error(
+            level=ErrorLevel.warning,
+            description=message,
+            details=details,
+        ).instant()
+
+    def error(
+        self,
+        message: str,
+        *,
+        details: str | None = None,
+        auto_log: bool = True,
+    ) -> None:
+        if auto_log:
+            self._logger.error(message)
+        self._new_error(
+            level=ErrorLevel.error,
+            description=message,
+            details=details,
+        ).instant()
+
+    def fatal(
+        self,
+        message: str,
+        *,
+        details: str | None = None,
+        auto_log: bool = True,
+    ) -> None:
+        if auto_log:
+            self._logger.critical(message)
+        self._new_error(
+            level=ErrorLevel.fatal,
+            description=message,
+            details=details,
+        ).instant()
+
+    def exception(
+        self,
+        message: str,
+        exception: Exception,
+        *,
+        level: ErrorLevel = ErrorLevel.error,
+        include_details: Literal["stack_trace"] | Literal["exception_message"] | bool = "exception_message",
+        auto_log: bool = True,
+    ) -> None:
+        if auto_log:
+            self._logger.log(level=level.log_level, msg=message, exc_info=exception)
+
+        details: str | None
+        match include_details:
+            case "stack_trace":
+                details = "".join(format_exception(exception))
+            case "exception_message" | True:
+                details = str(exception)
+            case False:
+                details = None
+            case _:
+                assert_never(include_details)
+
+        self._new_error(
+            level=level,
+            description=message,
+            details=details,
+        ).instant()

cognite/extractorutils/unstable/core/restart_policy.py

@@ -1,4 +1,4 @@
-from typing import Callable
+from collections.abc import Callable
 
 from cognite.extractorutils.unstable.core.tasks import ContinuousTask, Task
 

cognite/extractorutils/unstable/core/runtime.py

@@ -5,8 +5,7 @@ import time
 from argparse import ArgumentParser, Namespace
 from multiprocessing import Process, Queue
 from pathlib import Path
-from typing import Any, Generic, Type, TypeVar
-from uuid import uuid4
+from typing import Any, Generic, TypeVar
 
 from requests.exceptions import ConnectionError
 from typing_extensions import assert_never
@@ -17,7 +16,6 @@ from cognite.extractorutils.unstable.configuration.exceptions import InvalidConf
 from cognite.extractorutils.unstable.configuration.loaders import load_file, load_from_cdf
 from cognite.extractorutils.unstable.configuration.models import ConnectionConfig
 from cognite.extractorutils.unstable.core._dto import Error
-from cognite.extractorutils.util import now
 
 from ._messaging import RuntimeMessage
 from .base import ConfigRevision, ConfigType, Extractor, FullConfig
@@ -30,7 +28,7 @@ ExtractorType = TypeVar("ExtractorType", bound=Extractor)
 class Runtime(Generic[ExtractorType]):
     def __init__(
         self,
-        extractor: Type[ExtractorType],
+        extractor: type[ExtractorType],
     ) -> None:
         self._extractor_class = extractor
         self._cancellation_token = CancellationToken()
@@ -127,15 +125,13 @@ class Runtime(Generic[ExtractorType]):
         self,
         args: Namespace,
         connection_config: ConnectionConfig,
-    ) -> tuple[ConfigType, ConfigRevision, ConfigRevision]:
+    ) -> tuple[ConfigType, ConfigRevision]:
         current_config_revision: ConfigRevision
-        newest_config_revision: ConfigRevision
 
         if args.local_override:
             self.logger.info("Loading local application config")
 
             current_config_revision = "local"
-            newest_config_revision = "local"
             try:
                 application_config = load_file(args.local_override[0], self._extractor_class.CONFIG_TYPE)
             except InvalidConfigError as e:
@@ -153,50 +149,12 @@ class Runtime(Generic[ExtractorType]):
 
             errors: list[Error] = []
 
-            revision: int | None = None
             try:
-                while True:
-                    try:
-                        application_config, current_config_revision = load_from_cdf(
-                            client,
-                            connection_config.integration,
-                            self._extractor_class.CONFIG_TYPE,
-                            revision=revision,
-                        )
-                        break
-
-                    except InvalidConfigError as e:
-                        if e.attempted_revision is None:
-                            # Should never happen, attempted_revision is set in every handler in load_from_cdf, but it's
-                            # needed for type checks to pass
-                            raise e
-
-                        self.logger.error(f"Revision {e.attempted_revision} is invalid: {e.message}")
-
-                        t = now()
-                        errors.append(
-                            Error(
-                                external_id=str(uuid4()),
-                                level="error",
-                                description=f"Revision {e.attempted_revision} is invalid",
-                                details=e.message,
-                                start_time=t,
-                                end_time=t,
-                                task=None,
-                            )
-                        )
-
-                        if revision is None:
-                            revision = e.attempted_revision - 1
-                            newest_config_revision = e.attempted_revision
-                        else:
-                            revision -= 1
-
-                        if revision > 0:
-                            self.logger.info(f"Falling back to revision {revision}")
-                        else:
-                            self.logger.critical("No more revisions to fall back to")
-                            raise e
+                application_config, current_config_revision = load_from_cdf(
+                    client,
+                    connection_config.integration,
+                    self._extractor_class.CONFIG_TYPE,
+                )
 
             finally:
                 if errors:
@@ -209,7 +167,7 @@ class Runtime(Generic[ExtractorType]):
                         headers={"cdf-version": "alpha"},
                     )
 
-        return application_config, current_config_revision, newest_config_revision
+        return application_config, current_config_revision
 
     def _verify_connection_config(self, connection_config: ConnectionConfig) -> bool:
         client = connection_config.get_cognite_client(
@@ -281,9 +239,7 @@ class Runtime(Generic[ExtractorType]):
         application_config: Any
         while not self._cancellation_token.is_cancelled:
             try:
-                application_config, current_config_revision, newest_config_revision = self._get_application_config(
-                    args, connection_config
-                )
+                application_config, current_config_revision = self._get_application_config(args, connection_config)
 
             except InvalidConfigError:
                 self.logger.critical("Could not get a valid application config file. Shutting down")
@@ -295,7 +251,6 @@ class Runtime(Generic[ExtractorType]):
                     connection_config=connection_config,
                     application_config=application_config,
                     current_config_revision=current_config_revision,
-                    newest_config_revision=newest_config_revision,
                 )
             )
             process.join()