cognite-extractor-utils 5.1.0__py3-none-any.whl → 5.2.0__py3-none-any.whl

cognite/extractorutils/statestore.py

@@ -89,14 +89,13 @@ import json
 import logging
 import threading
 from abc import ABC, abstractmethod
-from decimal import Decimal
-from threading import Lock
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from types import TracebackType
+from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union
 
-from cognite.client import CogniteClient
-from cognite.client.exceptions import CogniteAPIError
 from requests.exceptions import ConnectionError
 
+from cognite.client import CogniteClient
+from cognite.client.exceptions import CogniteAPIError
 from cognite.extractorutils.uploader import DataPointList
 
 from ._inner_util import _DecimalDecoder, _DecimalEncoder, _resolve_log_level
@@ -140,8 +139,6 @@ class AbstractStateStore(ABC):
 
         self._deleted: List[str] = []
 
-        self.lock = Lock()
-
     def start(self) -> None:
         """
         Start saving state periodically if save_interval is set.
@@ -203,12 +200,12 @@ class AbstractStateStore(ABC):
         """
         with self.lock:
             if isinstance(external_id, list):
-                l = []
+                states = []
                 for e in external_id:
                     state = self._local_state.get(e, {})
-                    l.append((state.get("low"), state.get("high")))
+                    states.append((state.get("low"), state.get("high")))
 
-                return l
+                return states
 
             else:
                 state = self._local_state.get(external_id, {})
@@ -263,14 +260,14 @@ class AbstractStateStore(ABC):
             A function that expands the current states with the values given
         """
 
-        def callback(uploaded_points: List[Dict[str, Union[str, DataPointList]]]):
+        def callback(uploaded_points: List[Dict[str, Union[str, DataPointList]]]) -> None:
             for time_series in uploaded_points:
                 # Use CDF timestamps
                 data_points = time_series["datapoints"]
                 if data_points:
                     high = max(data_points)[0]
                     low = min(data_points)[0]
-                    external_id = time_series["externalId"]
+                    external_id: str = time_series["externalId"]  # type: ignore  # known to be str from where we set it
                     self.expand_state(external_id, low, high)
 
         return callback
@@ -301,7 +298,7 @@ class AbstractStateStore(ABC):
         return False
 
     def __getitem__(self, external_id: str) -> Tuple[Any, Any]:
-        return self.get_state(external_id)
+        return self.get_state(external_id)  # type: ignore  # will not be list if input is single str
 
     def __setitem__(self, key: str, value: Tuple[Any, Any]) -> None:
         self.set_state(external_id=key, low=value[0], high=value[1])
@@ -382,7 +379,8 @@ class RawStateStore(AbstractStateStore):
         if self._initialized and not force:
             return
 
-        rows = self._cdf_client.raw.rows.list(db_name=self.database, table_name=self.table, limit=None)
+        # ignore type since list _is_ optional, sdk types are wrong
+        rows = self._cdf_client.raw.rows.list(db_name=self.database, table_name=self.table, limit=None)  # type: ignore
 
         with self.lock:
             self._local_state.clear()
@@ -422,7 +420,9 @@ class RawStateStore(AbstractStateStore):
         self.start()
         return self
 
-    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+    def __exit__(
+        self, exc_type: Optional[Type[BaseException]], exc_val: Optional[BaseException], exc_tb: Optional[TracebackType]
+    ) -> None:
         """
         Wraps around stop method, for use as context manager
 
@@ -500,7 +500,9 @@ class LocalStateStore(AbstractStateStore):
         self.start()
         return self
 
-    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+    def __exit__(
+        self, exc_type: Optional[Type[BaseException]], exc_val: Optional[BaseException], exc_tb: Optional[TracebackType]
+    ) -> None:
         """
         Wraps around stop method, for use as context manager
 
@@ -517,7 +519,7 @@ class NoStateStore(AbstractStateStore):
     A state store that only keeps states in memory and never stores or initializes from external sources.
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__()
 
     def initialize(self, force: bool = False) -> None:

cognite/extractorutils/uploader/_base.py

@@ -19,8 +19,8 @@ from dataclasses import dataclass
 from typing import Any, Callable, List, Optional
 
 from arrow import Arrow
-from cognite.client import CogniteClient
 
+from cognite.client import CogniteClient
 from cognite.extractorutils._inner_util import _resolve_log_level
 
 
@@ -88,16 +88,9 @@ class AbstractUploadQueue(ABC):
         if self.post_upload_function is not None:
             try:
                 self.post_upload_function(uploaded)
-            except Exception as e:
+            except Exception:
                 logging.getLogger(__name__).exception("Error during upload callback")
 
-    @abstractmethod
-    def add_to_upload_queue(self, *args) -> None:
-        """
-        Adds an element to the upload queue. The queue will be uploaded if the queue byte size is larger than the
-        threshold specified in the config.
-        """
-
     @abstractmethod
     def upload(self) -> None:
         """

cognite/extractorutils/uploader/events.py

@@ -13,14 +13,15 @@
 #  limitations under the License.
 
 import threading
-from typing import Callable, List, Optional
+from types import TracebackType
+from typing import Callable, List, Optional, Type
 
 import arrow
-from cognite.client import CogniteClient
-from cognite.client.data_classes import Event
-from cognite.client.exceptions import CogniteAPIError
 from requests import ConnectionError
 
+from cognite.client import CogniteClient
+from cognite.client.data_classes import Event
+from cognite.client.exceptions import CogniteAPIError, CogniteDuplicatedError
 from cognite.extractorutils.uploader._base import (
     RETRIES,
     RETRY_BACKOFF_FACTOR,
@@ -131,8 +132,23 @@ class EventUploadQueue(AbstractUploadQueue):
         max_delay=RETRY_MAX_DELAY,
         backoff=RETRY_BACKOFF_FACTOR,
     )
-    def _upload_batch(self):
-        self.cdf_client.events.create([e for e in self.upload_queue])
+    def _upload_batch(self) -> None:
+        try:
+            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]
+            to_create = []
+            to_update = []
+            for evt in failed:
+                if evt.external_id is not None and evt.external_id in duplicated_ids:
+                    to_update.append(evt)
+                else:
+                    to_create.append(evt)
+            if to_create:
+                self.cdf_client.events.create(to_create)
+            if to_update:
+                self.cdf_client.events.update(to_update)
 
     def __enter__(self) -> "EventUploadQueue":
         """
@@ -144,7 +160,9 @@ class EventUploadQueue(AbstractUploadQueue):
         self.start()
         return self
 
-    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+    def __exit__(
+        self, exc_type: Optional[Type[BaseException]], exc_val: Optional[BaseException], exc_tb: Optional[TracebackType]
+    ) -> None:
         """
         Wraps around stop method, for use as context manager
 

cognite/extractorutils/uploader/files.py

@@ -15,14 +15,15 @@
 import threading
 from concurrent.futures import ThreadPoolExecutor
 from os import PathLike
-from typing import Any, Callable, List, Optional, Tuple, Union
+from types import TracebackType
+from typing import Any, Callable, List, Optional, Tuple, Type, Union
 
 import arrow
+from requests import ConnectionError
+
 from cognite.client import CogniteClient
 from cognite.client.data_classes import Event, FileMetadata
 from cognite.client.exceptions import CogniteAPIError
-from requests import ConnectionError
-
 from cognite.extractorutils.uploader._base import (
     RETRIES,
     RETRY_BACKOFF_FACTOR,
@@ -89,7 +90,7 @@ class FileUploadQueue(AbstractUploadQueue):
         self.latency = FILES_UPLOADER_LATENCY
         self.latency_zero_point = arrow.utcnow()
 
-    def add_to_upload_queue(self, file_meta: FileMetadata, file_name: Union[str, PathLike] = None) -> None:
+    def add_to_upload_queue(self, file_meta: FileMetadata, 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
         specified in the __init__.
@@ -141,14 +142,14 @@ class FileUploadQueue(AbstractUploadQueue):
         max_delay=RETRY_MAX_DELAY,
         backoff=RETRY_BACKOFF_FACTOR,
     )
-    def _upload_single(self, index, file_name, file_meta):
+    def _upload_single(self, index: int, file_name: Union[str, PathLike], file_meta: FileMetadata) -> None:
         # Upload file
-        file_meta = self.cdf_client.files.upload(file_name, overwrite=self.overwrite_existing, **file_meta.dump())
+        file_meta = self.cdf_client.files.upload(str(file_name), overwrite=self.overwrite_existing, **file_meta.dump())  # type: ignore
 
         # Update meta-object in queue
         self.upload_queue[index] = (file_meta, file_name)
 
-    def _upload_batch(self):
+    def _upload_batch(self) -> None:
         # Concurrently execute file-uploads
 
         with ThreadPoolExecutor(self.cdf_client.config.max_workers) as pool:
@@ -165,7 +166,9 @@ class FileUploadQueue(AbstractUploadQueue):
         self.start()
         return self
 
-    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+    def __exit__(
+        self, exc_type: Optional[Type[BaseException]], exc_val: Optional[BaseException], exc_tb: Optional[TracebackType]
+    ) -> None:
         """
         Wraps around stop method, for use as context manager
 
@@ -212,7 +215,7 @@ class BytesUploadQueue(AbstractUploadQueue):
         thread_name: Optional[str] = None,
         overwrite_existing: bool = False,
         cancellation_token: threading.Event = threading.Event(),
-    ):
+    ) -> None:
         super().__init__(
             cdf_client,
             post_upload_function,
@@ -277,7 +280,7 @@ class BytesUploadQueue(AbstractUploadQueue):
             self.logger.info(f"Uploaded {self.upload_queue_size} files")
             self.queue_size.set(self.upload_queue_size)
 
-    def _upload_batch(self):
+    def _upload_batch(self) -> None:
         # Concurrently execute bytes-uploads
         with ThreadPoolExecutor(self.cdf_client.config.max_workers) as pool:
             for i, (frame, metadata) in enumerate(self.upload_queue):
@@ -290,7 +293,7 @@ class BytesUploadQueue(AbstractUploadQueue):
         max_delay=RETRY_MAX_DELAY,
         backoff=RETRY_BACKOFF_FACTOR,
     )
-    def _upload_single(self, index: int, content: bytes, metadata: FileMetadata):
+    def _upload_single(self, index: int, content: bytes, metadata: FileMetadata) -> None:
         # Upload object
         file_meta_data: FileMetadata = self.cdf_client.files.upload_bytes(
             content, overwrite=self.overwrite_existing, **metadata.dump()
@@ -309,7 +312,9 @@ class BytesUploadQueue(AbstractUploadQueue):
         self.start()
         return self
 
-    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+    def __exit__(
+        self, exc_type: Optional[Type[BaseException]], exc_val: Optional[BaseException], exc_tb: Optional[TracebackType]
+    ) -> None:
         """
         Wraps around stop method, for use as context manager
 

cognite/extractorutils/uploader/raw.py

@@ -13,15 +13,16 @@
 #  limitations under the License.
 
 import threading
-from typing import Any, Callable, Dict, List, Optional
+from types import TracebackType
+from typing import Any, Callable, Dict, List, Optional, Type
 
 import arrow
 from arrow import Arrow
+from requests import ConnectionError
+
 from cognite.client import CogniteClient
 from cognite.client.data_classes import Row
 from cognite.client.exceptions import CogniteAPIError, CogniteReadTimeout
-from requests import ConnectionError
-
 from cognite.extractorutils.uploader._base import (
     RETRIES,
     RETRY_BACKOFF_FACTOR,
@@ -75,7 +76,7 @@ class RawUploadQueue(AbstractUploadQueue):
             thread_name,
             cancellation_token,
         )
-        self.upload_queue: Dict[str, Dict[str, List[TimestampedObject]]] = dict()
+        self.upload_queue: Dict[str, Dict[str, List[TimestampedObject]]] = {}
 
         # It is a hack since Prometheus client registers metrics on object creation, so object has to be created once
         self.rows_queued = RAW_UPLOADER_ROWS_QUEUED
@@ -97,7 +98,7 @@ class RawUploadQueue(AbstractUploadQueue):
         with self.lock:
             # Ensure that the dicts has correct keys
             if database not in self.upload_queue:
-                self.upload_queue[database] = dict()
+                self.upload_queue[database] = {}
             if table not in self.upload_queue[database]:
                 self.upload_queue[database][table] = []
 
@@ -152,7 +153,7 @@ class RawUploadQueue(AbstractUploadQueue):
         max_delay=RETRY_MAX_DELAY,
         backoff=RETRY_BACKOFF_FACTOR,
     )
-    def _upload_batch(self, database: str, table: str, patch: List[Row]):
+    def _upload_batch(self, database: str, table: str, patch: List[Row]) -> None:
         # Upload
         self.cdf_client.raw.rows.insert(db_name=database, table_name=table, row=patch, ensure_parent=True)
 
@@ -166,7 +167,9 @@ class RawUploadQueue(AbstractUploadQueue):
         self.start()
         return self
 
-    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+    def __exit__(
+        self, exc_type: Optional[Type[BaseException]], exc_val: Optional[BaseException], exc_tb: Optional[TracebackType]
+    ) -> None:
         """
         Wraps around stop method, for use as context manager