cognite-extractor-utils 7.5.8__tar.gz → 7.5.10__tar.gz

{cognite_extractor_utils-7.5.8 → cognite_extractor_utils-7.5.10}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cognite-extractor-utils
-Version: 7.5.8
+Version: 7.5.10
 Summary: Utilities for easier development of extractors for CDF
 Project-URL: repository, https://github.com/cognitedata/python-extractor-utils
 Author-email: Mathias Lohne <mathias.lohne@cognite.com>
@@ -14,7 +14,7 @@ Requires-Dist: azure-identity>=1.14.0
 Requires-Dist: azure-keyvault-secrets>=4.7.0
 Requires-Dist: cognite-sdk>=7.59.0
 Requires-Dist: croniter>=6.0.0
-Requires-Dist: dacite>=1.6.0
+Requires-Dist: dacite<1.9.0,>=1.6.0
 Requires-Dist: decorator>=5.1.1
 Requires-Dist: httpx<1,>=0.27.0
 Requires-Dist: jsonlines>=4.0.0

{cognite_extractor_utils-7.5.8 → cognite_extractor_utils-7.5.10}/cognite/extractorutils/__init__.py

@@ -16,7 +16,7 @@
 Cognite extractor utils is a Python package that simplifies the development of new extractors.
 """
 
-__version__ = "7.5.8"
+__version__ = "7.5.10"
 from .base import Extractor
 
 __all__ = ["Extractor"]

{cognite_extractor_utils-7.5.8 → cognite_extractor_utils-7.5.10}/cognite/extractorutils/configtools/__init__.py

@@ -16,7 +16,7 @@
 Module containing tools for loading and verifying config files, and a YAML loader to automatically serialize these
 dataclasses from a config file.
 
-Configs are described as ``dataclass``\es, and use the ``BaseConfig`` class as a superclass to get a few things
+Configs are described as ``dataclass`` es, and use the ``BaseConfig`` class as a superclass to get a few things
 built-in: config version, Cognite project and logging. Use type hints to specify types, use the ``Optional`` type to
 specify that a config parameter is optional, and give the attribute a value to give it a default.
 

{cognite_extractor_utils-7.5.8 → cognite_extractor_utils-7.5.10}/cognite/extractorutils/unstable/core/base.py

@@ -2,11 +2,10 @@ 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, TypeVar
 
@@ -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
 
@@ -51,7 +51,7 @@ class FullConfig(Generic[_T]):
         self.current_config_revision = current_config_revision
 
 
-class Extractor(Generic[ConfigType]):
+class Extractor(Generic[ConfigType], CogniteLogger):
     NAME: str
     EXTERNAL_ID: str
     DESCRIPTION: str
@@ -62,6 +62,8 @@ class Extractor(Generic[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()
 
@@ -80,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:
@@ -175,36 +173,34 @@ class Extractor(Generic[ConfigType]):
     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()
@@ -217,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
             """
@@ -227,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(
@@ -265,7 +250,16 @@ 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(
@@ -281,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
                 ],
@@ -310,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:
@@ -333,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_extractor_utils-7.5.8 → cognite_extractor_utils-7.5.10}/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__(

cognite_extractor_utils-7.5.10/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_extractor_utils-7.5.10/cognite/extractorutils/unstable/core/tasks.py

@@ -0,0 +1,116 @@
+import logging
+from collections.abc import Callable
+from typing import TYPE_CHECKING
+
+from cognite.extractorutils.unstable.configuration.models import (
+    CronConfig,
+    IntervalConfig,
+    ScheduleConfig,
+    TimeIntervalConfig,
+)
+from cognite.extractorutils.unstable.core.errors import Error, ErrorLevel
+from cognite.extractorutils.unstable.core.logger import CogniteLogger
+
+if TYPE_CHECKING:
+    from cognite.extractorutils.unstable.core.base import Extractor
+
+__all__ = ["ScheduledTask", "ContinuousTask", "StartupTask", "Task", "TaskContext"]
+
+
+class TaskContext(CogniteLogger):
+    def __init__(self, task: "Task", extractor: "Extractor"):
+        super().__init__()
+        self._task = task
+        self._extractor = extractor
+
+        self._logger = logging.getLogger(f"{self._extractor.EXTERNAL_ID}.{self._task.name.replace(' ', '')}")
+
+    def _new_error(
+        self,
+        level: ErrorLevel,
+        description: str,
+        *,
+        details: str | None = None,
+        task_name: str | None = None,
+    ) -> Error:
+        return self._extractor._new_error(
+            level=level,
+            description=description,
+            details=details,
+            task_name=self._task.name,
+        )
+
+
+TaskTarget = Callable[[TaskContext], None]
+
+
+class _Task:
+    def __init__(
+        self,
+        *,
+        name: str,
+        target: TaskTarget,
+        description: str | None = None,
+    ) -> None:
+        self.name = name
+        self.target = target
+        self.description = description
+
+
+class ScheduledTask(_Task):
+    def __init__(
+        self,
+        *,
+        name: str,
+        target: TaskTarget,
+        description: str | None = None,
+        schedule: ScheduleConfig,
+    ):
+        super().__init__(name=name, target=target, description=description)
+        self.schedule = schedule
+
+    @classmethod
+    def from_interval(
+        cls, *, interval: str, name: str, target: TaskTarget, description: str | None = None
+    ) -> "ScheduledTask":
+        return ScheduledTask(
+            name=name,
+            target=target,
+            description=description,
+            schedule=IntervalConfig(type="interval", expression=TimeIntervalConfig(interval)),
+        )
+
+    @classmethod
+    def from_cron(cls, *, cron: str, name: str, target: TaskTarget, description: str | None = None) -> "ScheduledTask":
+        return ScheduledTask(
+            name=name,
+            target=target,
+            description=description,
+            schedule=CronConfig(type="cron", expression=cron),
+        )
+
+
+class ContinuousTask(_Task):
+    def __init__(
+        self,
+        *,
+        name: str,
+        target: TaskTarget,
+        description: str | None = None,
+    ) -> None:
+        super().__init__(name=name, target=target, description=description)
+
+
+class StartupTask(_Task):
+    def __init__(
+        self,
+        *,
+        name: str,
+        target: TaskTarget,
+        description: str | None = None,
+    ) -> None:
+        super().__init__(name=name, target=target, description=description)
+
+
+# Making a type union to help with exhaustion checks in matches
+Task = ScheduledTask | ContinuousTask | StartupTask

{cognite_extractor_utils-7.5.8 → cognite_extractor_utils-7.5.10}/cognite/extractorutils/unstable/scheduling/_scheduler.py

@@ -71,7 +71,7 @@ class TaskScheduler:
                 with self._running_lock:
                     self._running.remove(job)
 
-        Thread(target=wrap, name=f"{pascalize(job.name)}").start()
+        Thread(target=wrap, name=f"{pascalize(job.name.replace(' ', '_'))}").start()
         return True
 
     def trigger(self, name: str) -> bool:

{cognite_extractor_utils-7.5.8 → cognite_extractor_utils-7.5.10}/cognite/extractorutils/uploader/files.py

@@ -25,6 +25,7 @@ from typing import (
 )
 from urllib.parse import ParseResult, urlparse
 
+from h11._util import LocalProtocolError
 from httpx import URL, Client, Headers, Request, StreamConsumed, SyncByteStream
 from requests.utils import super_len
 
@@ -481,6 +482,19 @@ class IOFileUploadQueue(AbstractUploadQueue):
             try:
                 upload_file(read_file, file_meta)
 
+            except LocalProtocolError as e:
+                error_message = str(e)
+                if (
+                    "Too much data for declared Content" in error_message
+                    or "Too little data for declared Content-Length" in error_message
+                ):
+                    self.logger.error(
+                        f"Content Length Mismatch while uploading file: {file_meta} error message: {error_message}"
+                    )
+
+                self.add_entry_failure_logger(file_name=str(file_meta.name), error=e)
+                self.errors.append(e)
+
             except Exception as e:
                 self.logger.exception(
                     f"Unexpected error while uploading file: {file_meta.external_id} {file_meta.name}"

{cognite_extractor_utils-7.5.8 → cognite_extractor_utils-7.5.10}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "cognite-extractor-utils"
-version = "7.5.8"
+version = "7.5.10"
 description = "Utilities for easier development of extractors for CDF"
 authors = [
     {name = "Mathias Lohne", email = "mathias.lohne@cognite.com"}
@@ -19,7 +19,7 @@ dependencies = [
     "prometheus-client>=0.7.0,<=1.0.0",
     "arrow>=1.0.0",
     "pyyaml>=5.3.0,<7",
-    "dacite>=1.6.0",
+    "dacite>=1.6.0,<1.9.0",
     "psutil>=6.0.0",
     "decorator>=5.1.1",
     "more-itertools>=10.0.0",

cognite_extractor_utils-7.5.8/cognite/extractorutils/unstable/core/tasks.py

@@ -1,31 +0,0 @@
-from abc import ABC
-from collections.abc import Callable
-from dataclasses import dataclass
-
-from cognite.extractorutils.unstable.configuration.models import ScheduleConfig
-
-__all__ = ["ScheduledTask", "ContinuousTask", "StartupTask", "Task"]
-
-
-@dataclass
-class _Task(ABC):
-    name: str
-    target: Callable[[], None]
-
-
-@dataclass
-class ScheduledTask(_Task):
-    schedule: ScheduleConfig
-
-
-@dataclass
-class ContinuousTask(_Task):
-    pass
-
-
-class StartupTask(_Task):
-    pass
-
-
-# Making a type union to help with exhaustion checks in matches
-Task = ScheduledTask | ContinuousTask | StartupTask