cognite-extractor-utils 7.4.9__tar.gz → 7.5.0__tar.gz

{cognite_extractor_utils-7.4.9 → cognite_extractor_utils-7.5.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: cognite-extractor-utils
-Version: 7.4.9
+Version: 7.5.0
 Summary: Utilities for easier development of extractors for CDF
 Home-page: https://github.com/cognitedata/python-extractor-utils
 License: Apache-2.0

{cognite_extractor_utils-7.4.9 → cognite_extractor_utils-7.5.0}/cognite/extractorutils/__init__.py

@@ -16,5 +16,5 @@
 Cognite extractor utils is a Python package that simplifies the development of new extractors.
 """
 
-__version__ = "7.4.9"
+__version__ = "7.5.0"
 from .base import Extractor

cognite_extractor_utils-7.5.0/cognite/extractorutils/unstable/core/_dto.py

@@ -0,0 +1,34 @@
+"""
+Temporary holding place for DTOs against Extraction Pipelines 2.0 until it's in the SDK
+"""
+
+from typing import Any, Literal
+
+from humps import camelize
+from pydantic import BaseModel, ConfigDict
+
+
+class CogniteModel(BaseModel):
+    """
+    Base class for DTO classes based on pydantic, but with a few tweaks to make it inline with the CDF API guidelines:
+      * camelCase instead of snake_case when serializing/deserializing into/from JSON
+      * exclude Nones from serialized JSON instead of having nulls in the response text
+    """
+
+    def model_dump(self, *args: Any, **kwargs: Any) -> dict[str, Any]:
+        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]:
+        return self.model_dump(*args, **kwargs)
+
+    model_config = ConfigDict(alias_generator=camelize, populate_by_name=True, extra="forbid")
+
+
+class TaskUpdate(CogniteModel):
+    type: Literal["started"] | Literal["ended"]
+    name: str
+    timestamp: int

cognite_extractor_utils-7.5.0/cognite/extractorutils/unstable/core/base.py

@@ -0,0 +1,203 @@
+import logging
+from concurrent.futures import ThreadPoolExecutor
+from multiprocessing import Queue
+from threading import RLock, Thread
+from types import TracebackType
+from typing import Generic, Literal, Optional, Type, TypeVar, Union
+
+from humps import pascalize
+from typing_extensions import Self, assert_never
+
+from cognite.extractorutils.threading import CancellationToken
+from cognite.extractorutils.unstable.configuration.models import ConnectionConfig, ExtractorConfig
+from cognite.extractorutils.unstable.core._dto import TaskUpdate
+from cognite.extractorutils.unstable.core._messaging import RuntimeMessage
+from cognite.extractorutils.unstable.core.tasks import ContinuousTask, ScheduledTask, StartupTask, Task
+from cognite.extractorutils.unstable.scheduling import TaskScheduler
+from cognite.extractorutils.util import now
+
+ConfigType = TypeVar("ConfigType", bound=ExtractorConfig)
+ConfigRevision = Union[Literal["local"], int]
+
+
+class Extractor(Generic[ConfigType]):
+    NAME: str
+    EXTERNAL_ID: str
+    DESCRIPTION: str
+    VERSION: str
+
+    CONFIG_TYPE: Type[ConfigType]
+
+    def __init__(
+        self,
+        connection_config: ConnectionConfig,
+        application_config: ConfigType,
+        current_config_revision: ConfigRevision,
+    ) -> None:
+        self.cancellation_token = CancellationToken()
+        self.cancellation_token.cancel_on_interrupt()
+
+        self.connection_config = connection_config
+        self.application_config = application_config
+        self.current_config_revision = current_config_revision
+
+        self.cognite_client = self.connection_config.get_cognite_client(f"{self.EXTERNAL_ID}-{self.VERSION}")
+
+        self._checkin_lock = RLock()
+        self._runtime_messages: Optional[Queue[RuntimeMessage]] = None
+
+        self._scheduler = TaskScheduler(self.cancellation_token.create_child_token())
+
+        self._tasks: list[Task] = []
+        self._task_updates: list[TaskUpdate] = []
+
+        self.logger = logging.getLogger(f"{self.EXTERNAL_ID}.main")
+
+    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()
+
+        res = self.cognite_client.post(
+            f"/api/v1/projects/{self.cognite_client.config.project}/odin/checkin",
+            json={
+                "externalId": self.connection_config.extraction_pipeline,
+                "taskEvents": task_updates,
+            },
+            headers={"cdf-version": "alpha"},
+        )
+        new_config_revision = res.json().get("lastConfigRevision")
+
+        if new_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._checkin()
+            except Exception:
+                self.logger.exception("Error during checkin")
+            self.cancellation_token.wait(10)
+
+    def restart(self) -> None:
+        if self._runtime_messages:
+            self._runtime_messages.put(RuntimeMessage.RESTART)
+        self.cancellation_token.cancel()
+
+    @classmethod
+    def init_from_runtime(
+        cls,
+        connection_config: ConnectionConfig,
+        application_config: ConfigType,
+        current_config_revision: ConfigRevision,
+    ) -> Self:
+        return cls(connection_config, application_config, current_config_revision)
+
+    def add_task(self, task: Task) -> None:
+        target = task.target
+
+        def wrapped() -> None:
+            with self._checkin_lock:
+                self._task_updates.append(
+                    TaskUpdate(type="started", name=task.name, timestamp=now()),
+                )
+
+            try:
+                target()
+
+            finally:
+                with self._checkin_lock:
+                    self._task_updates.append(
+                        TaskUpdate(type="ended", name=task.name, timestamp=now()),
+                    )
+
+        task.target = wrapped
+        self._tasks.append(task)
+
+        match task:
+            case ScheduledTask() as t:
+                self._scheduler.schedule_task(name=t.name, schedule=t.schedule, task=t.target)
+
+    def _report_extractor_info(self) -> None:
+        self.cognite_client.post(
+            f"/api/v1/projects/{self.cognite_client.config.project}/odin/extractorinfo",
+            json={
+                "externalId": self.connection_config.extraction_pipeline,
+                "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",
+                    }
+                    for t in self._tasks
+                ],
+            },
+            headers={"cdf-version": "alpha"},
+        )
+
+    def start(self) -> None:
+        self._report_extractor_info()
+        Thread(target=self._run_checkin, name="ExtractorCheckin", daemon=True).start()
+
+    def stop(self) -> None:
+        self.cancellation_token.cancel()
+
+    def __enter__(self) -> Self:
+        self.start()
+        return self
+
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[TracebackType],
+    ) -> bool:
+        self.stop()
+        with self._checkin_lock:
+            self._checkin()
+
+        return exc_val is None
+
+    def run(self) -> None:
+        has_scheduled = False
+
+        startup: list[StartupTask] = []
+        continuous: list[ContinuousTask] = []
+
+        for task in self._tasks:
+            match task:
+                case ScheduledTask():
+                    has_scheduled = True
+
+                case StartupTask() as t:
+                    startup.append(t)
+
+                case ContinuousTask() as t:
+                    continuous.append(t)
+
+                case _:
+                    assert_never(task)
+
+        self.logger.info("Starting up extractor")
+        if startup:
+            with ThreadPoolExecutor() as pool:
+                for task in startup:
+                    pool.submit(task.target)
+        self.logger.info("Startup done")
+
+        for task in continuous:
+            Thread(name=pascalize(task.name), target=task.target).start()
+
+        if has_scheduled:
+            self._scheduler.run()
+
+        else:
+            self.cancellation_token.wait()

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

@@ -0,0 +1,29 @@
+from abc import ABC
+from dataclasses import dataclass
+from typing import Callable
+
+from cognite.extractorutils.unstable.configuration.models import ScheduleConfig
+
+
+@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

{cognite_extractor_utils-7.4.9 → cognite_extractor_utils-7.5.0}/cognite/extractorutils/unstable/scheduling/_scheduler.py

@@ -42,6 +42,8 @@ class TaskScheduler:
                 parsed_schedule = IntervalSchedule(interval=interval_config.expression.seconds)
 
         with self._jobs_lock:
+            if name in self._jobs:
+                raise KeyError(f"Job '{name}' is already added to the scheduler")
             self._jobs[name] = Job(name=name, call=task, schedule=parsed_schedule)
 
     def _get_next(self) -> list[Job]:

{cognite_extractor_utils-7.4.9 → cognite_extractor_utils-7.5.0}/cognite/extractorutils/util.py

@@ -17,10 +17,12 @@ The ``util`` package contains miscellaneous functions and classes that can some
 extractors.
 """
 
+import io
 import logging
 import random
 from datetime import datetime, timezone
 from functools import partial, wraps
+from io import RawIOBase
 from threading import Thread
 from time import time
 from typing import Any, Callable, Dict, Generator, Iterable, List, Optional, Tuple, Type, TypeVar, Union
@@ -510,3 +512,113 @@ def datetime_to_timestamp(dt: datetime) -> int:
 
 def timestamp_to_datetime(ts: int) -> datetime:
     return datetime.fromtimestamp(ts / 1000, tz=timezone.utc)
+
+
+def now() -> int:
+    """
+    Current time in CDF format (milliseonds since 1970-01-01 00:00:00 UTC)
+    """
+    return int(time() * 1000)
+
+
+def truncate_byte_len(item: str, ln: int) -> str:
+    """Safely truncate an arbitrary utf-8 string.
+    Used to sanitize metadata.
+
+    Args:
+        item (str): string to be truncated
+        ln (int): length (bytes)
+
+    Returns:
+        str: truncated string
+    """
+
+    bts = item.encode("utf-8")
+    if len(bts) <= ln:
+        return item
+    bts = bts[:ln]
+    last_codepoint_index = len(bts) - 1
+    # Find the last byte that's the start of an UTF-8 codepoint
+    while last_codepoint_index > 0 and (bts[last_codepoint_index] & 0b11000000) == 0b10000000:
+        last_codepoint_index -= 1
+
+    last_codepoint_start = bts[last_codepoint_index]
+    last_codepoint_len = 0
+    if last_codepoint_start & 0b11111000 == 0b11110000:
+        last_codepoint_len = 4
+    elif last_codepoint_start & 0b11110000 == 0b11100000:
+        last_codepoint_len = 3
+    elif last_codepoint_start & 0b11100000 == 0b11000000:
+        last_codepoint_len = 2
+    elif last_codepoint_start & 0b10000000 == 0:
+        last_codepoint_len = 1
+    else:
+        if last_codepoint_index - 2 <= 0:
+            return ""
+        # Somehow a longer codepoint? In this case just use the previous codepoint.
+        return bts[: (last_codepoint_index - 2)].decode("utf-8")
+
+    last_codepoint_end_index = last_codepoint_index + last_codepoint_len - 1
+    if last_codepoint_end_index > ln - 1:
+        if last_codepoint_index - 2 <= 0:
+            return ""
+        # We're in the middle of a codepoint, cut to the previous one
+        return bts[:last_codepoint_index].decode("utf-8")
+    else:
+        return bts.decode("utf-8")
+
+
+class BufferedReadWithLength(io.BufferedReader):
+    def __init__(
+        self, raw: RawIOBase, buffer_size: int, len: int, on_close: Optional[Callable[[], None]] = None
+    ) -> None:
+        super().__init__(raw, buffer_size)
+        # Do not remove even if it appears to be unused. :P
+        # Requests uses this to add the content-length header, which is necessary for writing to files in azure clusters
+        self.len = len
+        self.on_close = on_close
+
+    def close(self) -> None:
+        if self.on_close:
+            self.on_close()
+        return super().close()
+
+
+def iterable_to_stream(
+    iterator: Iterable[bytes],
+    file_size_bytes: int,
+    buffer_size: int = io.DEFAULT_BUFFER_SIZE,
+    on_close: Optional[Callable[[], None]] = None,
+) -> BufferedReadWithLength:
+    class ChunkIteratorStream(io.RawIOBase):
+        def __init__(self) -> None:
+            self.last_chunk = None
+            self.loaded_bytes = 0
+            self.file_size_bytes = file_size_bytes
+
+        def tell(self) -> int:
+            return self.loaded_bytes
+
+        def __len__(self) -> int:
+            return self.file_size_bytes
+
+        def readable(self) -> bool:
+            return True
+
+        def readinto(self, buffer: Any) -> int | None:
+            try:
+                # Bytes to return
+                ln = len(buffer)
+                chunk = self.last_chunk or next(iterator)  # type: ignore
+                output, self.last_chunk = chunk[:ln], chunk[ln:]
+                if len(self.last_chunk) == 0:  # type: ignore
+                    self.last_chunk = None
+                buffer[: len(output)] = output
+                self.loaded_bytes += len(output)
+                return len(output)
+            except StopIteration:
+                return 0
+
+    return BufferedReadWithLength(
+        ChunkIteratorStream(), buffer_size=buffer_size, len=file_size_bytes, on_close=on_close
+    )

{cognite_extractor_utils-7.4.9 → cognite_extractor_utils-7.5.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "cognite-extractor-utils"
-version = "7.4.9"
+version = "7.5.0"
 description = "Utilities for easier development of extractors for CDF"
 authors = ["Mathias Lohne <mathias.lohne@cognite.com>"]
 license = "Apache-2.0"
@@ -85,8 +85,8 @@ ruff = "^0.6.0"
 pytest = "^8.0.0"
 pytest-cov = "^5.0.0"
 sphinx = "^7.0.0"
-sphinx-rtd-theme = "^2.0.0"
-pre-commit = "^3.3.0"
+sphinx-rtd-theme = "^3.0.0"
+pre-commit = "^4.0.0"
 SecretStorage = "^3.1.2"
 twine = "^5.0.0"
 pytest-order = "^1.0.1"

cognite_extractor_utils-7.4.9/cognite/extractorutils/unstable/core/base.py

@@ -1,116 +0,0 @@
-import logging
-from multiprocessing import Queue
-from threading import RLock, Thread
-from types import TracebackType
-from typing import Generic, Literal, Optional, Type, TypeVar, Union
-
-from typing_extensions import Self
-
-from cognite.extractorutils.threading import CancellationToken
-from cognite.extractorutils.unstable.configuration.models import ConnectionConfig, ExtractorConfig
-from cognite.extractorutils.unstable.core._messaging import RuntimeMessage
-
-ConfigType = TypeVar("ConfigType", bound=ExtractorConfig)
-ConfigRevision = Union[Literal["local"], int]
-
-
-class Extractor(Generic[ConfigType]):
-    NAME: str
-    EXTERNAL_ID: str
-    DESCRIPTION: str
-    VERSION: str
-
-    CONFIG_TYPE: Type[ConfigType]
-
-    def __init__(
-        self,
-        connection_config: ConnectionConfig,
-        application_config: ConfigType,
-        current_config_revision: ConfigRevision,
-    ) -> None:
-        self.cancellation_token = CancellationToken()
-        self.cancellation_token.cancel_on_interrupt()
-
-        self.connection_config = connection_config
-        self.application_config = application_config
-        self.current_config_revision = current_config_revision
-
-        self.cognite_client = self.connection_config.get_cognite_client(f"{self.EXTERNAL_ID}-{self.VERSION}")
-
-        self._checkin_lock = RLock()
-        self._runtime_messages: Optional[Queue[RuntimeMessage]] = None
-
-        self.logger = logging.getLogger(f"{self.EXTERNAL_ID}.main")
-
-    def _set_runtime_message_queue(self, queue: Queue) -> None:
-        self._runtime_messages = queue
-
-    def _run_checkin(self) -> None:
-        def checkin() -> None:
-            body = {"externalId": self.connection_config.extraction_pipeline}
-
-            with self._checkin_lock:
-                res = self.cognite_client.post(
-                    f"/api/v1/projects/{self.cognite_client.config.project}/odin/checkin",
-                    json=body,
-                    headers={"cdf-version": "alpha"},
-                )
-                new_config_revision = res.json().get("lastConfigRevision")
-
-                if new_config_revision and new_config_revision != self.current_config_revision:
-                    self.restart()
-
-        while not self.cancellation_token.is_cancelled:
-            try:
-                checkin()
-            except Exception:
-                self.logger.exception("Error during checkin")
-            self.cancellation_token.wait(10)
-
-    def restart(self) -> None:
-        if self._runtime_messages:
-            self._runtime_messages.put(RuntimeMessage.RESTART)
-        self.cancellation_token.cancel()
-
-    @classmethod
-    def init_from_runtime(
-        cls,
-        connection_config: ConnectionConfig,
-        application_config: ConfigType,
-        current_config_revision: ConfigRevision,
-    ) -> Self:
-        return cls(connection_config, application_config, current_config_revision)
-
-    def start(self) -> None:
-        self.cognite_client.post(
-            f"/api/v1/projects/{self.cognite_client.config.project}/odin/extractorinfo",
-            json={
-                "externalId": self.connection_config.extraction_pipeline,
-                "activeConfigRevision": self.current_config_revision,
-                "extractor": {
-                    "version": self.VERSION,
-                    "externalId": self.EXTERNAL_ID,
-                },
-            },
-            headers={"cdf-version": "alpha"},
-        )
-        Thread(target=self._run_checkin, name="ExtractorCheckin", daemon=True).start()
-
-    def stop(self) -> None:
-        self.cancellation_token.cancel()
-
-    def __enter__(self) -> Self:
-        self.start()
-        return self
-
-    def __exit__(
-        self,
-        exc_type: Optional[Type[BaseException]],
-        exc_val: Optional[BaseException],
-        exc_tb: Optional[TracebackType],
-    ) -> bool:
-        self.stop()
-        return exc_val is None
-
-    def run(self) -> None:
-        raise NotImplementedError()