cognite-extractor-utils 7.4.9__py3-none-any.whl → 7.5.0__py3-none-any.whl

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/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/extractorutils/unstable/core/base.py

@@ -1,14 +1,20 @@
 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 typing_extensions import Self
+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]
@@ -40,29 +46,39 @@ class Extractor(Generic[ConfigType]):
         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 _run_checkin(self) -> None:
-        def checkin() -> None:
-            body = {"externalId": self.connection_config.extraction_pipeline}
+    def _checkin(self) -> None:
+        with self._checkin_lock:
+            task_updates = [t.model_dump() for t in self._task_updates]
+            self._task_updates.clear()
 
-            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")
+        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()
+        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:
-                checkin()
+                self.logger.debug("Running checkin")
+                self._checkin()
             except Exception:
                 self.logger.exception("Error during checkin")
             self.cancellation_token.wait(10)
@@ -81,7 +97,32 @@ class Extractor(Generic[ConfigType]):
     ) -> Self:
         return cls(connection_config, application_config, current_config_revision)
 
-    def start(self) -> None:
+    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={
@@ -91,9 +132,19 @@ class Extractor(Generic[ConfigType]):
                     "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:
@@ -110,7 +161,43 @@ class Extractor(Generic[ConfigType]):
         exc_tb: Optional[TracebackType],
     ) -> bool:
         self.stop()
+        with self._checkin_lock:
+            self._checkin()
+
         return exc_val is None
 
     def run(self) -> None:
-        raise NotImplementedError()
+        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/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/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/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.dist-info → cognite_extractor_utils-7.5.0.dist-info}/METADATA

@@ -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.dist-info → cognite_extractor_utils-7.5.0.dist-info}/RECORD

@@ -1,4 +1,4 @@
-cognite/extractorutils/__init__.py,sha256=Cxd38BeR0YorUxKXaS9dz_FKoKa2trhrZJFhiHh0EuY,739
+cognite/extractorutils/__init__.py,sha256=CRWoZhLx2q6fib6_EUce3Das2gwRS9CrDqVlqnOWTPE,739
 cognite/extractorutils/_inner_util.py,sha256=gmz6aqS7jDNsg8z4RHgJjMFohDLOMiaU4gMWBhg3xcE,1558
 cognite/extractorutils/base.py,sha256=q6NU2bPec3WOasVnnIFoh-aUJudVZWZ2R6emz3IRj8Q,16391
 cognite/extractorutils/configtools/__init__.py,sha256=YEpFGJoza23eM8Zj5DqqUj7sEstERV_QYsN6Nw4dKCg,3092
@@ -20,11 +20,13 @@ cognite/extractorutils/unstable/configuration/loaders.py,sha256=9TqVLKGiFl7L-6SL
 cognite/extractorutils/unstable/configuration/models.py,sha256=fPu56TkFKMRr5uc5R3X4WlHNDofvyH0JZfORvwwDT44,7727
 cognite/extractorutils/unstable/core/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 cognite/extractorutils/unstable/core/__main__.py,sha256=8Tb0RqeRBs47dU5xOikZog-IM2IKVonj9dnPG7Ia710,728
+cognite/extractorutils/unstable/core/_dto.py,sha256=OP_cWZ2EvRbq3Tcczc-NWbYsxPZJ0HPgnKyZw-OAZ-I,1117
 cognite/extractorutils/unstable/core/_messaging.py,sha256=D9rOW8fijryXffbm90d8VTf2vy5FmwVGU-H0O-cn-EI,68
-cognite/extractorutils/unstable/core/base.py,sha256=y8KYppvC-9EhdwxZ4XWhYY4VRYGVS1jDs7x_QqVpS_4,3955
+cognite/extractorutils/unstable/core/base.py,sha256=nvvMc8ddi742M-_CIzsHQSlUswMyn4uAX3ceHp3Db_k,6745
 cognite/extractorutils/unstable/core/runtime.py,sha256=sRrPsICawgEHFU9DYArDpIMxebudR0BZ5K3h6HAW9As,6054
+cognite/extractorutils/unstable/core/tasks.py,sha256=JjzfQDpCmNZxWnSJxmhrK3TvxNbmzRxEjgjHJQQO__c,515
 cognite/extractorutils/unstable/scheduling/__init__.py,sha256=L90_rCZNHvti-PInne0r7W9edIkifctELjiaxEoQiSc,67
-cognite/extractorutils/unstable/scheduling/_scheduler.py,sha256=6Qui9v0BcSq-wJONxeqPs6W5SS_5paGW8nvzA3LESVA,3574
+cognite/extractorutils/unstable/scheduling/_scheduler.py,sha256=w2Hs1u3-cNjxrZHtoNFvCmLCd0GNU52K4uUd-Yo_RrM,3691
 cognite/extractorutils/unstable/scheduling/_schedules.py,sha256=y0NVeXYZOFcAyzBgAe8jqK0W-SZL5m99UwXAacGzqIw,677
 cognite/extractorutils/uploader/__init__.py,sha256=W22u6QHA4cR0j78LN5LTL5YGbfC-uTApagTyP5ab7uQ,3110
 cognite/extractorutils/uploader/_base.py,sha256=wktbV8dpb8zBOsNaECZkBNoJSpOz437NlNMER3-a3xQ,5304
@@ -37,8 +39,8 @@ cognite/extractorutils/uploader/raw.py,sha256=wFjF90PFTjmByOWx_Y4_YfDJ2w2jl0EQJ2
 cognite/extractorutils/uploader/time_series.py,sha256=HBtQdsQoIOaL-EG5lMsaY-ORwVb0kGiXG86VjE5-_Bg,26815
 cognite/extractorutils/uploader_extractor.py,sha256=E-mpVvbPg_Tk90U4S9JybV0duptJ2SXE88HB6npE3zI,7732
 cognite/extractorutils/uploader_types.py,sha256=wxfrsiKPTzG5lmoYtQsxt8Xyj-s5HnaLl8WDzJNrazg,1020
-cognite/extractorutils/util.py,sha256=T6ef5b7aYJ8yq9swQwybYaLe3YGr3hElsJQy8E-d5Rs,17469
-cognite_extractor_utils-7.4.9.dist-info/LICENSE,sha256=psuoW8kuDP96RQsdhzwOqi6fyWv0ct8CR6Jr7He_P_k,10173
-cognite_extractor_utils-7.4.9.dist-info/METADATA,sha256=I5rq1WJzAnrchQKwS75G30Xx_qZkJ-ZDYP9w_8RZ2jw,5598
-cognite_extractor_utils-7.4.9.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-cognite_extractor_utils-7.4.9.dist-info/RECORD,,
+cognite/extractorutils/util.py,sha256=cFrbO8XCFT7QBs2vvOk4dWyOvaHclvsuPlrDJYFG0xA,21160
+cognite_extractor_utils-7.5.0.dist-info/LICENSE,sha256=psuoW8kuDP96RQsdhzwOqi6fyWv0ct8CR6Jr7He_P_k,10173
+cognite_extractor_utils-7.5.0.dist-info/METADATA,sha256=kvTwuhJDL9kiV5LDf9S-UzgxBAvV08tAfKaWqnhjlOE,5598
+cognite_extractor_utils-7.5.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+cognite_extractor_utils-7.5.0.dist-info/RECORD,,