cognite-extractor-utils 7.6.0__py3-none-any.whl → 7.8.0__py3-none-any.whl

cognite/examples/unstable/extractors/simple_extractor/config/config.yaml

@@ -0,0 +1,3 @@
+log-handlers:
+  - type: console
+    level: INFO

cognite/examples/unstable/extractors/simple_extractor/config/connection_config.yaml

@@ -0,0 +1,10 @@
+project: ${COGNITE_PROJECT}
+base_url: ${COGNITE_BASE_URL}
+integration:
+  external_id: ${COGNITE_INTEGRATION_ID}
+authentication:
+  type: "client-credentials"
+  client_id: ${COGNITE_CLIENT_ID}
+  client_secret: ${COGNITE_CLIENT_SECRET}
+  token_url: ${COGNITE_TOKEN_URL}
+  scopes: ${COGNITE_SCOPES}

cognite/examples/unstable/extractors/simple_extractor/main.py

@@ -0,0 +1,81 @@
+"""
+An example extractor that logs messages at various levels.
+"""
+
+from cognite.extractorutils.unstable.configuration.models import ExtractorConfig, IntervalConfig, TimeIntervalConfig
+from cognite.extractorutils.unstable.core.base import Extractor, StartupTask, TaskContext
+from cognite.extractorutils.unstable.core.runtime import Runtime
+from cognite.extractorutils.unstable.core.tasks import ScheduledTask
+
+
+class SimpleConfig(ExtractorConfig):
+    """
+    Defines the configuration for the SimpleExtractor.
+    """
+
+    pass
+
+
+class SimpleExtractor(Extractor[SimpleConfig]):
+    """
+    An example extractor that logs messages at various levels.
+    """
+
+    NAME = "SimpleTestExtractor"
+    EXTERNAL_ID = "test-extractor"
+    DESCRIPTION = "An extractor for testing log levels"
+    VERSION = "1.0.0"
+    CONFIG_TYPE = SimpleConfig
+    SUPPORTS_DRY_RUN = True
+
+    def __init_tasks__(self) -> None:
+        """
+        Initializes and adds tasks to the extractor.
+        """
+        self.add_task(StartupTask(name="main_task", target=self.run_my_task))
+        self.add_task(
+            ScheduledTask(
+                name="scheduled_task",
+                target=self.scheduled_task,
+                schedule=IntervalConfig(type="interval", expression=TimeIntervalConfig("3s")),
+            )
+        )
+
+    # example task that logs messages at different levels
+    def run_my_task(self, ctx: TaskContext) -> None:
+        """
+        An example task that logs messages at different levels.
+
+        Args:
+            ctx: The context for the task execution, used for logging.
+        """
+        ctx.debug("This is a detailed debug message.")
+        ctx.info("This is an informational message.")
+        ctx.warning("This is a warning message.")
+        ctx.info("Test finished.")
+
+    def scheduled_task(self, ctx: TaskContext) -> None:
+        """
+        An example scheduled task that logs a message.
+
+        Args:
+            ctx: The context for the task execution, used for logging.
+        """
+        ctx.info("This is a scheduled task running.")
+        ctx.warning("This is a warning from the scheduled task.")
+        ctx.debug("Debugging the scheduled task execution.")
+        ctx.error("This is an error message from the scheduled task.")
+
+    # add more tasks as needed
+
+
+def main() -> None:
+    """
+    Main function to run the SimpleExtractor.
+    """
+    runtime = Runtime(SimpleExtractor)
+    runtime.run()
+
+
+if __name__ == "__main__":
+    main()

cognite/extractorutils/__init__.py

@@ -16,7 +16,7 @@
 Cognite extractor utils is a Python package that simplifies the development of new extractors.
 """
 
-__version__ = "7.6.0"
+__version__ = "7.8.0"
 from .base import Extractor
 
 __all__ = ["Extractor"]

cognite/extractorutils/_inner_util.py

@@ -37,14 +37,14 @@ def resolve_log_level_for_httpx(level: str) -> str:
 
 
 class _DecimalEncoder(json.JSONEncoder):
-    def default(self, obj: Any) -> dict[str, str]:
+    def default(self, obj: Any) -> dict[str, str]:  # noqa: ANN401
         if isinstance(obj, Decimal):
             return {"type": "decimal_encoded", "value": str(obj)}
         return super().default(obj)
 
 
 class _DecimalDecoder(json.JSONDecoder):
-    def __init__(self, *args: Any, **kwargs: Any) -> None:
+    def __init__(self, *args: Any, **kwargs: Any) -> None:  # noqa: ANN401
         json.JSONDecoder.__init__(self, *args, object_hook=self.object_hook, **kwargs)
 
     def object_hook(self, obj_dict: dict[str, str]) -> dict[str, str] | Decimal:

cognite/extractorutils/base.py

@@ -112,7 +112,7 @@ class Extractor(Generic[CustomConfigClass]):
         reload_config_interval: int | None = 300,
         reload_config_action: ReloadConfigAction = ReloadConfigAction.DO_NOTHING,
         success_message: str = "Successful shutdown",
-    ):
+    ) -> None:
         self.name = name
         self.description = description
         self.run_handle = run_handle

cognite/extractorutils/configtools/elements.py

@@ -29,6 +29,7 @@ from urllib.parse import urljoin, urlparse
 
 import yaml
 from prometheus_client import REGISTRY, start_http_server
+from typing_extensions import Self
 
 from cognite.client import ClientConfig, CogniteClient
 from cognite.client.credentials import (
@@ -604,6 +605,7 @@ class LoggingConfig:
                 when="midnight",
                 utc=True,
                 backupCount=self.file.retention,
+                encoding="utf-8",
             )
             file_handler.setLevel(self.file.level)
             file_handler.setFormatter(fmt)
@@ -926,7 +928,7 @@ class CastableInt(int):
     file.
     """
 
-    def __new__(cls, value: Any) -> "CastableInt":
+    def __new__(cls, value: int | str | bytes) -> Self:
         """
         Returns value as is if it's int.
 
@@ -955,7 +957,7 @@ class PortNumber(CastableInt):
     not a valid port number raises a ValueError at instantiation.
     """
 
-    def __new__(cls, value: Any) -> "PortNumber":
+    def __new__(cls, value: int | str | bytes) -> Self:
         """
         Try to cast the value to an integer and validate it as a port number.
 

cognite/extractorutils/configtools/loaders.py

@@ -72,7 +72,7 @@ class KeyVaultLoader:
         config: A dictionary containing the configuration for the keyvault.
     """
 
-    def __init__(self, config: dict | None):
+    def __init__(self, config: dict | None) -> None:
         self.config = config
 
         self.client: SecretClient | None = None
@@ -374,7 +374,7 @@ class ConfigResolver(Generic[CustomConfigClass]):
     Automatically reloads the configuration file if it has changed
     """
 
-    def __init__(self, config_path: str, config_type: type[CustomConfigClass]):
+    def __init__(self, config_path: str, config_type: type[CustomConfigClass]) -> None:
         self.config_path = config_path
         self.config_type = config_type
 
@@ -384,7 +384,7 @@ class ConfigResolver(Generic[CustomConfigClass]):
         self._cognite_client: CogniteClient | None = None
 
     def _reload_file(self) -> None:
-        with open(self.config_path) as stream:
+        with open(self.config_path, encoding="utf-8") as stream:
             self._config_text = stream.read()
 
     @property

cognite/extractorutils/exceptions.py

@@ -25,7 +25,7 @@ class InvalidConfigError(Exception):
       * Unknown fields
     """
 
-    def __init__(self, message: str, details: list[str] | None = None):
+    def __init__(self, message: str, details: list[str] | None = None) -> None:
         super().__init__()
         self.message = message
         self.details = details

cognite/extractorutils/metrics.py

@@ -67,7 +67,7 @@ _metrics_singularities = {}
 T = TypeVar("T")
 
 
-def safe_get(cls: type[T], *args: Any, **kwargs: Any) -> T:
+def safe_get(cls: type[T], *args: Any, **kwargs: Any) -> T:  # noqa: ANN401
     """
     A factory for instances of metrics collections.
 
@@ -122,7 +122,7 @@ class BaseMetrics:
         process_scrape_interval: Interval (in seconds) between each fetch of data for the ``process_*`` gauges
     """
 
-    def __init__(self, extractor_name: str, extractor_version: str, process_scrape_interval: float = 15):
+    def __init__(self, extractor_name: str, extractor_version: str, process_scrape_interval: float = 15) -> None:
         extractor_name = extractor_name.strip().replace(" ", "_")
 
         self.startup = Gauge(f"{extractor_name}_start_time", "Timestamp (seconds) of when the extractor last started")
@@ -187,7 +187,7 @@ class AbstractMetricsPusher(ABC):
         push_interval: int | None = None,
         thread_name: str | None = None,
         cancellation_token: CancellationToken | None = None,
-    ):
+    ) -> None:
         self.push_interval = push_interval
         self.thread_name = thread_name
 
@@ -274,7 +274,7 @@ class PrometheusPusher(AbstractMetricsPusher):
         password: str | None = None,
         thread_name: str | None = None,
         cancellation_token: CancellationToken | None = None,
-    ):
+    ) -> None:
         super().__init__(push_interval, thread_name, cancellation_token)
 
         self.username = username
@@ -283,7 +283,9 @@ 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: bytes
+    ) -> Callable[[], None]:
         """
         Returns a authentication handler against the Prometheus Pushgateway to use in the pushadd_to_gateway method.
 
@@ -350,7 +352,7 @@ class CognitePusher(AbstractMetricsPusher):
         data_set: EitherId | None = None,
         thread_name: str | None = None,
         cancellation_token: CancellationToken | None = None,
-    ):
+    ) -> None:
         super().__init__(push_interval, thread_name, cancellation_token)
 
         self.cdf_client = cdf_client

cognite/extractorutils/statestore/watermark.py

@@ -1,3 +1,6 @@
+# ruff: noqa: ANN401
+# TODO: the state stores should be generic over the type of state, not just Any.
+
 #  Copyright 2020 Cognite AS
 #
 #  Licensed under the Apache License, Version 2.0 (the "License");
@@ -66,7 +69,7 @@ class AbstractStateStore(_BaseStateStore, ABC):
         trigger_log_level: str = "DEBUG",
         thread_name: str | None = None,
         cancellation_token: CancellationToken | None = None,
-    ):
+    ) -> None:
         super().__init__(
             save_interval=save_interval,
             trigger_log_level=trigger_log_level,
@@ -249,7 +252,7 @@ class RawStateStore(AbstractStateStore):
         trigger_log_level: str = "DEBUG",
         thread_name: str | None = None,
         cancellation_token: CancellationToken | None = None,
-    ):
+    ) -> None:
         super().__init__(save_interval, trigger_log_level, thread_name, cancellation_token)
 
         self._cdf_client = cdf_client
@@ -395,7 +398,7 @@ class LocalStateStore(AbstractStateStore):
         trigger_log_level: str = "DEBUG",
         thread_name: str | None = None,
         cancellation_token: CancellationToken | None = None,
-    ):
+    ) -> None:
         super().__init__(save_interval, trigger_log_level, thread_name, cancellation_token)
 
         self._file_path = file_path

cognite/extractorutils/threading.py

@@ -6,7 +6,7 @@ import logging
 import signal
 from threading import Condition
 from time import time
-from typing import Any
+from types import FrameType
 
 
 class CancellationToken:
@@ -114,7 +114,7 @@ class CancellationToken:
         This will set the cancellation token instead of throwing a KeyboardInterrupt exception.
         """
 
-        def sigint_handler(sig_num: int, frame: Any) -> None:
+        def sigint_handler(sig_num: int, frame: FrameType | None) -> None:
             logger = logging.getLogger(__name__)
             logger.warning("Interrupt signal received, stopping extractor gracefully")
             self.cancel()

cognite/extractorutils/unstable/configuration/exceptions.py

@@ -12,7 +12,7 @@ class InvalidConfigError(Exception):
       * Unknown fields
     """
 
-    def __init__(self, message: str, details: list[str] | None = None):
+    def __init__(self, message: str, details: list[str] | None = None) -> None:
         super().__init__()
         self.message = message
         self.details = details
@@ -30,3 +30,30 @@ class InvalidConfigError(Exception):
         Underlying message prefixed with 'Invalid config:'.
         """
         return self.__str__()
+
+
+class InvalidArgumentError(Exception):
+    """
+    Exception thrown when an invalid argument is passed to the extractor.
+
+    This can be due to:
+      * Missing required arguments
+      * Invalid argument types
+      * Unsupported argument values
+    """
+
+    def __init__(self, message: str) -> None:
+        super().__init__(message)
+        self.message = message
+
+    def __str__(self) -> str:
+        """
+        Underlying message prefixed with 'Invalid argument:'.
+        """
+        return f"Invalid argument: {self.message}"
+
+    def __repr__(self) -> str:
+        """
+        Underlying message prefixed with 'Invalid argument:'.
+        """
+        return self.__str__()

cognite/extractorutils/unstable/configuration/models.py

@@ -4,10 +4,11 @@ Module containing pre-built models for common extractor configuration.
 
 import os
 import re
+from collections.abc import Iterator
 from datetime import timedelta
 from enum import Enum
 from pathlib import Path
-from typing import Annotated, Any, Literal
+from typing import Annotated, Any, Literal, TypeVar
 
 from humps import kebabize
 from pydantic import BaseModel, ConfigDict, Field, GetCoreSchemaHandler
@@ -23,6 +24,13 @@ from cognite.client.credentials import (
 )
 from cognite.extractorutils.configtools._util import _load_certificate_data
 from cognite.extractorutils.exceptions import InvalidConfigError
+from cognite.extractorutils.statestore import (
+    AbstractStateStore,
+    LocalStateStore,
+    NoStateStore,
+    RawStateStore,
+)
+from cognite.extractorutils.threading import CancellationToken
 
 __all__ = [
     "AuthenticationConfig",
@@ -53,23 +61,44 @@ class ConfigModel(BaseModel):
     )
 
 
-class _ClientCredentialsConfig(ConfigModel):
-    type: Literal["client-credentials"]
+class Scopes(str):
+    def __init__(self, scopes: str) -> None:
+        self._scopes = list(scopes.split(" "))
+
+    @classmethod
+    def __get_pydantic_core_schema__(cls, source_type: Any, handler: GetCoreSchemaHandler) -> CoreSchema:  # noqa: ANN401
+        return core_schema.no_info_after_validator_function(cls, handler(str))
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, Scopes):
+            return NotImplemented
+        return self._scopes == other._scopes
+
+    def __hash__(self) -> int:
+        return hash(self._scopes)
+
+    def __iter__(self) -> Iterator[str]:
+        return iter(self._scopes)
+
+
+class BaseCredentialsConfig(ConfigModel):
     client_id: str
+    scopes: Scopes
+
+
+class _ClientCredentialsConfig(BaseCredentialsConfig):
+    type: Literal["client-credentials"]
     client_secret: str
     token_url: str
-    scopes: list[str]
     resource: str | None = None
     audience: str | None = None
 
 
-class _ClientCertificateConfig(ConfigModel):
+class _ClientCertificateConfig(BaseCredentialsConfig):
     type: Literal["client-certificate"]
-    client_id: str
     path: Path
     password: str | None = None
     authority_url: str
-    scopes: list[str]
 
 
 AuthenticationConfig = Annotated[_ClientCredentialsConfig | _ClientCertificateConfig, Field(discriminator="type")]
@@ -84,7 +113,7 @@ class TimeIntervalConfig:
         self._interval, self._expression = TimeIntervalConfig._parse_expression(expression)
 
     @classmethod
-    def __get_pydantic_core_schema__(cls, source_type: Any, handler: GetCoreSchemaHandler) -> CoreSchema:
+    def __get_pydantic_core_schema__(cls, source_type: Any, handler: GetCoreSchemaHandler) -> CoreSchema:  # noqa: ANN401
         """
         Pydantic hook to define how this class should be serialized/deserialized.
 
@@ -191,18 +220,26 @@ class TimeIntervalConfig:
         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 RetriesConfig(ConfigModel):
+    max_retries: int = Field(default=10, ge=-1)
+    max_backoff: TimeIntervalConfig = Field(default_factory=lambda: TimeIntervalConfig("30s"))
     timeout: TimeIntervalConfig = Field(default_factory=lambda: TimeIntervalConfig("30s"))
 
 
+class SslCertificatesConfig(ConfigModel):
+    verify: bool = True
+    allow_list: list[str] | None = None
+
+
+class ConnectionParameters(ConfigModel):
+    retries: RetriesConfig = Field(default_factory=RetriesConfig)
+    ssl_certificates: SslCertificatesConfig = Field(default_factory=SslCertificatesConfig)
+
+
+class IntegrationConfig(ConfigModel):
+    external_id: str
+
+
 class ConnectionConfig(ConfigModel):
     """
     Configuration for connecting to a Cognite Data Fusion project.
@@ -216,11 +253,11 @@ class ConnectionConfig(ConfigModel):
     project: str
     base_url: str
 
-    integration: str
+    integration: IntegrationConfig
 
     authentication: AuthenticationConfig
 
-    connection: _ConnectionParameters = Field(default_factory=_ConnectionParameters)
+    connection: ConnectionParameters = Field(default_factory=ConnectionParameters)
 
     def get_cognite_client(self, client_name: str) -> CogniteClient:
         """
@@ -235,14 +272,9 @@ class ConnectionConfig(ConfigModel):
         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
+        global_config.max_retries = self.connection.retries.max_retries
+        global_config.max_retry_backoff = self.connection.retries.max_backoff.seconds
+        global_config.disable_ssl = not self.connection.ssl_certificates.verify
 
         credential_provider: CredentialProvider
         match self.authentication:
@@ -270,7 +302,7 @@ class ConnectionConfig(ConfigModel):
                     client_id=client_certificate.client_id,
                     cert_thumbprint=str(thumbprint),
                     certificate=str(key),
-                    scopes=client_certificate.scopes,
+                    scopes=list(client_certificate.scopes),
                 )
 
             case _:
@@ -280,7 +312,7 @@ class ConnectionConfig(ConfigModel):
             project=self.project,
             base_url=self.base_url,
             client_name=client_name,
-            timeout=self.connection.timeout.seconds,
+            timeout=self.connection.retries.timeout.seconds,
             credentials=credential_provider,
         )
 
@@ -315,7 +347,9 @@ class ConnectionConfig(ConfigModel):
                 client_id=os.environ["COGNITE_CLIENT_ID"],
                 client_secret=os.environ["COGNITE_CLIENT_SECRET"],
                 token_url=os.environ["COGNITE_TOKEN_URL"],
-                scopes=os.environ["COGNITE_TOKEN_SCOPES"].split(","),
+                scopes=Scopes(
+                    os.environ["COGNITE_TOKEN_SCOPES"],
+                ),
             )
         elif "COGNITE_CLIENT_CERTIFICATE_PATH" in os.environ:
             auth = _ClientCertificateConfig(
@@ -324,7 +358,9 @@ class ConnectionConfig(ConfigModel):
                 path=Path(os.environ["COGNITE_CLIENT_CERTIFICATE_PATH"]),
                 password=os.environ.get("COGNITE_CLIENT_CERTIFICATE_PATH"),
                 authority_url=os.environ["COGNITE_AUTHORITY_URL"],
-                scopes=os.environ["COGNITE_TOKEN_SCOPES"].split(","),
+                scopes=Scopes(
+                    os.environ["COGNITE_TOKEN_SCOPES"],
+                ),
             )
         else:
             raise KeyError("Missing auth, either COGNITE_CLIENT_SECRET or COGNITE_CLIENT_CERTIFICATE_PATH must be set")
@@ -332,7 +368,7 @@ class ConnectionConfig(ConfigModel):
         return ConnectionConfig(
             project=os.environ["COGNITE_PROJECT"],
             base_url=os.environ["COGNITE_BASE_URL"],
-            integration=os.environ["COGNITE_INTEGRATION"],
+            integration=IntegrationConfig(external_id=os.environ["COGNITE_INTEGRATION"]),
             authentication=auth,
         )
 
@@ -398,9 +434,98 @@ def _log_handler_default() -> list[LogHandlerConfig]:
     return [LogConsoleHandlerConfig(type="console", level=LogLevel.INFO)]
 
 
+class RawDestinationConfig(ConfigModel):
+    """
+    Configuration parameters for using Raw.
+    """
+
+    database: str
+    table: str
+
+
+class RawStateStoreConfig(RawDestinationConfig):
+    """
+    Configuration of a state store based on CDF RAW.
+    """
+
+    upload_interval: TimeIntervalConfig = Field(default_factory=lambda: TimeIntervalConfig("30s"))
+
+
+class LocalStateStoreConfig(ConfigModel):
+    """
+    Configuration of a state store using a local JSON file.
+    """
+
+    path: Path
+    save_interval: TimeIntervalConfig = Field(default_factory=lambda: TimeIntervalConfig("30s"))
+
+
+class StateStoreConfig(ConfigModel):
+    """
+    Configuration of the State Store, containing ``LocalStateStoreConfig`` or ``RawStateStoreConfig``.
+    """
+
+    raw: RawStateStoreConfig | None = None
+    local: LocalStateStoreConfig | None = None
+
+    def create_state_store(
+        self,
+        cdf_client: CogniteClient | None = None,
+        default_to_local: bool = True,
+        cancellation_token: CancellationToken | None = None,
+    ) -> AbstractStateStore:
+        """
+        Create a state store object based on the config.
+
+        Args:
+            cdf_client: CogniteClient object to use in case of a RAW state store (ignored otherwise)
+            default_to_local: If true, return a LocalStateStore if no state store is configured. Otherwise return a
+                NoStateStore
+            cancellation_token: Cancellation token to pass to created state stores
+
+        Returns:
+            An (uninitialized) state store
+        """
+        if self.raw and self.local:
+            raise ValueError("Only one state store can be used simultaneously")
+
+        if self.raw:
+            if cdf_client is None:
+                raise TypeError("A cognite client object must be provided when state store is RAW")
+
+            return RawStateStore(
+                cdf_client=cdf_client,
+                database=self.raw.database,
+                table=self.raw.table,
+                save_interval=self.raw.upload_interval.seconds,
+                cancellation_token=cancellation_token,
+            )
+
+        if self.local:
+            if self.local.path.is_dir():
+                raise IsADirectoryError(self.local.path)
+
+            return LocalStateStore(
+                file_path=str(self.local.path),
+                save_interval=self.local.save_interval.seconds,
+                cancellation_token=cancellation_token,
+            )
+
+        if default_to_local:
+            return LocalStateStore(file_path="states.json", cancellation_token=cancellation_token)
+
+        return NoStateStore()
+
+
 class ExtractorConfig(ConfigModel):
     """
     Base class for application configuration for extractors.
     """
 
+    state_store: StateStoreConfig | None = None
     log_handlers: list[LogHandlerConfig] = Field(default_factory=_log_handler_default)
+    retry_startup: bool = True
+
+
+ConfigType = TypeVar("ConfigType", bound=ExtractorConfig)
+ConfigRevision = Literal["local"] | int