Add CDM Timeseries Upload Queue

cognite/extractorutils/uploader/time_series.py

@@ -16,7 +16,7 @@
 from collections.abc import Callable
 from datetime import datetime
 from types import TracebackType
-from typing import Any, TypeVar
+from typing import Any, Generic, TypeVar
 
 from cognite.client import CogniteClient
 from cognite.client.data_classes import (
@@ -62,6 +62,7 @@
 DataPointList = list[DataPoint]
 
 TQueue = TypeVar("TQueue", bound="BaseTimeSeriesUploadQueue")
+IdType = TypeVar("IdType", EitherId, NodeId)
 
 
 def default_time_series_factory(external_id: str, datapoints: DataPointList) -> TimeSeries:
@@ -83,7 +84,13 @@ def default_time_series_factory(external_id: str, datapoints: DataPointList) ->
     return TimeSeries(external_id=external_id, is_string=is_string)
 
 
-class BaseTimeSeriesUploadQueue(AbstractUploadQueue):
+def _apply_cognite_timeseries(cdf_client: CogniteClient, timeseries_apply: CogniteExtractorTimeSeriesApply) -> NodeId:
+    instance_result = cdf_client.data_modeling.instances.apply(timeseries_apply)
+    node = instance_result.nodes[0]
+    return node.as_id()
+
+
+class BaseTimeSeriesUploadQueue(AbstractUploadQueue, Generic[IdType]):
     """
     Abstract base upload queue for time series
 
@@ -120,7 +127,7 @@ def __init__(
             cancellation_token,
         )
 
-        self.upload_queue: dict[EitherId, DataPointList] = {}
+        self.upload_queue: dict[IdType, DataPointList] = {}
 
         self.points_queued = TIMESERIES_UPLOADER_POINTS_QUEUED
         self.points_written = TIMESERIES_UPLOADER_POINTS_WRITTEN
@@ -204,7 +211,7 @@ def __enter__(self: TQueue) -> TQueue:
         return self
 
 
-class TimeSeriesUploadQueue(BaseTimeSeriesUploadQueue):
+class TimeSeriesUploadQueue(BaseTimeSeriesUploadQueue[EitherId]):
     """
     Upload queue for time series
 
@@ -388,7 +395,7 @@ def _upload_batch(upload_this: list[dict], retries: int = 5) -> list[dict]:
             self.queue_size.set(self.upload_queue_size)
 
 
-class CDMTimeSeriesUploadQueue(BaseTimeSeriesUploadQueue):
+class CDMTimeSeriesUploadQueue(BaseTimeSeriesUploadQueue[NodeId]):
     """
     Upload queue for CDM time series
 
@@ -412,6 +419,7 @@ def __init__(
         max_upload_interval: int | None = None,
         trigger_log_level: str = "DEBUG",
         thread_name: str | None = None,
+        create_missing: Callable[[CogniteClient, CogniteExtractorTimeSeriesApply], NodeId] | bool = False,
         cancellation_token: CancellationToken | None = None,
     ):
         super().__init__(
@@ -424,10 +432,15 @@ def __init__(
             cancellation_token,
         )
 
-    def _apply_cognite_timeseries(self, timeseries_apply: CogniteExtractorTimeSeriesApply) -> NodeId:
-        instance_result = self.cdf_client.data_modeling.instances.apply(timeseries_apply)
-        node = instance_result.nodes[0]
-        return node.as_id()
+        self.missing_factory: Callable[[CogniteClient, CogniteExtractorTimeSeriesApply], NodeId]
+        self.timeseries_apply_dict: dict[NodeId, CogniteExtractorTimeSeriesApply] = {}
+
+        if isinstance(create_missing, bool):
+            self.create_missing = create_missing
+            self.missing_factory = _apply_cognite_timeseries
+        else:
+            self.create_missing = True
+            self.missing_factory = create_missing
 
     def add_to_upload_queue(
         self, *, timeseries_apply: CogniteExtractorTimeSeriesApply, datapoints: DataPointList | None = None
@@ -442,14 +455,16 @@ def add_to_upload_queue(
         """
         datapoints = self._sanitize_datapoints(datapoints)
 
-        instance_id = self._apply_cognite_timeseries(timeseries_apply)
-        either_id = EitherId(instance_id=instance_id)
+        instance_id = NodeId(timeseries_apply.space, timeseries_apply.external_id)
 
         with self.lock:
-            if either_id not in self.upload_queue:
-                self.upload_queue[either_id] = []
+            if instance_id not in self.upload_queue:
+                self.upload_queue[instance_id] = []
 
-            self.upload_queue[either_id].extend(datapoints)
+            if instance_id not in self.timeseries_apply_dict:
+                self.timeseries_apply_dict[instance_id] = timeseries_apply
+
+            self.upload_queue[instance_id].extend(datapoints)
             self.points_queued.inc(len(datapoints))
             self.upload_queue_size += len(datapoints)
             self.queue_size.set(self.upload_queue_size)
@@ -469,9 +484,61 @@ def upload(self) -> None:
             max_delay=RETRY_MAX_DELAY,
             backoff=RETRY_BACKOFF_FACTOR,
         )
-        def _upload_batch(upload_this: list[dict]) -> list[dict]:
-            if len(upload_this) > 0:
+        def _upload_batch(upload_this: list[dict], retries: int = 5) -> list[dict]:
+            if len(upload_this) == 0:
+                return upload_this
+
+            try:
                 self.cdf_client.time_series.data.insert_multiple(upload_this)
+            except CogniteNotFoundError as ex:
+                if not retries:
+                    raise ex
+
+                if not self.create_missing:
+                    self.logger.error("Could not upload data points to %s: %s", str(ex.not_found), str(ex))
+
+                # Get IDs of time series that exists, but failed because of the non-existing time series
+                retry_these = [
+                    NodeId(id_dict["instanceId"]["space"], id_dict["instanceId"]["externalId"])
+                    for id_dict in ex.failed
+                    if id_dict not in ex.not_found
+                ]
+
+                if self.create_missing:
+                    # Get the time series that can be created
+                    create_these_ids = set(
+                        [
+                            NodeId(id_dict["instanceId"]["space"], id_dict["instanceId"]["externalId"])
+                            for id_dict in ex.not_found
+                        ]
+                    )
+                    self.logger.info(f"Creating {len(create_these_ids)} time series")
+                    to_create: list[NodeId] = [
+                        self.missing_factory(self.cdf_client, self.timeseries_apply_dict[nodeId])
+                        for nodeId in create_these_ids
+                    ]
+
+                    retry_these.extend([nodeId for nodeId in to_create])
+
+                    if len(ex.not_found) != len(create_these_ids):
+                        missing = [
+                            id_dict
+                            for id_dict in ex.not_found
+                            if NodeId(id_dict["instanceId"]["space"], id_dict["instanceId"]["externalId"])
+                            not in retry_these
+                        ]
+                        missing_num = len(ex.not_found) - len(create_these_ids)
+                        self.logger.error(
+                            f"{missing_num} time series not found, and could not be created automatically: "
+                            + str(missing)
+                            + " Data will be dropped"
+                        )
+
+                # Remove entries with non-existing time series from upload queue
+                upload_this = [entry for entry in upload_this if entry["instanceId"] in retry_these]
+
+                # Upload remaining
+                _upload_batch(upload_this, retries - 1)
 
             return upload_this
 
@@ -481,13 +548,13 @@ def _upload_batch(upload_this: list[dict]) -> list[dict]:
         with self.lock:
             upload_this = _upload_batch(
                 [
-                    {either_id.type(): either_id.content(), "datapoints": list(datapoints)}
-                    for either_id, datapoints in self.upload_queue.items()
+                    {"instanceId": instance_id, "datapoints": list(datapoints)}
+                    for instance_id, datapoints in self.upload_queue.items()
                     if len(datapoints) > 0
                 ]
             )
 
-            for _either_id, datapoints in self.upload_queue.items():
+            for _instance_id, datapoints in self.upload_queue.items():
                 self.points_written.inc(len(datapoints))
 
             try:

cognite/extractorutils/util.py

@@ -32,7 +32,6 @@
 
 from cognite.client import CogniteClient
 from cognite.client.data_classes import Asset, ExtractionPipelineRun, TimeSeries
-from cognite.client.data_classes.data_modeling import NodeId
 from cognite.client.exceptions import CogniteAPIError, CogniteException, CogniteFileUploadError, CogniteNotFoundError
 from cognite.extractorutils.threading import CancellationToken
 
@@ -80,41 +79,35 @@ def ensure_assets(cdf_client: CogniteClient, assets: Iterable[Asset]) -> None:
 
 class EitherId:
     """
-    Class representing an ID in CDF, which can either be an external ID, internal ID or instance ID.
-    An EitherId can only hold one ID type.
+    Class representing an ID in CDF, which can either be an external or internal ID. An EitherId can only hold one ID
+    type, not both.
 
     Args:
         id: Internal ID
         external_id: external ID. It can be `external_id` or `externalId`
-        instance_id: Instance ID. It can be `instance_id` or `instanceId`
 
     Raises:
-        TypeError: If none or more than one of the id types are set.
+        TypeError: If none of both of id types are set.
     """
 
-    def __init__(self, **kwargs: int | str | NodeId | None):
+    def __init__(self, **kwargs: int | str | None):
         internal_id = kwargs.get("id")
         external_id = kwargs.get("externalId") or kwargs.get("external_id")
-        instance_id = kwargs.get("instanceId") or kwargs.get("instance_id")
 
-        identifiers = [internal_id, external_id, instance_id]
-        provided_ids = [id_val for id_val in identifiers if id_val is not None]
+        if internal_id is None and external_id is None:
+            raise TypeError("Either id or external_id must be set")
 
-        if not provided_ids:
-            raise TypeError("Either id, external_id, or instance_id must be set")
-        if len(provided_ids) > 1:
-            raise TypeError("Only one of id, external_id, or instance_id can be set")
+        if internal_id is not None and external_id is not None:
+            raise TypeError("Only one of id and external_id can be set")
 
         if internal_id is not None and not isinstance(internal_id, int):
             raise TypeError("Internal IDs must be integers")
+
         if external_id is not None and not isinstance(external_id, str):
             raise TypeError("External IDs must be strings")
-        if instance_id is not None and not isinstance(instance_id, NodeId):
-            raise TypeError("Instance IDs must be NodeId objects")
 
         self.internal_id: int | None = internal_id
         self.external_id: str | None = external_id
-        self.instance_id: NodeId | None = instance_id
 
     def type(self) -> str:
         """
@@ -123,21 +116,16 @@ def type(self) -> str:
         Returns:
             'id' if the EitherId represents an internal ID, 'externalId' if the EitherId represents an external ID
         """
-        if self.internal_id is not None:
-            return "id"
-        elif self.instance_id is not None:
-            return "instanceId"
-        else:
-            return "externalId"
+        return "id" if self.internal_id is not None else "externalId"
 
-    def content(self) -> int | str | NodeId:
+    def content(self) -> int | str:
         """
         Get the value of the ID
 
         Returns:
             The ID
         """
-        return self.internal_id or self.external_id or self.instance_id  # type: ignore
+        return self.internal_id or self.external_id  # type: ignore  # checked to be not None in init
 
     def __eq__(self, other: Any) -> bool:
         """
@@ -152,11 +140,7 @@ def __eq__(self, other: Any) -> bool:
         if not isinstance(other, EitherId):
             return False
 
-        return (
-            self.internal_id == other.internal_id
-            and self.external_id == other.external_id
-            and self.instance_id == other.instance_id
-        )
+        return self.internal_id == other.internal_id and self.external_id == other.external_id
 
     def __hash__(self) -> int:
         """
@@ -165,7 +149,7 @@ def __hash__(self) -> int:
         Returns:
             Hash code of ID
         """
-        return hash((self.internal_id, self.external_id, self.instance_id))
+        return hash((self.internal_id, self.external_id))
 
     def __str__(self) -> str:
         """