Cognite-SDK v7 (#262)

* build: Set cognite-sdk verson to 7 pre-release * refactor: Support v7 SDK * build: changelog * style * Update pyproject.toml Co-authored-by: Håkon V. Treider <[email protected]> * build: bump version
cognitedata · Nov 21, 2023 · 4833a3c · 4833a3c
1 parent eeed8fa
commit 4833a3c
Show file tree

Hide file tree

Showing 5 changed files with 49 additions and 24 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -12,6 +12,12 @@ Changes are grouped as follows
 - `Fixed` for any bug fixes.
 - `Security` in case of vulnerabilities.
 
+## [6.0.0]
+
+### Changed
+
+* `cognite-sdk` to `v7`
+
 ## [5.5.1]
 
 ### Added

diff --git a/cognite/extractorutils/__init__.py b/cognite/extractorutils/__init__.py
@@ -16,5 +16,5 @@
 Cognite extractor utils is a Python package that simplifies the development of new extractors.
 """
 
-__version__ = "5.5.1"
+__version__ = "6.0.0"
 from .base import Extractor
diff --git a/cognite/extractorutils/uploader/time_series.py b/cognite/extractorutils/uploader/time_series.py
@@ -21,7 +21,12 @@
 from requests import ConnectionError
 
 from cognite.client import CogniteClient
-from cognite.client.data_classes import Sequence, SequenceData, TimeSeries
+from cognite.client.data_classes import (
+    Sequence,
+    SequenceData,
+    SequenceRows,
+    TimeSeries,
+)
 from cognite.client.exceptions import CogniteAPIError, CogniteDuplicatedError, CogniteNotFoundError
 from cognite.extractorutils.uploader._base import (
     RETRIES,
@@ -361,7 +366,7 @@ def __init__(
             thread_name,
             cancellation_token,
         )
-        self.upload_queue: Dict[EitherId, SequenceData] = {}
+        self.upload_queue: Dict[EitherId, SequenceRows] = {}
         self.sequence_metadata: Dict[EitherId, Dict[str, Union[str, int, float]]] = {}
         self.sequence_asset_external_ids: Dict[EitherId, str] = {}
         self.sequence_dataset_external_ids: Dict[EitherId, str] = {}
@@ -435,6 +440,7 @@ def add_to_upload_queue(
             List[Tuple[int, Union[int, float, str]]],
             List[Dict[str, Any]],
             SequenceData,
+            SequenceRows,
         ],
         column_external_ids: Optional[List[dict]] = None,
         id: Optional[int] = None,
@@ -459,28 +465,33 @@ def add_to_upload_queue(
 
         either_id = EitherId(id=id, external_id=external_id)
 
-        if isinstance(rows, SequenceData):
-            # Already in desired format
+        if isinstance(rows, SequenceRows):
+            # Already in the desired format
             pass
-        elif isinstance(rows, dict):
-            rows = [{"rowNumber": row_number, "values": values} for row_number, values in rows.items()]
-
-            rows = SequenceData(id=id, external_id=id, rows=rows, columns=column_external_ids)  # type: ignore
-
-        elif isinstance(rows, list):
-            if isinstance(rows[0], tuple) or isinstance(rows[0], list):
-                rows = [{"rowNumber": row_number, "values": values} for row_number, values in rows]
-
-            rows = SequenceData(id=id, external_id=id, rows=rows, columns=column_external_ids)  # type: ignore
+        elif isinstance(rows, (dict, list)):
+            rows_raw: List[Dict[str, Any]]
+            if isinstance(rows, dict):
+                rows_raw = [{"rowNumber": row_number, "values": values} for row_number, values in rows.items()]
+            elif isinstance(rows, list) and rows and isinstance(rows[0], (tuple, list)):
+                rows_raw = [{"rowNumber": row_number, "values": values} for row_number, values in rows]
+            else:
+                rows_raw = rows  # type: ignore[assignment]
+            rows = SequenceRows.load(
+                {
+                    "rows": rows_raw,
+                    "columns": column_external_ids,
+                    "id": id,
+                    "externalId": external_id,
+                }
+            )
         else:
             raise TypeError("Unsupported type for sequence rows: {}".format(type(rows)))
 
         with self.lock:
             seq = self.upload_queue.get(either_id)
             if seq is not None:
                 # Update sequence
-                seq.values.extend(rows.values)  # type: ignore  # type is list, mypy is wrong
-                seq.row_numbers.extend(rows.row_numbers)  # type: ignore
+                seq.rows.extend(rows.rows)  # type: ignore[attr-defined]
 
                 self.upload_queue[either_id] = seq
             else:
@@ -585,7 +596,7 @@ def _create_or_update(self, either_id: EitherId) -> None:
 
         # Update definition of cached sequence
         cseq = self.upload_queue[either_id]
-        cseq.columns = seq.columns
+        cseq.columns = seq.columns  # type: ignore[assignment]
 
     def _resolve_asset_ids(self) -> None:
         """

diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "cognite-extractor-utils"
-version = "5.5.1"
+version = "6.0.0"
 description = "Utilities for easier development of extractors for CDF"
 authors = ["Mathias Lohne <[email protected]>"]
 license = "Apache-2.0"
@@ -51,7 +51,7 @@ exclude = "tests/*"
 
 [tool.poetry.dependencies]
 python = "^3.8.0"
-cognite-sdk = ">=6.0, <7"
+cognite-sdk = "^7"
 prometheus-client = ">0.7.0, <=1.0.0"
 arrow = "^1.0.0"
 pyyaml = ">=5.3.0, <7"

diff --git a/tests/tests_unit/test_cdf_upload_queues.py b/tests/tests_unit/test_cdf_upload_queues.py
@@ -221,11 +221,15 @@ def post(x):
         queue.start()
 
         queue.add_to_upload_queue(
-            rows=[{"rowNumber": 1, "values": {"field": "Hello"}}], column_external_ids=[], external_id="seq-1"
+            rows=[{"rowNumber": 1, "values": ["Hello"]}],
+            column_external_ids=[{"externalId": "field", "valueType": "String"}],
+            external_id="seq-1",
         )
 
         queue.add_to_upload_queue(
-            rows=[{"rowNumber": 2, "values": {"field": "World"}}], column_external_ids=[], external_id="seq-1"
+            rows=[{"rowNumber": 2, "values": ["World"]}],
+            column_external_ids=[{"externalId": "field", "valueType": "String"}],
+            external_id="seq-1",
         )
 
         time.sleep(2.1)
@@ -249,11 +253,15 @@ def post(x):
         queue.start()
 
         queue.add_to_upload_queue(
-            rows=[{"rowNumber": 1, "values": {"field": "Hello"}}], column_external_ids=[], external_id="seq-1"
+            rows=[{"rowNumber": 1, "values": ["Hello"]}],
+            column_external_ids=[{"externalId": "field", "valueType": "String"}],
+            external_id="seq-1",
         )
 
         queue.add_to_upload_queue(
-            rows=[{"rowNumber": 2, "values": {"field": "World"}}], column_external_ids=[], external_id="seq-2"
+            rows=[{"rowNumber": 2, "values": ["World"]}],
+            column_external_ids=[{"externalId": "field", "valueType": "String"}],
+            external_id="seq-2",
         )
 
         time.sleep(2.1)