cognite-extractor-utils 7.4.6__tar.gz → 7.4.8__tar.gz

{cognite_extractor_utils-7.4.6 → cognite_extractor_utils-7.4.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: cognite-extractor-utils
-Version: 7.4.6
+Version: 7.4.8
 Summary: Utilities for easier development of extractors for CDF
 Home-page: https://github.com/cognitedata/python-extractor-utils
 License: Apache-2.0
@@ -18,6 +18,7 @@ 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.59.0,<8.0.0)
+Requires-Dist: croniter (>=3.0.3,<4.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)

{cognite_extractor_utils-7.4.6 → cognite_extractor_utils-7.4.8}/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.6"
+__version__ = "7.4.8"
 from .base import Extractor

{cognite_extractor_utils-7.4.6 → cognite_extractor_utils-7.4.8}/cognite/extractorutils/configtools/_util.py

@@ -81,8 +81,10 @@ def _to_snake_case(dictionary: Dict[str, Any], case_style: str) -> Dict[str, Any
         raise ValueError(f"Invalid case style: {case_style}")
 
 
-def _load_certificate_data(cert_path: str, password: Optional[str]) -> Union[Tuple[str, str], Tuple[bytes, bytes]]:
-    path = Path(cert_path)
+def _load_certificate_data(
+    cert_path: str | Path, password: Optional[str]
+) -> Union[Tuple[str, str], Tuple[bytes, bytes]]:
+    path = Path(cert_path) if isinstance(cert_path, str) else cert_path
     cert_data = Path(path).read_bytes()
 
     if path.suffix == ".pem":

{cognite_extractor_utils-7.4.6 → cognite_extractor_utils-7.4.8}/cognite/extractorutils/configtools/elements.py

@@ -696,3 +696,51 @@ class StateStoreConfig:
             return LocalStateStore(file_path="states.json", cancellation_token=cancellation_token)
         else:
             return NoStateStore()
+
+
+class RegExpFlag(Enum):
+    IGNORECASE = "ignore-case"
+    IC = "i"
+    ASCII = "ascii-only"
+    A = "a"
+
+    def get_regex_flag(self) -> int:
+        if self in (RegExpFlag.IGNORECASE, RegExpFlag.IC):
+            return re.IGNORECASE
+        elif self.value in (RegExpFlag.ASCII, RegExpFlag.A):
+            return re.ASCII
+        return 0
+
+
+@dataclass
+class IgnorePattern:
+    """
+    Configuration for regexp for ignore pattern
+    """
+
+    pattern: str
+    options: Optional[list[RegExpFlag]] = None
+    flags: Optional[list[RegExpFlag]] = None
+
+    def compile(self) -> re.Pattern[str]:
+        """
+        Compile RegExp pattern.
+
+        Returns:
+            Compiled pattern.
+        """
+        flag = 0
+        for f in self.options or []:
+            flag |= f.get_regex_flag()
+        return re.compile(self.pattern, flag)
+
+    def __post_init__(self) -> None:
+        if self.options is not None and self.flags is not None:
+            raise ValueError("Only one of either 'options' or 'flags' can be specified.")
+        if self.options is None and self.flags is None:
+            raise ValueError("'options' is required.")
+
+        if self.flags is not None:
+            _logger.warning("'options' is preferred over 'flags' as this may be removed in a future release")
+            self.options = self.flags
+            self.flags = None

{cognite_extractor_utils-7.4.6 → cognite_extractor_utils-7.4.8}/cognite/extractorutils/configtools/loaders.py

@@ -22,7 +22,7 @@ import sys
 from enum import Enum
 from hashlib import sha256
 from pathlib import Path
-from typing import Any, Callable, Dict, Generic, Iterable, Optional, TextIO, Type, TypeVar, Union, cast
+from typing import Any, Callable, Dict, Generic, Iterable, List, Optional, TextIO, Type, TypeVar, Union, cast
 
 import dacite
 import yaml
@@ -37,6 +37,7 @@ from cognite.extractorutils.configtools._util import _to_snake_case
 from cognite.extractorutils.configtools.elements import (
     BaseConfig,
     ConfigType,
+    IgnorePattern,
     TimeIntervalConfig,
     _BaseConfig,
 )
@@ -320,6 +321,25 @@ def load_yaml_dict(
     )
 
 
+def compile_patterns(ignore_patterns: List[Union[str, IgnorePattern]]) -> list[re.Pattern[str]]:
+    """
+    List of patterns to compile
+
+    Args:
+        ignore_patterns: A list of strings or IgnorePattern to be compiled.
+
+    Returns:
+        A list of compiled RegExp patterns.
+    """
+    compiled = []
+    for p in ignore_patterns:
+        if isinstance(p, IgnorePattern):
+            compiled.append(re.compile(p.compile()))
+        else:
+            compiled.append(re.compile(p))
+    return compiled
+
+
 class ConfigResolver(Generic[CustomConfigClass]):
     def __init__(self, config_path: str, config_type: Type[CustomConfigClass]):
         self.config_path = config_path

cognite_extractor_utils-7.4.8/cognite/extractorutils/configtools/validators.py

@@ -0,0 +1,37 @@
+import logging
+import re
+from typing import Union
+
+_logger = logging.getLogger(__name__)
+
+
+def matches_patterns(patterns: list[Union[str, re.Pattern[str]]], string: str) -> bool:
+    """
+    Check string against list of RegExp patterns.
+
+    Args:
+        patterns: A list of (re) patterns to match string against.
+        string: String to which we match the pattern.
+
+    Returns:
+        boolean value indicating whether string matches any of the patterns.
+    """
+    return any([matches_pattern(pattern, string) for pattern in patterns])
+
+
+def matches_pattern(pattern: Union[str, re.Pattern[str]], string: str) -> bool:
+    """
+    Match pattern against a string.
+
+    Args:
+        pattern: (re) Pattern to match against a string.
+        string: String to which we match the pattern.
+
+    Returns:
+        boolean value indicating a match or otherwise.
+    """
+    try:
+        return re.search(pattern, string) is not None
+    except re.error as e:
+        _logger.warning(f"Could not apply RegExp: {pattern}\nReason: {e}")
+        return False

{cognite_extractor_utils-7.4.6 → cognite_extractor_utils-7.4.8}/cognite/extractorutils/unstable/configuration/loaders.py

@@ -2,7 +2,7 @@ import json
 from enum import Enum
 from io import StringIO
 from pathlib import Path
-from typing import Dict, Optional, TextIO, Type, TypeVar, Union
+from typing import Dict, Optional, TextIO, Tuple, Type, TypeVar, Union
 
 from pydantic import ValidationError
 
@@ -33,7 +33,7 @@ def load_file(path: Path, schema: Type[_T]) -> _T:
 
 def load_from_cdf(
     cognite_client: CogniteClient, external_id: str, schema: Type[_T], revision: Optional[int] = None
-) -> _T:
+) -> Tuple[_T, int]:
     params: Dict[str, Union[str, int]] = {"externalId": external_id}
     if revision:
         params["revision"] = revision
@@ -44,7 +44,7 @@ def load_from_cdf(
     )
     response.raise_for_status()
     data = response.json()
-    return load_io(StringIO(data["config"]), ConfigFormat.YAML, schema)
+    return load_io(StringIO(data["config"]), ConfigFormat.YAML, schema), data["revision"]
 
 
 def load_io(stream: TextIO, format: ConfigFormat, schema: Type[_T]) -> _T:

{cognite_extractor_utils-7.4.6 → cognite_extractor_utils-7.4.8}/cognite/extractorutils/unstable/configuration/models.py

@@ -7,7 +7,16 @@ 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 typing_extensions import assert_never
+
+from cognite.client import CogniteClient
+from cognite.client.config import ClientConfig
+from cognite.client.credentials import (
+    CredentialProvider,
+    OAuthClientCertificate,
+    OAuthClientCredentials,
+)
+from cognite.extractorutils.configtools._util import _load_certificate_data
 from cognite.extractorutils.exceptions import InvalidConfigError
 
 
@@ -33,7 +42,9 @@ class _ClientCredentialsConfig(ConfigModel):
 class _ClientCertificateConfig(ConfigModel):
     type: Literal["client-certificate"]
     client_id: str
-    certificate_path: Path
+    path: Path
+    password: Optional[str] = None
+    authority_url: str
     scopes: List[str]
 
 
@@ -121,6 +132,7 @@ class _ConnectionParameters(ConfigModel):
     max_connection_pool_size: int = 50
     ssl_verify: bool = True
     proxies: Dict[str, str] = Field(default_factory=dict)
+    timeout: TimeIntervalConfig = Field(default_factory=lambda: TimeIntervalConfig("30s"))
 
 
 class ConnectionConfig(ConfigModel):
@@ -133,6 +145,74 @@ class ConnectionConfig(ConfigModel):
 
     connection: _ConnectionParameters = Field(default_factory=_ConnectionParameters)
 
+    def get_cognite_client(self, client_name: str) -> CogniteClient:
+        from cognite.client.config import global_config
+
+        global_config.disable_pypi_version_check = True
+        global_config.disable_gzip = not self.connection.gzip_compression
+        global_config.status_forcelist = set(self.connection.status_forcelist)
+        global_config.max_retries = self.connection.max_retries
+        global_config.max_retries_connect = self.connection.max_retries_connect
+        global_config.max_retry_backoff = self.connection.max_retry_backoff.seconds
+        global_config.max_connection_pool_size = self.connection.max_connection_pool_size
+        global_config.disable_ssl = not self.connection.ssl_verify
+        global_config.proxies = self.connection.proxies
+
+        credential_provider: CredentialProvider
+        match self.authentication:
+            case _ClientCredentialsConfig() as client_credentials:
+                kwargs = {
+                    "token_url": client_credentials.token_url,
+                    "client_id": client_credentials.client_id,
+                    "client_secret": client_credentials.client_secret,
+                    "scopes": client_credentials.scopes,
+                }
+                if client_credentials.audience is not None:
+                    kwargs["audience"] = client_credentials.audience
+                if client_credentials.resource is not None:
+                    kwargs["resource"] = client_credentials.resource
+
+                credential_provider = OAuthClientCredentials(**kwargs)  # type: ignore  # I know what I'm doing
+
+            case _ClientCertificateConfig() as client_certificate:
+                thumbprint, key = _load_certificate_data(
+                    client_certificate.path,
+                    client_certificate.password,
+                )
+                credential_provider = OAuthClientCertificate(
+                    authority_url=client_certificate.authority_url,
+                    client_id=client_certificate.client_id,
+                    cert_thumbprint=str(thumbprint),
+                    certificate=str(key),
+                    scopes=client_certificate.scopes,
+                )
+
+            case _:
+                assert_never(self.authentication)
+
+        client_config = ClientConfig(
+            project=self.project,
+            base_url=self.base_url,
+            client_name=client_name,
+            timeout=self.connection.timeout.seconds,
+            credentials=credential_provider,
+        )
+
+        return CogniteClient(client_config)
+
+
+class CronConfig(ConfigModel):
+    type: Literal["cron"]
+    expression: str
+
+
+class IntervalConfig(ConfigModel):
+    type: Literal["interval"]
+    expression: TimeIntervalConfig
+
+
+ScheduleConfig = Annotated[CronConfig | IntervalConfig, Field(discriminator="type")]
+
 
 class LogLevel(Enum):
     CRITICAL = "CRITICAL"

cognite_extractor_utils-7.4.8/cognite/extractorutils/unstable/core/__main__.py

@@ -0,0 +1,31 @@
+"""
+Example of how you would build an extractor with the new base class
+"""
+
+from cognite.extractorutils.unstable.configuration.models import ExtractorConfig
+
+from .base import Extractor
+from .runtime import Runtime
+
+
+class MyConfig(ExtractorConfig):
+    parameter_one: int
+    parameter_two: str
+
+
+class MyExtractor(Extractor[MyConfig]):
+    NAME = "Test extractor"
+    EXTERNAL_ID = "test-extractor"
+    DESCRIPTION = "Test of the new runtime"
+    VERSION = "1.0.0"
+    CONFIG_TYPE = MyConfig
+
+    def run(self) -> None:
+        self.logger.info("Started!")
+        if not self.cancellation_token.wait(10):
+            raise ValueError("Oops")
+
+
+if __name__ == "__main__":
+    runtime = Runtime(MyExtractor)
+    runtime.run()

cognite_extractor_utils-7.4.8/cognite/extractorutils/unstable/core/_messaging.py

@@ -0,0 +1,5 @@
+from enum import Enum
+
+
+class RuntimeMessage(Enum):
+    RESTART = 1

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

@@ -0,0 +1,116 @@
+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()

cognite_extractor_utils-7.4.8/cognite/extractorutils/unstable/core/runtime.py

@@ -0,0 +1,171 @@
+import logging
+import os
+import sys
+import time
+from argparse import ArgumentParser, Namespace
+from multiprocessing import Process, Queue
+from pathlib import Path
+from typing import Any, Generic, Type, TypeVar
+
+from typing_extensions import assert_never
+
+from cognite.extractorutils.threading import CancellationToken
+from cognite.extractorutils.unstable.configuration.loaders import load_file, load_from_cdf
+from cognite.extractorutils.unstable.configuration.models import ConnectionConfig
+
+from ._messaging import RuntimeMessage
+from .base import ConfigRevision, ConfigType, Extractor
+
+ExtractorType = TypeVar("ExtractorType", bound=Extractor)
+
+
+class Runtime(Generic[ExtractorType]):
+    def __init__(
+        self,
+        extractor: Type[ExtractorType],
+    ) -> None:
+        self._extractor_class = extractor
+        self._cancellation_token = CancellationToken()
+        self._cancellation_token.cancel_on_interrupt()
+        self._message_queue: Queue[RuntimeMessage] = Queue()
+        self.logger = logging.getLogger(f"{self._extractor_class.EXTERNAL_ID}.runtime")
+        self._setup_logging()
+
+    def _create_argparser(self) -> ArgumentParser:
+        argparser = ArgumentParser(
+            prog=sys.argv[0],
+            description=self._extractor_class.DESCRIPTION,
+        )
+        argparser.add_argument(
+            "-v",
+            "--version",
+            action="version",
+            version=f"{self._extractor_class.NAME} v{self._extractor_class.VERSION}",
+        )
+        argparser.add_argument(
+            "-c",
+            "--connection-config",
+            nargs=1,
+            type=Path,
+            required=True,
+            help="Connection parameters",
+        )
+        argparser.add_argument(
+            "-l",
+            "--local-override",
+            nargs=1,
+            type=Path,
+            required=False,
+            default=None,
+            help="Include to use a local application configuration instead of fetching it from CDF",
+        )
+
+        return argparser
+
+    def _setup_logging(self) -> None:
+        # TODO: Figure out file logging for runtime
+        fmt = logging.Formatter(
+            "%(asctime)s.%(msecs)03d UTC [%(levelname)-8s] %(threadName)s - %(message)s",
+            "%Y-%m-%d %H:%M:%S",
+        )
+        # Set logging to UTC
+        fmt.converter = time.gmtime
+
+        root = logging.getLogger()
+        root.setLevel(logging.INFO)
+
+        console_handler = logging.StreamHandler()
+        console_handler.setFormatter(fmt)
+
+        root.addHandler(console_handler)
+
+    def _inner_run(
+        self,
+        message_queue: Queue,
+        connection_config: ConnectionConfig,
+        application_config: ConfigType,
+        current_config_revision: ConfigRevision,
+    ) -> None:
+        # This code is run inside the new extractor process
+        extractor = self._extractor_class.init_from_runtime(
+            connection_config,
+            application_config,
+            current_config_revision,
+        )
+        extractor._set_runtime_message_queue(message_queue)
+
+        try:
+            with extractor:
+                extractor.run()
+
+        except Exception:
+            self.logger.exception("Extractor crashed, will attempt restart")
+            message_queue.put(RuntimeMessage.RESTART)
+
+    def _spawn_extractor(
+        self,
+        connection_config: ConnectionConfig,
+        application_config: ConfigType,
+        current_config_revision: ConfigRevision,
+    ) -> Process:
+        self._message_queue = Queue()
+        process = Process(
+            target=self._inner_run,
+            args=(self._message_queue, connection_config, application_config, current_config_revision),
+        )
+
+        process.start()
+        self.logger.info(f"Started extractor as {process.pid}")
+        return process
+
+    def _get_application_config(
+        self,
+        args: Namespace,
+        connection_config: ConnectionConfig,
+    ) -> tuple[ConfigType, ConfigRevision]:
+        current_config_revision: ConfigRevision
+        if args.local_override:
+            current_config_revision = "local"
+            application_config = load_file(args.local_override[0], self._extractor_class.CONFIG_TYPE)
+        else:
+            client = connection_config.get_cognite_client(
+                f"{self._extractor_class.EXTERNAL_ID}-{self._extractor_class.VERSION}"
+            )
+            application_config, current_config_revision = load_from_cdf(
+                client,
+                connection_config.extraction_pipeline,
+                self._extractor_class.CONFIG_TYPE,
+            )
+
+        return application_config, current_config_revision
+
+    def run(self) -> None:
+        argparser = self._create_argparser()
+        args = argparser.parse_args()
+
+        self.logger.info(f"Started runtime as {os.getpid()}")
+
+        connection_config = load_file(args.connection_config[0], ConnectionConfig)
+
+        # This has to be Any. We don't know the type of the extractors' config at type checking since the sel doesn't
+        # exist yet, and I have not found a way to represent it in a generic way that isn't just an Any in disguise.
+        application_config: Any
+        while not self._cancellation_token.is_cancelled:
+            application_config, current_config_revision = self._get_application_config(args, connection_config)
+            # Start extractor in separate process, and wait for it to end
+            process = self._spawn_extractor(connection_config, application_config, current_config_revision)
+            process.join()
+
+            # Check if we are asked to restart the extractor, shut down otherwise
+            if not self._message_queue.empty():
+                message = self._message_queue.get_nowait()
+                match message:
+                    case RuntimeMessage.RESTART:
+                        continue
+
+                    case _:
+                        assert_never(message)
+
+            else:
+                self.logger.info("Shutting down runtime")
+                self._cancellation_token.cancel()

cognite_extractor_utils-7.4.8/cognite/extractorutils/unstable/scheduling/__init__.py

@@ -0,0 +1,3 @@
+from ._scheduler import TaskScheduler
+
+__all__ = ["TaskScheduler"]

cognite_extractor_utils-7.4.8/cognite/extractorutils/unstable/scheduling/_scheduler.py

@@ -0,0 +1,102 @@
+from dataclasses import dataclass
+from logging import getLogger
+from threading import RLock, Thread
+from time import time
+from typing import Callable
+
+import arrow
+from humps import pascalize
+
+from cognite.extractorutils.threading import CancellationToken
+from cognite.extractorutils.unstable.configuration.models import CronConfig, IntervalConfig, ScheduleConfig
+from cognite.extractorutils.unstable.scheduling._schedules import CronSchedule, IntervalSchedule, Schedule
+
+
+@dataclass
+class Job:
+    name: str
+    call: Callable[[], None]
+    schedule: Schedule
+
+    def __hash__(self) -> int:
+        return hash(self.name)
+
+
+class TaskScheduler:
+    def __init__(self, cancellation_token: CancellationToken) -> None:
+        self._cancellation_token = cancellation_token
+        self._jobs: dict[str, Job] = {}
+        self._jobs_lock = RLock()
+        self._running: set[Job] = set()
+        self._running_lock = RLock()
+
+        self._logger = getLogger()
+
+    def schedule_task(self, name: str, schedule: ScheduleConfig, task: Callable[[], None]) -> None:
+        parsed_schedule: Schedule
+        match schedule:
+            case CronConfig() as cron_config:
+                parsed_schedule = CronSchedule(expression=cron_config.expression)
+
+            case IntervalConfig() as interval_config:
+                parsed_schedule = IntervalSchedule(interval=interval_config.expression.seconds)
+
+        with self._jobs_lock:
+            self._jobs[name] = Job(name=name, call=task, schedule=parsed_schedule)
+
+    def _get_next(self) -> list[Job]:
+        if not self._jobs:
+            return []
+        with self._jobs_lock:
+            next_runs = sorted([(j.schedule.next(), j) for j in self._jobs.values()], key=lambda tup: tup[0])
+        return [job for (next, job) in next_runs if next == next_runs[0][0]] if next_runs else []
+
+    def _run_job(self, job: Job) -> bool:
+        with self._running_lock:
+            if job in self._running:
+                self._logger.warning(f"Job {job.name} already running")
+                return False
+
+        def wrap() -> None:
+            with self._running_lock:
+                self._running.add(job)
+            try:
+                job.call()
+
+                self._logger.info(f"Job {job.name} done. Next run at {arrow.get(job.schedule.next()).isoformat()}")
+
+            finally:
+                with self._running_lock:
+                    self._running.remove(job)
+
+        Thread(target=wrap, name=f"Run{pascalize(job.name)}").start()
+        return True
+
+    def trigger(self, name: str) -> bool:
+        return self._run_job(self._jobs[name])
+
+    def run(self) -> None:
+        if not self._jobs:
+            raise ValueError("Can't run scheduler without any scheduled tasks")
+
+        # Run all interval jobs on startup since the first next() is one interval from now
+        for job in [j for j in self._jobs.values() if isinstance(j.schedule, IntervalSchedule)]:
+            self.trigger(job.name)
+
+        while not self._cancellation_token.is_cancelled:
+            next_runs = self._get_next()
+
+            next_time = next_runs[0].schedule.next()
+            wait_time = max(next_time - time(), 0)
+
+            if wait_time:
+                self._logger.info(f"Waiting until {arrow.get(next_time).isoformat()}")
+                if self._cancellation_token.wait(wait_time):
+                    break
+
+            for job in next_runs:
+                self._logger.info(f"Starting job {job.name}")
+                self._run_job(job)
+
+    def stop(self) -> None:
+        self._cancellation_token.cancel()

cognite_extractor_utils-7.4.8/cognite/extractorutils/unstable/scheduling/_schedules.py

@@ -0,0 +1,31 @@
+from abc import ABC, abstractmethod
+from time import time
+
+from croniter import croniter
+
+
+class Schedule(ABC):
+    @abstractmethod
+    def next(self) -> int:
+        pass
+
+
+class CronSchedule(Schedule):
+    def __init__(self, expression: str) -> None:
+        self._cron = croniter(expression)
+
+    def next(self) -> int:
+        return int(self._cron.get_next(start_time=time()))
+
+
+class IntervalSchedule(Schedule):
+    def __init__(self, interval: int) -> None:
+        self._interval = interval
+        self._next = int(time())
+
+    def next(self) -> int:
+        t = time()
+        while t > self._next:
+            self._next += self._interval
+
+        return self._next

{cognite_extractor_utils-7.4.6 → cognite_extractor_utils-7.4.8}/cognite/extractorutils/uploader/files.py

@@ -19,6 +19,7 @@ from math import ceil
 from os import PathLike
 from types import TracebackType
 from typing import Any, BinaryIO, Callable, Dict, Iterator, List, Optional, Tuple, Type, Union
+from urllib.parse import ParseResult, urlparse
 
 from httpx import URL, Client, Headers, Request, StreamConsumed, SyncByteStream
 from requests.utils import super_len
@@ -255,18 +256,20 @@ class IOFileUploadQueue(AbstractUploadQueue):
     ) -> tuple[FileMetadataOrCogniteExtractorFile, str]:
         if isinstance(file_meta, CogniteExtractorFileApply):
             node_id = self._apply_cognite_file(file_meta)
-            file_meta, url = self._create_cdm(instance_id=node_id)
+            file_meta_response, url = self._create_cdm(instance_id=node_id)
         else:
-            file_meta, url = self.cdf_client.files.create(file_metadata=file_meta, overwrite=self.overwrite_existing)
+            file_meta_response, url = self.cdf_client.files.create(
+                file_metadata=file_meta, overwrite=self.overwrite_existing
+            )
             # trigger update after creation (upsert =P)
             basic_attributes = set(["externalId", "name"])
             attr = set(file_meta.dump().keys())
             diff = attr - basic_attributes
 
             if len(diff) >= 1 and "externalId" in attr:
-                file_meta = self.cdf_client.files.update(file_meta)
+                file_meta_response = self.cdf_client.files.update(file_meta)
 
-        return file_meta, url
+        return file_meta_response, url
 
     def _upload_bytes(self, size: int, file: BinaryIO, file_meta: FileMetadataOrCogniteExtractorFile) -> None:
         file_meta, url = self._upload_empty(file_meta)
@@ -406,12 +409,22 @@ class IOFileUploadQueue(AbstractUploadQueue):
         self, url_str: str, stream: BinaryIO, size: int, mime_type: Optional[str] = None
     ) -> Request:
         url = URL(url_str)
+        base_url = URL(self.cdf_client.config.base_url)
+
+        if url.host == base_url.host:
+            upload_url = url
+        else:
+            parsed_url: ParseResult = urlparse(url_str)
+            parsed_base_url: ParseResult = urlparse(self.cdf_client.config.base_url)
+            replaced_upload_url = parsed_url._replace(netloc=parsed_base_url.netloc).geturl()
+            upload_url = URL(replaced_upload_url)
+
         headers = Headers(self._httpx_client.headers)
         headers.update(
             {
                 "Accept": "*/*",
                 "Content-Length": str(size),
-                "Host": url.netloc.decode("ascii"),
+                "Host": upload_url.netloc.decode("ascii"),
                 "x-cdp-app": self.cdf_client._config.client_name,
             }
         )
@@ -421,7 +434,7 @@ class IOFileUploadQueue(AbstractUploadQueue):
 
         return Request(
             method="PUT",
-            url=url,
+            url=upload_url,
             stream=IOByteStream(stream),
             headers=headers,
         )

{cognite_extractor_utils-7.4.6 → cognite_extractor_utils-7.4.8}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "cognite-extractor-utils"
-version = "7.4.6"
+version = "7.4.8"
 description = "Utilities for easier development of extractors for CDF"
 authors = ["Mathias Lohne <mathias.lohne@cognite.com>"]
 license = "Apache-2.0"
@@ -74,6 +74,7 @@ orjson = "^3.10.3"
 httpx = "^0.27.0"
 pydantic = "^2.8.2"
 pyhumps = "^3.8.0"
+croniter = "^3.0.3"
 
 [tool.poetry.extras]
 experimental = ["cognite-sdk-experimental"]
@@ -93,7 +94,7 @@ parameterized = "*"
 requests = "^2.31.0"
 types-requests = "^2.31.0.20240125"
 httpx = "^0.27.0"
-faker = "^28.0.0"
+faker = "^30.0.0"
 
 [build-system]
 requires = ["poetry-core>=1.0.0"]