feat(records): add upsert endpoint (#2679)

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Co-authored-by: Håkon V. Treider <haakonvt@gmail.com>

Author: 3 people

cognite/client/_api/data_modeling/records.py

@@ -35,7 +35,9 @@ def _get_semaphore(self, operation: Literal["write", "delete"]) -> asyncio.Bound
         )
 
     def _records_url(self, stream_id: str, suffix: str = "") -> str:
-        return interpolate_and_url_encode("/streams/{}/records{}", stream_id, suffix)
+        # Encode only stream_id; the suffix is a literal path segment (e.g. "/upsert"),
+        # so it must not be percent-encoded.
+        return interpolate_and_url_encode("/streams/{}/records", stream_id) + suffix
 
     async def delete(
         self,
@@ -127,3 +129,57 @@ async def ingest(
             resource_path=self._records_url(stream_id),
             no_response=True,
         )
+
+    async def upsert(
+        self,
+        items: RecordWrite | Sequence[RecordWrite],
+        *,
+        stream_id: str,
+        upsert_mode: Literal["replace"] = "replace",
+    ) -> None:
+        """`Upsert records into a stream <https://api-docs.cognite.com/20230101/tag/Records/operation/upsertRecords>`_.
+
+        Creates or fully updates records. Only valid for mutable streams (returns 422 on
+        immutable). When a record with the same ``space + externalId`` already exists it is
+        fully replaced (this endpoint does not do partial property updates); otherwise it is
+        created.
+
+        Args:
+            items (RecordWrite | Sequence[RecordWrite]): One or more records to upsert.
+            stream_id (str): External ID of the stream to upsert into.
+            upsert_mode (Literal['replace']): How existing records are updated. Currently only ``"replace"`` is supported, which fully replaces the existing record. Defaults to ``"replace"``.
+
+        Examples:
+
+            Upsert a single record:
+
+                >>> from cognite.client import CogniteClient
+                >>> from cognite.client.data_classes.data_modeling.records import (
+                ...     RecordWrite,
+                ...     RecordContainerId,
+                ...     RecordSource,
+                ... )
+                >>> client = CogniteClient()
+                >>> client.data_modeling.records.upsert(
+                ...     RecordWrite(
+                ...         space="my-space",
+                ...         external_id="rec-1",
+                ...         sources=[
+                ...             RecordSource(
+                ...                 source=RecordContainerId(
+                ...                     space="my-space", external_id="my-container"
+                ...                 ),
+                ...                 properties={"temperature": 23.0},
+                ...             )
+                ...         ],
+                ...     ),
+                ...     stream_id="my-stream",
+                ... )
+        """
+        self._warning.warn()
+        item_list: list[RecordWrite] = [items] if isinstance(items, RecordWrite) else list(items)
+        await self._create_multiple(
+            items=item_list,
+            resource_path=self._records_url(stream_id, "/upsert"),
+            no_response=True,
+        )

cognite/client/_sync_api/data_modeling/records.py

@@ -45,6 +45,16 @@ def mock_ingest(httpx_mock: HTTPXMock, ingest_url_pattern: re.Pattern) -> None:
     httpx_mock.add_response(method="POST", url=ingest_url_pattern, status_code=202)
 
 
+@pytest.fixture
+def upsert_url_pattern(records_base_url: str) -> re.Pattern:
+    return re.compile(re.escape(records_base_url) + r"/upsert$")
+
+
+@pytest.fixture
+def mock_upsert(httpx_mock: HTTPXMock, upsert_url_pattern: re.Pattern) -> None:
+    httpx_mock.add_response(method="POST", url=upsert_url_pattern, status_code=202)
+
+
 @pytest.fixture
 def write_item() -> RecordWrite:
     return RecordWrite(
@@ -152,6 +162,70 @@ def test_ingest_chunks_over_1000(
         assert len(jsgz_load(requests[1].content)["items"]) == 1
 
 
+class TestRecordsAPIUpsert:
+    def test_upsert_single_posts_correct_body(
+        self,
+        cognite_client: CogniteClient,
+        httpx_mock: HTTPXMock,
+        mock_upsert: None,
+        stream_id: str,
+        write_item: RecordWrite,
+    ) -> None:
+        cognite_client.data_modeling.records.upsert(write_item, stream_id=stream_id)
+        requests = httpx_mock.get_requests()
+        assert len(requests) == 1
+        assert requests[0].url.path.endswith(f"/streams/{stream_id}/records/upsert")
+        body = jsgz_load(requests[0].content)
+        assert body == {
+            "items": [
+                {
+                    "space": "sp",
+                    "externalId": "rec-1",
+                    "sources": [
+                        {
+                            "source": {"type": "container", "space": "sp", "externalId": "container-x"},
+                            "properties": {"temp": 22.5},
+                        }
+                    ],
+                }
+            ]
+        }
+
+    def test_upsert_accepts_sequence(
+        self,
+        cognite_client: CogniteClient,
+        httpx_mock: HTTPXMock,
+        mock_upsert: None,
+        stream_id: str,
+    ) -> None:
+        items = [
+            RecordWrite(space="sp", external_id="rec-1", sources=[]),
+            RecordWrite(space="sp", external_id="rec-2", sources=[]),
+        ]
+        cognite_client.data_modeling.records.upsert(items, stream_id=stream_id)
+        body = jsgz_load(httpx_mock.get_requests()[0].content)
+        assert [item["externalId"] for item in body["items"]] == ["rec-1", "rec-2"]
+
+    def test_upsert_chunks(
+        self,
+        cognite_client: CogniteClient,
+        async_client: AsyncCogniteClient,
+        httpx_mock: HTTPXMock,
+        upsert_url_pattern: re.Pattern,
+        stream_id: str,
+        monkeypatch: pytest.MonkeyPatch,
+    ) -> None:
+        monkeypatch.setattr(async_client.data_modeling.records, "_CREATE_LIMIT", 10)
+        httpx_mock.add_response(method="POST", url=upsert_url_pattern, status_code=202)
+        httpx_mock.add_response(method="POST", url=upsert_url_pattern, status_code=202)
+        items = [RecordWrite(space="sp", external_id=f"r-{i}", sources=[]) for i in range(11)]
+        cognite_client.data_modeling.records.upsert(items, stream_id=stream_id)
+        requests = httpx_mock.get_requests()
+        assert len(requests) == 2
+        assert len(jsgz_load(requests[0].content)["items"]) == 10
+        assert len(jsgz_load(requests[1].content)["items"]) == 1
+
+
 class TestRecordDTOs:
     def test_record_write_as_id(self, write_item: RecordWrite) -> None:
         rid = write_item.as_id()