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

cognite/extractorutils/unstable/core/runtime.py

@@ -6,15 +6,23 @@ from argparse import ArgumentParser, Namespace
 from multiprocessing import Process, Queue
 from pathlib import Path
 from typing import Any, Generic, Type, TypeVar
+from uuid import uuid4
 
+from requests.exceptions import ConnectionError
 from typing_extensions import assert_never
 
+from cognite.client.exceptions import CogniteAPIError, CogniteAuthError, CogniteConnectionError
 from cognite.extractorutils.threading import CancellationToken
+from cognite.extractorutils.unstable.configuration.exceptions import InvalidConfigError
 from cognite.extractorutils.unstable.configuration.loaders import load_file, load_from_cdf
 from cognite.extractorutils.unstable.configuration.models import ConnectionConfig
+from cognite.extractorutils.unstable.core._dto import Error
+from cognite.extractorutils.util import now
 
 from ._messaging import RuntimeMessage
-from .base import ConfigRevision, ConfigType, Extractor
+from .base import ConfigRevision, ConfigType, Extractor, FullConfig
+
+__all__ = ["Runtime", "ExtractorType"]
 
 ExtractorType = TypeVar("ExtractorType", bound=Extractor)
 
@@ -59,13 +67,18 @@ class Runtime(Generic[ExtractorType]):
             default=None,
             help="Include to use a local application configuration instead of fetching it from CDF",
         )
+        argparser.add_argument(
+            "--skip-init-checks",
+            action="store_true",
+            help="Skip any checks during startup. Useful for debugging, not recommended for production deployments.",
+        )
 
         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",
+            "%(asctime)s.%(msecs)03d UTC [%(levelname)-8s] %(process)d %(threadName)s - %(message)s",
             "%Y-%m-%d %H:%M:%S",
         )
         # Set logging to UTC
@@ -82,16 +95,10 @@ class Runtime(Generic[ExtractorType]):
     def _inner_run(
         self,
         message_queue: Queue,
-        connection_config: ConnectionConfig,
-        application_config: ConfigType,
-        current_config_revision: ConfigRevision,
+        config: FullConfig,
     ) -> 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 = self._extractor_class._init_from_runtime(config)
         extractor._set_runtime_message_queue(message_queue)
 
         try:
@@ -104,56 +111,193 @@ class Runtime(Generic[ExtractorType]):
 
     def _spawn_extractor(
         self,
-        connection_config: ConnectionConfig,
-        application_config: ConfigType,
-        current_config_revision: ConfigRevision,
+        config: FullConfig,
     ) -> Process:
         self._message_queue = Queue()
         process = Process(
             target=self._inner_run,
-            args=(self._message_queue, connection_config, application_config, current_config_revision),
+            args=(self._message_queue, config),
         )
 
         process.start()
-        self.logger.info(f"Started extractor as {process.pid}")
+        self.logger.info(f"Started extractor with PID {process.pid}")
         return process
 
     def _get_application_config(
         self,
         args: Namespace,
         connection_config: ConnectionConfig,
-    ) -> tuple[ConfigType, ConfigRevision]:
+    ) -> tuple[ConfigType, ConfigRevision, ConfigRevision]:
         current_config_revision: ConfigRevision
+        newest_config_revision: ConfigRevision
+
         if args.local_override:
+            self.logger.info("Loading local application config")
+
             current_config_revision = "local"
-            application_config = load_file(args.local_override[0], self._extractor_class.CONFIG_TYPE)
+            newest_config_revision = "local"
+            try:
+                application_config = load_file(args.local_override[0], self._extractor_class.CONFIG_TYPE)
+            except InvalidConfigError as e:
+                self.logger.critical(str(e))
+                raise e
+            except FileNotFoundError as e:
+                self.logger.critical(str(e))
+                raise InvalidConfigError(str(e)) from e
+
         else:
+            self.logger.info("Loading application config from CDF")
             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,
+
+            errors: list[Error] = []
+
+            revision: int | None = None
+            try:
+                while True:
+                    try:
+                        application_config, current_config_revision = load_from_cdf(
+                            client,
+                            connection_config.integration,
+                            self._extractor_class.CONFIG_TYPE,
+                            revision=revision,
+                        )
+                        break
+
+                    except InvalidConfigError as e:
+                        if e.attempted_revision is None:
+                            # Should never happen, attempted_revision is set in every handler in load_from_cdf, but it's
+                            # needed for type checks to pass
+                            raise e
+
+                        self.logger.error(f"Revision {e.attempted_revision} is invalid: {e.message}")
+
+                        t = now()
+                        errors.append(
+                            Error(
+                                external_id=str(uuid4()),
+                                level="error",
+                                description=f"Revision {e.attempted_revision} is invalid",
+                                details=e.message,
+                                start_time=t,
+                                end_time=t,
+                                task=None,
+                            )
+                        )
+
+                        if revision is None:
+                            revision = e.attempted_revision - 1
+                            newest_config_revision = e.attempted_revision
+                        else:
+                            revision -= 1
+
+                        if revision > 0:
+                            self.logger.info(f"Falling back to revision {revision}")
+                        else:
+                            self.logger.critical("No more revisions to fall back to")
+                            raise e
+
+            finally:
+                if errors:
+                    client.post(
+                        f"/api/v1/projects/{client.config.project}/odin/checkin",
+                        json={
+                            "externalId": connection_config.integration,
+                            "errors": [e.model_dump() for e in errors],
+                        },
+                        headers={"cdf-version": "alpha"},
+                    )
+
+        return application_config, current_config_revision, newest_config_revision
+
+    def _verify_connection_config(self, connection_config: ConnectionConfig) -> bool:
+        client = connection_config.get_cognite_client(
+            f"{self._extractor_class.EXTERNAL_ID}-{self._extractor_class.VERSION}"
+        )
+        try:
+            client.post(
+                f"/api/v1/projects/{client.config.project}/odin/checkin",
+                json={
+                    "externalId": connection_config.integration,
+                },
+                headers={"cdf-version": "alpha"},
             )
 
-        return application_config, current_config_revision
+        except CogniteConnectionError as e:
+            if e.__cause__ is not None:
+                self.logger.error(str(e.__cause__))
+            self.logger.critical("Could not connect to CDF. Please check your configuration.")
+            return False
+
+        except CogniteAuthError as e:
+            # Error while fetching auth token
+            self.logger.error(str(e))
+            self.logger.critical("Could not get an access token. Please check your configuration.")
+            return False
+
+        except CogniteAPIError as e:
+            # Error response from the CDF API
+            if e.code == 401:
+                self.logger.critical(
+                    "Got a 401 error from CDF. Please check your configuration. "
+                    "Make sure the credentials and project is correct."
+                )
+
+            elif e.message:
+                self.logger.critical(str(e.message))
+
+            else:
+                self.logger.critical(f"Error while connecting to CDF {str(e)}")
+
+            return False
+
+        except ConnectionError as e:
+            # This is sometime thrown, I've seen it when trying to get an auth token but it might happen elsewhere too
+            self.logger.error(str(e))
+            self.logger.critical("Could not initiate connection. Please check your configuration.")
+            return False
+
+        return True
 
     def run(self) -> None:
         argparser = self._create_argparser()
         args = argparser.parse_args()
 
-        self.logger.info(f"Started runtime as {os.getpid()}")
+        self.logger.info(f"Started runtime with PID {os.getpid()}")
+
+        try:
+            connection_config = load_file(args.connection_config[0], ConnectionConfig)
+        except InvalidConfigError as e:
+            self.logger.error(str(e))
+            self.logger.critical("Could not load connection config")
+            sys.exit(1)
 
-        connection_config = load_file(args.connection_config[0], ConnectionConfig)
+        if not args.skip_init_checks and not self._verify_connection_config(connection_config):
+            sys.exit(1)
 
         # 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)
+            try:
+                application_config, current_config_revision, newest_config_revision = self._get_application_config(
+                    args, connection_config
+                )
+
+            except InvalidConfigError:
+                self.logger.critical("Could not get a valid application config file. Shutting down")
+                sys.exit(1)
+
             # Start extractor in separate process, and wait for it to end
-            process = self._spawn_extractor(connection_config, application_config, current_config_revision)
+            process = self._spawn_extractor(
+                FullConfig(
+                    connection_config=connection_config,
+                    application_config=application_config,
+                    current_config_revision=current_config_revision,
+                    newest_config_revision=newest_config_revision,
+                )
+            )
             process.join()
 
             # Check if we are asked to restart the extractor, shut down otherwise

cognite/extractorutils/unstable/core/tasks.py

@@ -4,6 +4,8 @@ from typing import Callable
 
 from cognite.extractorutils.unstable.configuration.models import ScheduleConfig
 
+__all__ = ["ScheduledTask", "ContinuousTask", "StartupTask", "Task"]
+
 
 @dataclass
 class _Task(ABC):

cognite/extractorutils/unstable/scheduling/_scheduler.py

@@ -56,7 +56,7 @@ class TaskScheduler:
     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")
+                self._logger.warning(f"Job '{job.name}' already running")
                 return False
 
         def wrap() -> None:
@@ -65,13 +65,13 @@ class TaskScheduler:
             try:
                 job.call()
 
-                self._logger.info(f"Job {job.name} done. Next run at {arrow.get(job.schedule.next()).isoformat()}")
+                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()
+        Thread(target=wrap, name=f"{pascalize(job.name)}").start()
         return True
 
     def trigger(self, name: str) -> bool:
@@ -97,7 +97,7 @@ class TaskScheduler:
                     break
 
             for job in next_runs:
-                self._logger.info(f"Starting job {job.name}")
+                self._logger.info(f"Starting job '{job.name}'")
                 self._run_job(job)
 
     def stop(self) -> None:

cognite/extractorutils/uploader/__init__.py

@@ -75,3 +75,17 @@ from .time_series import (
     TimeSeriesUploadQueue,
     default_time_series_factory,
 )
+
+__all__ = [
+    "AssetUploadQueue",
+    "EventUploadQueue",
+    "BytesUploadQueue",
+    "FileUploadQueue",
+    "IOFileUploadQueue",
+    "RawUploadQueue",
+    "DataPoint",
+    "DataPointList",
+    "SequenceUploadQueue",
+    "TimeSeriesUploadQueue",
+    "default_time_series_factory",
+]

cognite/extractorutils/uploader/_base.py

@@ -16,7 +16,7 @@ import logging
 import threading
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
-from typing import Any, Callable, List, Optional
+from typing import Any, Callable
 
 from arrow import Arrow
 
@@ -43,12 +43,12 @@ class AbstractUploadQueue(ABC):
     def __init__(
         self,
         cdf_client: CogniteClient,
-        post_upload_function: Optional[Callable[[List[Any]], None]] = None,
-        max_queue_size: Optional[int] = None,
-        max_upload_interval: Optional[int] = None,
+        post_upload_function: Callable[[list[Any]], None] | None = None,
+        max_queue_size: int | None = None,
+        max_upload_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,
     ):
         self.cdf_client = cdf_client
 
@@ -81,12 +81,12 @@ class AbstractUploadQueue(ABC):
 
         return None
 
-    def _post_upload(self, uploaded: List[Any]) -> None:
+    def _post_upload(self, uploaded: list[Any]) -> None:
         """
         Perform post_upload_function to uploaded data, if applicable
 
         Args:
-            uploaded: List of uploaded data
+            uploaded: list of uploaded data
         """
         if self.post_upload_function is not None:
             try:

cognite/extractorutils/uploader/assets.py

@@ -12,7 +12,8 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
-from typing import Any, Callable, List, Optional, Type
+from types import TracebackType
+from typing import Any, Callable, Type
 
 from cognite.client import CogniteClient
 from cognite.client.data_classes.assets import Asset
@@ -52,12 +53,12 @@ class AssetUploadQueue(AbstractUploadQueue):
     def __init__(
         self,
         cdf_client: CogniteClient,
-        post_upload_function: Optional[Callable[[List[Any]], None]] = None,
-        max_queue_size: Optional[int] = None,
-        max_upload_interval: Optional[int] = None,
+        post_upload_function: Callable[[list[Any]], None] | None = None,
+        max_queue_size: int | None = None,
+        max_upload_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__(
             cdf_client,
@@ -68,7 +69,7 @@ class AssetUploadQueue(AbstractUploadQueue):
             thread_name,
             cancellation_token,
         )
-        self.upload_queue: List[Asset] = []
+        self.upload_queue: list[Asset] = []
         self.assets_queued = ASSETS_UPLOADER_QUEUED
         self.assets_written = ASSETS_UPLOADER_WRITTEN
         self.queue_size = ASSETS_UPLOADER_QUEUE_SIZE
@@ -106,7 +107,7 @@ class AssetUploadQueue(AbstractUploadQueue):
                 self.cdf_client.assets.create(self.upload_queue)
             except CogniteDuplicatedError as e:
                 duplicated_ids = set([dup["externalId"] for dup in e.duplicated if "externalId" in dup])
-                failed: List[Asset] = [e for e in e.failed]
+                failed: list[Asset] = [e for e in e.failed]
                 to_create = []
                 to_update = []
                 for asset in failed:
@@ -144,7 +145,12 @@ class AssetUploadQueue(AbstractUploadQueue):
         self.start()
         return self
 
-    def __exit__(self, exc_type: Optional[Type[BaseException]], exc_val: Optional[BaseException]) -> None:
+    def __exit__(
+        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/uploader/data_modeling.py

@@ -1,5 +1,5 @@
 from types import TracebackType
-from typing import Any, Callable, List, Optional, Type
+from typing import Any, Callable, Type
 
 from cognite.client import CogniteClient
 from cognite.client.data_classes.data_modeling import EdgeApply, NodeApply
@@ -18,12 +18,12 @@ class InstanceUploadQueue(AbstractUploadQueue):
     def __init__(
         self,
         cdf_client: CogniteClient,
-        post_upload_function: Optional[Callable[[List[Any]], None]] = None,
-        max_queue_size: Optional[int] = None,
-        max_upload_interval: Optional[int] = None,
+        post_upload_function: Callable[[list[Any]], None] | None = None,
+        max_queue_size: int | None = None,
+        max_upload_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,
         auto_create_start_nodes: bool = True,
         auto_create_end_nodes: bool = True,
         auto_create_direct_relations: bool = True,
@@ -42,14 +42,14 @@ class InstanceUploadQueue(AbstractUploadQueue):
         self.auto_create_end_nodes = auto_create_end_nodes
         self.auto_create_direct_relations = auto_create_direct_relations
 
-        self.node_queue: List[NodeApply] = []
-        self.edge_queue: List[EdgeApply] = []
+        self.node_queue: list[NodeApply] = []
+        self.edge_queue: list[EdgeApply] = []
 
     def add_to_upload_queue(
         self,
         *,
-        node_data: Optional[List[NodeApply]] = None,
-        edge_data: Optional[List[EdgeApply]] = None,
+        node_data: list[NodeApply] | None = None,
+        edge_data: list[EdgeApply] | None = None,
     ) -> None:
         if node_data:
             with self.lock:
@@ -100,9 +100,9 @@ class InstanceUploadQueue(AbstractUploadQueue):
 
     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/uploader/events.py

@@ -13,7 +13,7 @@
 #  limitations under the License.
 
 from types import TracebackType
-from typing import Callable, List, Optional, Type
+from typing import Callable, Type
 
 from cognite.client import CogniteClient
 from cognite.client.data_classes import Event
@@ -52,12 +52,12 @@ class EventUploadQueue(AbstractUploadQueue):
     def __init__(
         self,
         cdf_client: CogniteClient,
-        post_upload_function: Optional[Callable[[List[Event]], None]] = None,
-        max_queue_size: Optional[int] = None,
-        max_upload_interval: Optional[int] = None,
+        post_upload_function: Callable[[list[Event]], None] | None = None,
+        max_queue_size: int | None = None,
+        max_upload_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 sets post_upload and threshold
         super().__init__(
@@ -70,7 +70,7 @@ class EventUploadQueue(AbstractUploadQueue):
             cancellation_token,
         )
 
-        self.upload_queue: List[Event] = []
+        self.upload_queue: list[Event] = []
 
         self.events_queued = EVENTS_UPLOADER_QUEUED
         self.events_written = EVENTS_UPLOADER_WRITTEN
@@ -110,7 +110,7 @@ class EventUploadQueue(AbstractUploadQueue):
                 self.cdf_client.events.create([e for e in self.upload_queue])
             except CogniteDuplicatedError as e:
                 duplicated_ids = set([dup["externalId"] for dup in e.duplicated if "externalId" in dup])
-                failed: List[Event] = [e for e in e.failed]
+                failed: list[Event] = [e for e in e.failed]
                 to_create = []
                 to_update = []
                 for evt in failed:
@@ -151,7 +151,7 @@ class EventUploadQueue(AbstractUploadQueue):
         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