cognite-extractor-utils 7.4.2__tar.gz → 7.4.4__tar.gz

{cognite_extractor_utils-7.4.2 → cognite_extractor_utils-7.4.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: cognite-extractor-utils
-Version: 7.4.2
+Version: 7.4.4
 Summary: Utilities for easier development of extractors for CDF
 Home-page: https://github.com/cognitedata/python-extractor-utils
 License: Apache-2.0
@@ -17,7 +17,7 @@ Provides-Extra: experimental
 Requires-Dist: arrow (>=1.0.0,<2.0.0)
 Requires-Dist: azure-identity (>=1.14.0,<2.0.0)
 Requires-Dist: azure-keyvault-secrets (>=4.7.0,<5.0.0)
-Requires-Dist: cognite-sdk (>=7.58.4,<8.0.0)
+Requires-Dist: cognite-sdk (>=7.59.0,<8.0.0)
 Requires-Dist: dacite (>=1.6.0,<2.0.0)
 Requires-Dist: decorator (>=5.1.1,<6.0.0)
 Requires-Dist: httpx (>=0.27.0,<0.28.0)
@@ -25,6 +25,8 @@ Requires-Dist: more-itertools (>=10.0.0,<11.0.0)
 Requires-Dist: orjson (>=3.10.3,<4.0.0)
 Requires-Dist: prometheus-client (>0.7.0,<=1.0.0)
 Requires-Dist: psutil (>=6.0.0,<7.0.0)
+Requires-Dist: pydantic (>=2.8.2,<3.0.0)
+Requires-Dist: pyhumps (>=3.8.0,<4.0.0)
 Requires-Dist: python-dotenv (>=1.0.0,<2.0.0)
 Requires-Dist: pyyaml (>=5.3.0,<7)
 Requires-Dist: typing-extensions (>=3.7.4,<5)

{cognite_extractor_utils-7.4.2 → cognite_extractor_utils-7.4.4}/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.2"
+__version__ = "7.4.4"
 from .base import Extractor

{cognite_extractor_utils-7.4.2 → cognite_extractor_utils-7.4.4}/cognite/extractorutils/exceptions.py

@@ -13,6 +13,9 @@
 #  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
@@ -22,9 +25,10 @@ class InvalidConfigError(Exception):
       * Unkown fields
     """
 
-    def __init__(self, message: str):
+    def __init__(self, message: str, details: Optional[List[str]] = None):
         super(InvalidConfigError, self).__init__()
         self.message = message
+        self.details = details
 
     def __str__(self) -> str:
         return f"Invalid config: {self.message}"

cognite_extractor_utils-7.4.4/cognite/extractorutils/unstable/configuration/loaders.py

@@ -0,0 +1,111 @@
+import json
+from enum import Enum
+from io import StringIO
+from pathlib import Path
+from typing import Dict, Optional, TextIO, Type, TypeVar, Union
+
+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.unstable.configuration.models import ConfigModel
+
+_T = TypeVar("_T", bound=ConfigModel)
+
+
+class ConfigFormat(Enum):
+    JSON = "json"
+    YAML = "yaml"
+
+
+def load_file(path: Path, schema: Type[_T]) -> _T:
+    if path.suffix in [".yaml", ".yml"]:
+        format = ConfigFormat.YAML
+    elif path.suffix == ".json":
+        format = ConfigFormat.JSON
+    else:
+        raise InvalidConfigError(f"Unknown file type {path.suffix}")
+
+    with open(path, "r") as stream:
+        return load_io(stream, format, schema)
+
+
+def load_from_cdf(
+    cognite_client: CogniteClient, external_id: str, schema: Type[_T], revision: Optional[int] = None
+) -> _T:
+    params: Dict[str, Union[str, int]] = {"externalId": external_id}
+    if revision:
+        params["revision"] = revision
+    response = cognite_client.get(
+        f"/api/v1/projects/{cognite_client.config.project}/odin/config",
+        params=params,
+        headers={"cdf-version": "alpha"},
+    )
+    response.raise_for_status()
+    data = response.json()
+    return load_io(StringIO(data["config"]), ConfigFormat.YAML, schema)
+
+
+def load_io(stream: TextIO, format: ConfigFormat, schema: Type[_T]) -> _T:
+    if format == ConfigFormat.JSON:
+        data = json.load(stream)
+
+    elif format == ConfigFormat.YAML:
+        data = _load_yaml_dict_raw(stream)
+
+        if "azure-keyvault" in data:
+            data.pop("azure-keyvault")
+        if "key-vault" in data:
+            data.pop("key-vault")
+
+    return load_dict(data, schema)
+
+
+def _make_loc_str(loc: tuple) -> str:
+    # Remove the body parameter if it is present
+    if loc[0] == "body":
+        loc = loc[1:]
+
+    # Create a string from the loc parameter
+    loc_str = ""
+    needs_sep = False
+    for lo in loc:
+        if not needs_sep:
+            loc_str = f"{loc_str}{lo}"
+            needs_sep = True
+        else:
+            if isinstance(lo, int):
+                loc_str = f"{loc_str}[{lo}]"
+            else:
+                loc_str = f"{loc_str}.{lo}"
+
+    return loc_str
+
+
+def load_dict(data: dict, schema: Type[_T]) -> _T:
+    try:
+        return schema.model_validate(data)
+
+    except ValidationError as e:
+        messages = []
+        for err in e.errors():
+            loc = err.get("loc")
+            if loc is None:
+                continue
+
+            # Create a string from the loc parameter
+            loc_str = _make_loc_str(loc)
+
+            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)}")
+                    continue
+
+            if err.get("type") == "json_invalid":
+                messages.append(f"{err.get('msg')}: {loc_str}")
+            else:
+                messages.append(f"{loc_str}: {err.get('msg')}")
+
+        raise InvalidConfigError(", ".join(messages), details=messages) from e

cognite_extractor_utils-7.4.4/cognite/extractorutils/unstable/configuration/models.py

@@ -0,0 +1,159 @@
+import re
+from datetime import timedelta
+from enum import Enum
+from pathlib import Path
+from typing import Annotated, Any, Dict, List, Literal, Optional, Union
+
+from humps import kebabize
+from pydantic import BaseModel, ConfigDict, Field, GetCoreSchemaHandler
+from pydantic_core import CoreSchema, core_schema
+
+from cognite.extractorutils.exceptions import InvalidConfigError
+
+
+class ConfigModel(BaseModel):
+    model_config = ConfigDict(
+        alias_generator=kebabize,
+        populate_by_name=True,
+        extra="forbid",
+        # arbitrary_types_allowed=True,
+    )
+
+
+class _ClientCredentialsConfig(ConfigModel):
+    type: Literal["client-credentials"]
+    client_id: str
+    client_secret: str
+    token_url: str
+    scopes: List[str]
+    resource: Optional[str] = None
+    audience: Optional[str] = None
+
+
+class _ClientCertificateConfig(ConfigModel):
+    type: Literal["client-certificate"]
+    client_id: str
+    certificate_path: Path
+    scopes: List[str]
+
+
+AuthenticationConfig = Annotated[Union[_ClientCredentialsConfig, _ClientCertificateConfig], Field(discriminator="type")]
+
+
+class TimeIntervalConfig:
+    """
+    Configuration parameter for setting a time interval
+    """
+
+    def __init__(self, expression: str) -> None:
+        self._interval, self._expression = TimeIntervalConfig._parse_expression(expression)
+
+    @classmethod
+    def __get_pydantic_core_schema__(cls, source_type: Any, handler: GetCoreSchemaHandler) -> CoreSchema:
+        return core_schema.no_info_after_validator_function(cls, handler(Union[str, int]))
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, TimeIntervalConfig):
+            return NotImplemented
+        return self._interval == other._interval
+
+    def __hash__(self) -> int:
+        return hash(self._interval)
+
+    @classmethod
+    def _parse_expression(cls, expression: str) -> tuple[int, str]:
+        # First, try to parse pure number and assume seconds (for backwards compatibility)
+        try:
+            return int(expression), f"{expression}s"
+        except ValueError:
+            pass
+
+        match = re.match(r"(\d+)[ \t]*(s|m|h|d)", expression)
+        if not match:
+            raise InvalidConfigError("Invalid interval pattern")
+
+        number, unit = match.groups()
+        numeric_unit = {"s": 1, "m": 60, "h": 60 * 60, "d": 60 * 60 * 24}[unit]
+
+        return int(number) * numeric_unit, expression
+
+    @property
+    def seconds(self) -> int:
+        return self._interval
+
+    @property
+    def minutes(self) -> float:
+        return self._interval / 60
+
+    @property
+    def hours(self) -> float:
+        return self._interval / (60 * 60)
+
+    @property
+    def days(self) -> float:
+        return self._interval / (60 * 60 * 24)
+
+    @property
+    def timedelta(self) -> timedelta:
+        days = self._interval // (60 * 60 * 24)
+        seconds = self._interval % (60 * 60 * 24)
+        return timedelta(days=days, seconds=seconds)
+
+    def __int__(self) -> int:
+        return int(self._interval)
+
+    def __float__(self) -> float:
+        return float(self._interval)
+
+    def __str__(self) -> str:
+        return self._expression
+
+    def __repr__(self) -> str:
+        return self._expression
+
+
+class _ConnectionParameters(ConfigModel):
+    gzip_compression: bool = False
+    status_forcelist: List[int] = Field(default_factory=lambda: [429, 502, 503, 504])
+    max_retries: int = 10
+    max_retries_connect: int = 3
+    max_retry_backoff: TimeIntervalConfig = Field(default_factory=lambda: TimeIntervalConfig("30s"))
+    max_connection_pool_size: int = 50
+    ssl_verify: bool = True
+    proxies: Dict[str, str] = Field(default_factory=dict)
+
+
+class ConnectionConfig(ConfigModel):
+    project: str
+    base_url: str
+
+    extraction_pipeline: str
+
+    authentication: AuthenticationConfig
+
+    connection: _ConnectionParameters = Field(default_factory=_ConnectionParameters)
+
+
+class LogLevel(Enum):
+    CRITICAL = "CRITICAL"
+    ERROR = "ERROR"
+    WARNING = "WARNING"
+    INFO = "INFO"
+    DEBUG = "DEBUG"
+
+
+class LogFileHandlerConfig(ConfigModel):
+    path: Path
+    level: LogLevel
+    retention: int = 7
+
+
+class LogConsoleHandlerConfig(ConfigModel):
+    level: LogLevel
+
+
+LogHandlerConfig = Union[LogFileHandlerConfig, LogConsoleHandlerConfig]
+
+
+class ExtractorConfig(ConfigModel):
+    log_handlers: List[LogHandlerConfig] = Field(default_factory=lambda: [LogConsoleHandlerConfig(level=LogLevel.INFO)])

{cognite_extractor_utils-7.4.2 → cognite_extractor_utils-7.4.4}/cognite/extractorutils/uploader/files.py

@@ -251,43 +251,38 @@ class IOFileUploadQueue(AbstractUploadQueue):
         return node.as_id()
 
     def _upload_empty(
-        self, meta_or_apply: FileMetadataOrCogniteExtractorFile
+        self, file_meta: FileMetadataOrCogniteExtractorFile
     ) -> tuple[FileMetadataOrCogniteExtractorFile, str]:
-        if isinstance(meta_or_apply, CogniteExtractorFileApply):
-            node_id = self._apply_cognite_file(meta_or_apply)
-            meta_or_apply, url = self._create_cdm(instance_id=node_id)
+        if isinstance(file_meta, CogniteExtractorFileApply):
+            node_id = self._apply_cognite_file(file_meta)
+            file_meta, url = self._create_cdm(instance_id=node_id)
         else:
-            meta_or_apply, url = self.cdf_client.files.create(
-                file_metadata=meta_or_apply, overwrite=self.overwrite_existing
-            )
-        return meta_or_apply, url
+            file_meta, url = self.cdf_client.files.create(file_metadata=file_meta, overwrite=self.overwrite_existing)
+        return file_meta, url
 
-    def _upload_bytes(self, size: int, file: BinaryIO, meta_or_apply: FileMetadataOrCogniteExtractorFile) -> None:
-        meta_or_apply, url = self._upload_empty(meta_or_apply)
-        resp = self._httpx_client.send(self._get_file_upload_request(url, file, size, meta_or_apply.mime_type))
+    def _upload_bytes(self, size: int, file: BinaryIO, file_meta: FileMetadataOrCogniteExtractorFile) -> None:
+        file_meta, url = self._upload_empty(file_meta)
+        resp = self._httpx_client.send(self._get_file_upload_request(url, file, size, file_meta.mime_type))
         resp.raise_for_status()
 
-    def _upload_multipart(self, size: int, file: BinaryIO, meta_or_apply: FileMetadataOrCogniteExtractorFile) -> None:
+    def _upload_multipart(self, size: int, file: BinaryIO, file_meta: FileMetadataOrCogniteExtractorFile) -> None:
         chunks = ChunkedStream(file, self.max_file_chunk_size, size)
         self.logger.debug(
-            f"File {meta_or_apply.external_id} is larger than 5GiB ({size})"
-            f", uploading in {chunks.chunk_count} chunks"
+            f"File {file_meta.external_id} is larger than 5GiB ({size})" f", uploading in {chunks.chunk_count} chunks"
         )
 
-        returned_file_metadata = self._create_multi_part(meta_or_apply, chunks)
+        returned_file_metadata = self._create_multi_part(file_meta, chunks)
         upload_urls = returned_file_metadata["uploadUrls"]
         upload_id = returned_file_metadata["uploadId"]
         file_meta = FileMetadata.load(returned_file_metadata)
 
         for url in upload_urls:
             chunks.next_chunk()
-            resp = self._httpx_client.send(
-                self._get_file_upload_request(url, chunks, len(chunks), meta_or_apply.mime_type)
-            )
+            resp = self._httpx_client.send(self._get_file_upload_request(url, chunks, len(chunks), file_meta.mime_type))
             resp.raise_for_status()
 
         completed_headers = (
-            _CDF_ALPHA_VERSION_HEADER if isinstance(meta_or_apply, CogniteExtractorFileApply) is not None else None
+            _CDF_ALPHA_VERSION_HEADER if isinstance(file_meta, CogniteExtractorFileApply) is not None else None
         )
 
         res = self.cdf_client.files._post(
@@ -297,9 +292,9 @@ class IOFileUploadQueue(AbstractUploadQueue):
         )
         res.raise_for_status()
 
-    def _create_multi_part(self, meta_or_apply: FileMetadataOrCogniteExtractorFile, chunks: ChunkedStream) -> dict:
-        if isinstance(meta_or_apply, CogniteExtractorFileApply):
-            node_id = self._apply_cognite_file(meta_or_apply)
+    def _create_multi_part(self, file_meta: FileMetadataOrCogniteExtractorFile, chunks: ChunkedStream) -> dict:
+        if isinstance(file_meta, CogniteExtractorFileApply):
+            node_id = self._apply_cognite_file(file_meta)
             identifiers = IdentifierSequence.load(instance_ids=node_id).as_singleton()
             self.cdf_client.files._warn_alpha()
             res = self.cdf_client.files._post(
@@ -313,7 +308,7 @@ class IOFileUploadQueue(AbstractUploadQueue):
         else:
             res = self.cdf_client.files._post(
                 url_path="/files/initmultipartupload",
-                json=meta_or_apply.dump(camel_case=True),
+                json=file_meta.dump(camel_case=True),
                 params={"overwrite": self.overwrite_existing, "parts": chunks.chunk_count},
             )
             res.raise_for_status()
@@ -321,7 +316,7 @@ class IOFileUploadQueue(AbstractUploadQueue):
 
     def add_io_to_upload_queue(
         self,
-        meta_or_apply: FileMetadataOrCogniteExtractorFile,
+        file_meta: FileMetadataOrCogniteExtractorFile,
         read_file: Callable[[], BinaryIO],
         extra_retries: Optional[
             Union[Tuple[Type[Exception], ...], Dict[Type[Exception], Callable[[Any], bool]]]
@@ -351,36 +346,34 @@ class IOFileUploadQueue(AbstractUploadQueue):
             max_delay=RETRY_MAX_DELAY,
             backoff=RETRY_BACKOFF_FACTOR,
         )
-        def upload_file(read_file: Callable[[], BinaryIO], meta_or_apply: FileMetadataOrCogniteExtractorFile) -> None:
+        def upload_file(read_file: Callable[[], BinaryIO], file_meta: FileMetadataOrCogniteExtractorFile) -> None:
             with read_file() as file:
                 size = super_len(file)
                 if size == 0:
                     # upload just the file metadata witout data
-                    meta_or_apply, _ = self._upload_empty(meta_or_apply)
+                    file_meta, _ = self._upload_empty(file_meta)
                 elif size >= self.max_single_chunk_file_size:
                     # The minimum chunk size is 4000MiB.
-                    self._upload_multipart(size, file, meta_or_apply)
+                    self._upload_multipart(size, file, file_meta)
 
                 else:
-                    self._upload_bytes(size, file, meta_or_apply)
+                    self._upload_bytes(size, file, file_meta)
 
-                if isinstance(meta_or_apply, CogniteExtractorFileApply):
-                    meta_or_apply.is_uploaded = True
+                if isinstance(file_meta, CogniteExtractorFileApply):
+                    file_meta.is_uploaded = True
 
             if self.post_upload_function:
                 try:
-                    self.post_upload_function([meta_or_apply])
+                    self.post_upload_function([file_meta])
                 except Exception as e:
                     self.logger.error("Error in upload callback: %s", str(e))
 
-        def wrapped_upload(
-            read_file: Callable[[], BinaryIO], meta_or_apply: FileMetadataOrCogniteExtractorFile
-        ) -> None:
+        def wrapped_upload(read_file: Callable[[], BinaryIO], file_meta: FileMetadataOrCogniteExtractorFile) -> None:
             try:
-                upload_file(read_file, meta_or_apply)
+                upload_file(read_file, file_meta)
 
             except Exception as e:
-                self.logger.exception(f"Unexpected error while uploading file: {meta_or_apply.external_id}")
+                self.logger.exception(f"Unexpected error while uploading file: {file_meta.external_id}")
                 self.errors.append(e)
 
             finally:
@@ -397,7 +390,7 @@ class IOFileUploadQueue(AbstractUploadQueue):
                     pass
 
         with self.lock:
-            self.upload_queue.append(self._pool.submit(wrapped_upload, read_file, meta_or_apply))
+            self.upload_queue.append(self._pool.submit(wrapped_upload, read_file, file_meta))
             self.upload_queue_size += 1
             self.files_queued.inc()
             self.queue_size.set(self.upload_queue_size)
@@ -522,7 +515,7 @@ class FileUploadQueue(IOFileUploadQueue):
         )
 
     def add_to_upload_queue(
-        self, meta_or_apply: FileMetadataOrCogniteExtractorFile, file_name: Union[str, PathLike]
+        self, file_meta: FileMetadataOrCogniteExtractorFile, file_name: Union[str, PathLike]
     ) -> None:
         """
         Add file to upload queue. The queue will be uploaded if the queue size is larger than the threshold
@@ -537,7 +530,7 @@ class FileUploadQueue(IOFileUploadQueue):
         def load_file_from_path() -> BinaryIO:
             return open(file_name, "rb")
 
-        self.add_io_to_upload_queue(meta_or_apply, load_file_from_path)
+        self.add_io_to_upload_queue(file_meta, load_file_from_path)
 
 
 class BytesUploadQueue(IOFileUploadQueue):
@@ -574,7 +567,7 @@ class BytesUploadQueue(IOFileUploadQueue):
             cancellation_token,
         )
 
-    def add_to_upload_queue(self, content: bytes, meta_or_apply: FileMetadataOrCogniteExtractorFile) -> None:
+    def add_to_upload_queue(self, content: bytes, file_meta: FileMetadataOrCogniteExtractorFile) -> None:
         """
         Add object to upload queue. The queue will be uploaded if the queue size is larger than the threshold
         specified in the __init__.
@@ -586,4 +579,4 @@ class BytesUploadQueue(IOFileUploadQueue):
         def get_byte_io() -> BinaryIO:
             return BytesIO(content)
 
-        self.add_io_to_upload_queue(meta_or_apply, get_byte_io)
+        self.add_io_to_upload_queue(file_meta, get_byte_io)

{cognite_extractor_utils-7.4.2 → cognite_extractor_utils-7.4.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "cognite-extractor-utils"
-version = "7.4.2"
+version = "7.4.4"
 description = "Utilities for easier development of extractors for CDF"
 authors = ["Mathias Lohne <mathias.lohne@cognite.com>"]
 license = "Apache-2.0"
@@ -58,7 +58,7 @@ exclude = "tests/*"
 
 [tool.poetry.dependencies]
 python = "^3.9.0"
-cognite-sdk = "^7.58.4"
+cognite-sdk = "^7.59.0"
 prometheus-client = ">0.7.0, <=1.0.0"
 arrow = "^1.0.0"
 pyyaml = ">=5.3.0, <7"
@@ -72,6 +72,8 @@ azure-identity = "^1.14.0"
 azure-keyvault-secrets = "^4.7.0"
 orjson = "^3.10.3"
 httpx = "^0.27.0"
+pydantic = "^2.8.2"
+pyhumps = "^3.8.0"
 
 [tool.poetry.extras]
 experimental = ["cognite-sdk-experimental"]