cognite-extractor-utils 7.5.4__py3-none-any.whl → 7.5.5__py3-none-any.whl

cognite/extractorutils/exceptions.py

@@ -13,9 +13,6 @@
 #  limitations under the License.
 
 
-from typing import List, Optional
-
-
 class InvalidConfigError(Exception):
     """
     Exception thrown from ``load_yaml`` and ``load_yaml_dict`` if config file is invalid. This can be due to
@@ -25,7 +22,7 @@ class InvalidConfigError(Exception):
       * Unkown fields
     """
 
-    def __init__(self, message: str, details: Optional[List[str]] = None):
+    def __init__(self, message: str, details: list[str] | None = None):
         super(InvalidConfigError, self).__init__()
         self.message = message
         self.details = details

cognite/extractorutils/metrics.py

@@ -43,7 +43,7 @@ import threading
 from abc import ABC, abstractmethod
 from time import sleep
 from types import TracebackType
-from typing import Any, Callable, Dict, List, Optional, Tuple, Type, TypeVar, Union
+from typing import Any, Callable, Type, TypeVar
 
 import arrow
 import psutil
@@ -177,14 +177,14 @@ class AbstractMetricsPusher(ABC):
 
     def __init__(
         self,
-        push_interval: Optional[int] = None,
-        thread_name: Optional[str] = None,
-        cancellation_token: Optional[CancellationToken] = None,
+        push_interval: int | None = None,
+        thread_name: str | None = None,
+        cancellation_token: CancellationToken | None = None,
     ):
         self.push_interval = push_interval
         self.thread_name = thread_name
 
-        self.thread: Optional[threading.Thread] = None
+        self.thread: threading.Thread | None = None
         self.thread_name = thread_name
         self.cancellation_token = cancellation_token.create_child_token() if cancellation_token else CancellationToken()
 
@@ -232,7 +232,7 @@ class AbstractMetricsPusher(ABC):
         return self
 
     def __exit__(
-        self, exc_type: Optional[Type[BaseException]], exc_val: Optional[BaseException], exc_tb: Optional[TracebackType]
+        self, exc_type: Type[BaseException] | None, exc_val: BaseException | None, exc_tb: TracebackType | None
     ) -> None:
         """
         Wraps around stop method, for use as context manager
@@ -264,10 +264,10 @@ class PrometheusPusher(AbstractMetricsPusher):
         job_name: str,
         url: str,
         push_interval: int,
-        username: Optional[str] = None,
-        password: Optional[str] = None,
-        thread_name: Optional[str] = None,
-        cancellation_token: Optional[CancellationToken] = None,
+        username: str | None = None,
+        password: str | None = None,
+        thread_name: str | None = None,
+        cancellation_token: CancellationToken | None = None,
     ):
         super(PrometheusPusher, self).__init__(push_interval, thread_name, cancellation_token)
 
@@ -277,7 +277,7 @@ class PrometheusPusher(AbstractMetricsPusher):
 
         self.url = url
 
-    def _auth_handler(self, url: str, method: str, timeout: int, headers: List[Tuple[str, str]], data: Any) -> Callable:
+    def _auth_handler(self, url: str, method: str, timeout: int, headers: list[tuple[str, str]], data: Any) -> Callable:
         """
         Returns a authentication handler against the Prometheus Pushgateway to use in the pushadd_to_gateway method.
 
@@ -340,10 +340,10 @@ class CognitePusher(AbstractMetricsPusher):
         cdf_client: CogniteClient,
         external_id_prefix: str,
         push_interval: int,
-        asset: Optional[Asset] = None,
-        data_set: Optional[EitherId] = None,
-        thread_name: Optional[str] = None,
-        cancellation_token: Optional[CancellationToken] = None,
+        asset: Asset | None = None,
+        data_set: EitherId | None = None,
+        thread_name: str | None = None,
+        cancellation_token: CancellationToken | None = None,
     ):
         super(CognitePusher, self).__init__(push_interval, thread_name, cancellation_token)
 
@@ -360,11 +360,11 @@ class CognitePusher(AbstractMetricsPusher):
         """
         Initialize the CDF tenant with the necessary time series and asset.
         """
-        time_series: List[TimeSeries] = []
+        time_series: list[TimeSeries] = []
 
         if self.asset is not None:
             # Ensure that asset exist, and retrieve internal ID
-            asset: Optional[Asset]
+            asset: Asset | None
             try:
                 asset = self.cdf_client.assets.create(self.asset)
             except CogniteDuplicatedError:
@@ -406,7 +406,7 @@ class CognitePusher(AbstractMetricsPusher):
         """
         timestamp = int(arrow.get().float_timestamp * 1000)
 
-        datapoints: List[Dict[str, Union[str, int, List[Any], Datapoints, DatapointsArray]]] = []
+        datapoints: list[dict[str, str | int | list[Any] | Datapoints | DatapointsArray]] = []
 
         for metric in REGISTRY.collect():
             if type(metric) == Metric and metric.type in ["gauge", "counter"]:

cognite/extractorutils/statestore/_base.py

@@ -1,7 +1,6 @@
 import logging
 import threading
 from abc import ABC, abstractmethod
-from typing import Optional
 
 from cognite.extractorutils._inner_util import _resolve_log_level
 from cognite.extractorutils.threading import CancellationToken
@@ -15,10 +14,10 @@ RETRIES = 10
 class _BaseStateStore(ABC):
     def __init__(
         self,
-        save_interval: Optional[int] = None,
+        save_interval: int | None = None,
         trigger_log_level: str = "DEBUG",
-        thread_name: Optional[str] = None,
-        cancellation_token: Optional[CancellationToken] = None,
+        thread_name: str | None = None,
+        cancellation_token: CancellationToken | None = None,
     ) -> None:
         self._initialized = False
 

cognite/extractorutils/statestore/hashing.py

@@ -2,7 +2,7 @@ import hashlib
 import json
 from abc import ABC
 from types import TracebackType
-from typing import Any, Dict, Iterable, Iterator, Optional, Set, Type
+from typing import Any, Iterable, Iterator, Type
 
 import orjson
 
@@ -19,10 +19,10 @@ from ._base import RETRIES, RETRY_BACKOFF_FACTOR, RETRY_DELAY, RETRY_MAX_DELAY,
 class AbstractHashStateStore(_BaseStateStore, ABC):
     def __init__(
         self,
-        save_interval: Optional[int] = None,
+        save_interval: int | None = None,
         trigger_log_level: str = "DEBUG",
-        thread_name: Optional[str] = None,
-        cancellation_token: Optional[CancellationToken] = None,
+        thread_name: str | None = None,
+        cancellation_token: CancellationToken | None = None,
     ) -> None:
         super().__init__(
             save_interval=save_interval,
@@ -31,31 +31,31 @@ class AbstractHashStateStore(_BaseStateStore, ABC):
             cancellation_token=cancellation_token,
         )
 
-        self._local_state: Dict[str, Dict[str, str]] = {}
-        self._seen: Set[str] = set()
+        self._local_state: dict[str, dict[str, str]] = {}
+        self._seen: set[str] = set()
 
-    def get_state(self, external_id: str) -> Optional[str]:
+    def get_state(self, external_id: str) -> str | None:
         with self.lock:
             return self._local_state.get(external_id, {}).get("digest")
 
-    def _hash_row(self, data: Dict[str, Any]) -> str:
+    def _hash_row(self, data: dict[str, Any]) -> str:
         return hashlib.sha256(orjson.dumps(data, option=orjson.OPT_SORT_KEYS)).hexdigest()
 
-    def set_state(self, external_id: str, data: Dict[str, Any]) -> None:
+    def set_state(self, external_id: str, data: dict[str, Any]) -> None:
         with self.lock:
             self._local_state[external_id] = {"digest": self._hash_row(data)}
 
-    def has_changed(self, external_id: str, data: Dict[str, Any]) -> bool:
+    def has_changed(self, external_id: str, data: dict[str, Any]) -> bool:
         with self.lock:
             if external_id not in self._local_state:
                 return True
 
             return self._hash_row(data) != self._local_state[external_id]["digest"]
 
-    def __getitem__(self, external_id: str) -> Optional[str]:
+    def __getitem__(self, external_id: str) -> str | None:
         return self.get_state(external_id)
 
-    def __setitem__(self, key: str, value: Dict[str, Any]) -> None:
+    def __setitem__(self, key: str, value: dict[str, Any]) -> None:
         self.set_state(external_id=key, data=value)
 
     def __contains__(self, external_id: str) -> bool:
@@ -76,10 +76,10 @@ class RawHashStateStore(AbstractHashStateStore):
         cdf_client: CogniteClient,
         database: str,
         table: str,
-        save_interval: Optional[int] = None,
+        save_interval: int | None = None,
         trigger_log_level: str = "DEBUG",
-        thread_name: Optional[str] = None,
-        cancellation_token: Optional[CancellationToken] = None,
+        thread_name: str | None = None,
+        cancellation_token: CancellationToken | None = None,
     ) -> None:
         super().__init__(
             save_interval=save_interval,
@@ -169,9 +169,9 @@ class RawHashStateStore(AbstractHashStateStore):
 
     def __exit__(
         self,
-        exc_type: Optional[Type[BaseException]],
-        exc_val: Optional[BaseException],
-        exc_tb: Optional[TracebackType],
+        exc_type: Type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
     ) -> None:
         """
         Wraps around stop method, for use as context manager
@@ -188,10 +188,10 @@ class LocalHashStateStore(AbstractHashStateStore):
     def __init__(
         self,
         file_path: str,
-        save_interval: Optional[int] = None,
+        save_interval: int | None = None,
         trigger_log_level: str = "DEBUG",
-        thread_name: Optional[str] = None,
-        cancellation_token: Optional[CancellationToken] = None,
+        thread_name: str | None = None,
+        cancellation_token: CancellationToken | None = None,
     ) -> None:
         super().__init__(
             save_interval=save_interval,
@@ -243,9 +243,9 @@ class LocalHashStateStore(AbstractHashStateStore):
 
     def __exit__(
         self,
-        exc_type: Optional[Type[BaseException]],
-        exc_val: Optional[BaseException],
-        exc_tb: Optional[TracebackType],
+        exc_type: Type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
     ) -> None:
         """
         Wraps around stop method, for use as context manager

cognite/extractorutils/statestore/watermark.py

@@ -88,7 +88,7 @@ You can set a state store to automatically update on upload triggers from an upl
 import json
 from abc import ABC
 from types import TracebackType
-from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple, Type, Union
+from typing import Any, Callable, Dict, Iterator, List, Tuple, Type, Union
 
 from cognite.client import CogniteClient
 from cognite.client.exceptions import CogniteAPIError
@@ -114,10 +114,10 @@ class AbstractStateStore(_BaseStateStore, ABC):
 
     def __init__(
         self,
-        save_interval: Optional[int] = None,
+        save_interval: int | None = None,
         trigger_log_level: str = "DEBUG",
-        thread_name: Optional[str] = None,
-        cancellation_token: Optional[CancellationToken] = None,
+        thread_name: str | None = None,
+        cancellation_token: CancellationToken | None = None,
     ):
         super().__init__(
             save_interval=save_interval,
@@ -152,7 +152,7 @@ class AbstractStateStore(_BaseStateStore, ABC):
                 state = self._local_state.get(external_id, {})
                 return state.get("low"), state.get("high")
 
-    def set_state(self, external_id: str, low: Optional[Any] = None, high: Optional[Any] = None) -> None:
+    def set_state(self, external_id: str, low: Any | None = None, high: Any | None = None) -> None:
         """
         Set/update state of a singe external ID.
 
@@ -166,7 +166,7 @@ class AbstractStateStore(_BaseStateStore, ABC):
             state["low"] = low if low is not None else state.get("low")
             state["high"] = high if high is not None else state.get("high")
 
-    def expand_state(self, external_id: str, low: Optional[Any] = None, high: Optional[Any] = None) -> None:
+    def expand_state(self, external_id: str, low: Any | None = None, high: Any | None = None) -> None:
         """
         Like set_state, but only sets state if the proposed state is outside the stored state. That is if e.g. low is
         lower than the stored low.
@@ -275,10 +275,10 @@ class RawStateStore(AbstractStateStore):
         cdf_client: CogniteClient,
         database: str,
         table: str,
-        save_interval: Optional[int] = None,
+        save_interval: int | None = None,
         trigger_log_level: str = "DEBUG",
-        thread_name: Optional[str] = None,
-        cancellation_token: Optional[CancellationToken] = None,
+        thread_name: str | None = None,
+        cancellation_token: CancellationToken | None = None,
     ):
         super().__init__(save_interval, trigger_log_level, thread_name, cancellation_token)
 
@@ -380,7 +380,7 @@ class RawStateStore(AbstractStateStore):
         return self
 
     def __exit__(
-        self, exc_type: Optional[Type[BaseException]], exc_val: Optional[BaseException], exc_tb: Optional[TracebackType]
+        self, exc_type: Type[BaseException] | None, exc_val: BaseException | None, exc_tb: TracebackType | None
     ) -> None:
         """
         Wraps around stop method, for use as context manager
@@ -409,10 +409,10 @@ class LocalStateStore(AbstractStateStore):
     def __init__(
         self,
         file_path: str,
-        save_interval: Optional[int] = None,
+        save_interval: int | None = None,
         trigger_log_level: str = "DEBUG",
-        thread_name: Optional[str] = None,
-        cancellation_token: Optional[CancellationToken] = None,
+        thread_name: str | None = None,
+        cancellation_token: CancellationToken | None = None,
     ):
         super().__init__(save_interval, trigger_log_level, thread_name, cancellation_token)
 
@@ -459,7 +459,10 @@ class LocalStateStore(AbstractStateStore):
         return self
 
     def __exit__(
-        self, exc_type: Optional[Type[BaseException]], exc_val: Optional[BaseException], exc_tb: Optional[TracebackType]
+        self,
+        exc_type: Type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
     ) -> None:
         """
         Wraps around stop method, for use as context manager

cognite/extractorutils/threading.py

@@ -2,7 +2,7 @@ import logging
 import signal
 from threading import Condition
 from time import time
-from typing import Any, Optional
+from typing import Any
 
 
 class CancellationToken:
@@ -14,10 +14,10 @@ class CancellationToken:
     cancelled if the parent is cancelled, but can be canceled alone without affecting the parent token.
     """
 
-    def __init__(self, condition: Optional[Condition] = None) -> None:
+    def __init__(self, condition: Condition | None = None) -> None:
         self._cv: Condition = condition or Condition()
         self._is_cancelled_int: bool = False
-        self._parent: Optional["CancellationToken"] = None
+        self._parent: "CancellationToken" | None = None
 
     def __repr__(self) -> str:
         cls = self.__class__
@@ -59,7 +59,7 @@ class CancellationToken:
         """
         self.cancel()
 
-    def wait(self, timeout: Optional[float] = None) -> bool:
+    def wait(self, timeout: float | None = None) -> bool:
         endtime = None
         if timeout is not None:
             endtime = time() + timeout

cognite/extractorutils/unstable/configuration/exceptions.py

@@ -0,0 +1,24 @@
+from typing import List, Optional
+
+
+class InvalidConfigError(Exception):
+    """
+    Exception thrown from ``load_yaml`` and ``load_yaml_dict`` if config file is invalid. This can be due to
+
+      * Missing fields
+      * Incompatible types
+      * Unkown fields
+    """
+
+    def __init__(self, message: str, details: Optional[List[str]] = None):
+        super(InvalidConfigError, self).__init__()
+        self.message = message
+        self.details = details
+
+        self.attempted_revision: int | None = None
+
+    def __str__(self) -> str:
+        return f"Invalid config: {self.message}"
+
+    def __repr__(self) -> str:
+        return self.__str__()

cognite/extractorutils/unstable/configuration/loaders.py

@@ -8,9 +8,13 @@ from pydantic import ValidationError
 
 from cognite.client import CogniteClient
 from cognite.extractorutils.configtools.loaders import _load_yaml_dict_raw
-from cognite.extractorutils.exceptions import InvalidConfigError
+from cognite.extractorutils.exceptions import InvalidConfigError as OldInvalidConfigError
+from cognite.extractorutils.unstable.configuration.exceptions import InvalidConfigError
 from cognite.extractorutils.unstable.configuration.models import ConfigModel
 
+__all__ = ["ConfigFormat", "load_file", "load_from_cdf", "load_io", "load_dict"]
+
+
 _T = TypeVar("_T", bound=ConfigModel)
 
 
@@ -44,7 +48,17 @@ def load_from_cdf(
     )
     response.raise_for_status()
     data = response.json()
-    return load_io(StringIO(data["config"]), ConfigFormat.YAML, schema), data["revision"]
+
+    try:
+        return load_io(StringIO(data["config"]), ConfigFormat.YAML, schema), data["revision"]
+
+    except InvalidConfigError as e:
+        e.attempted_revision = data["revision"]
+        raise e
+    except OldInvalidConfigError as e:
+        new_e = InvalidConfigError(e.message)
+        new_e.attempted_revision = data["revision"]
+        raise new_e from e
 
 
 def load_io(stream: TextIO, format: ConfigFormat, schema: Type[_T]) -> _T:
@@ -100,12 +114,9 @@ def load_dict(data: dict, schema: Type[_T]) -> _T:
             if "ctx" in err and "error" in err["ctx"]:
                 exc = err["ctx"]["error"]
                 if isinstance(exc, ValueError) or isinstance(exc, AssertionError):
-                    messages.append(f"{loc_str}: {str(exc)}")
+                    messages.append(f"{str(exc)}: {loc_str}")
                     continue
 
-            if err.get("type") == "json_invalid":
-                messages.append(f"{err.get('msg')}: {loc_str}")
-            else:
-                messages.append(f"{loc_str}: {err.get('msg')}")
+            messages.append(f"{err.get('msg')}: {loc_str}")
 
         raise InvalidConfigError(", ".join(messages), details=messages) from e

cognite/extractorutils/unstable/configuration/models.py

@@ -19,6 +19,21 @@ from cognite.client.credentials import (
 from cognite.extractorutils.configtools._util import _load_certificate_data
 from cognite.extractorutils.exceptions import InvalidConfigError
 
+__all__ = [
+    "ConfigModel",
+    "AuthenticationConfig",
+    "TimeIntervalConfig",
+    "ConnectionConfig",
+    "CronConfig",
+    "IntervalConfig",
+    "ScheduleConfig",
+    "LogLevel",
+    "LogFileHandlerConfig",
+    "LogConsoleHandlerConfig",
+    "LogHandlerConfig",
+    "ExtractorConfig",
+]
+
 
 class ConfigModel(BaseModel):
     model_config = ConfigDict(
@@ -139,7 +154,7 @@ class ConnectionConfig(ConfigModel):
     project: str
     base_url: str
 
-    extraction_pipeline: str
+    integration: str
 
     authentication: AuthenticationConfig
 
@@ -223,17 +238,24 @@ class LogLevel(Enum):
 
 
 class LogFileHandlerConfig(ConfigModel):
+    type: Literal["file"]
     path: Path
     level: LogLevel
     retention: int = 7
 
 
 class LogConsoleHandlerConfig(ConfigModel):
+    type: Literal["console"]
     level: LogLevel
 
 
-LogHandlerConfig = Union[LogFileHandlerConfig, LogConsoleHandlerConfig]
+LogHandlerConfig = Annotated[LogFileHandlerConfig | LogConsoleHandlerConfig, Field(discriminator="type")]
+
+
+# Mypy BS
+def _log_handler_default() -> List[LogHandlerConfig]:
+    return [LogConsoleHandlerConfig(type="console", level=LogLevel.INFO)]
 
 
 class ExtractorConfig(ConfigModel):
-    log_handlers: List[LogHandlerConfig] = Field(default_factory=lambda: [LogConsoleHandlerConfig(level=LogLevel.INFO)])
+    log_handlers: List[LogHandlerConfig] = Field(default_factory=_log_handler_default)

cognite/extractorutils/unstable/core/_dto.py

@@ -32,3 +32,13 @@ class TaskUpdate(CogniteModel):
     type: Literal["started"] | Literal["ended"]
     name: str
     timestamp: int
+
+
+class Error(CogniteModel):
+    external_id: str
+    level: str
+    description: str
+    details: str | None
+    start_time: int
+    end_time: int | None
+    task: str | None