fix(#1081): prevent add records of different task (#1085)

* fix(#1081): prevent add records of different task

* test: add tests

(cherry picked from commit 45d0ce6)

src/rubrix/server/datasets/service.py

@@ -191,6 +191,7 @@ def update(
         data: UpdateDatasetRequest,
         user: User,
         workspace: Optional[str],
+        task: Optional[str] = None,
     ) -> Dataset:
         """
         Updates an existing dataset. Fields in update data are
@@ -206,13 +207,15 @@ def update(
             The current user
         workspace:
             The workspace where dataset belongs to
+        task:
+            If provided, is used to match the existing dataset
 
         Returns
         -------
             The updated dataset
         """
 
-        found = self.find_by_name(name, task=None, user=user, workspace=workspace)
+        found = self.find_by_name(name, task=task, user=user, workspace=workspace)
 
         data.tags = {**found.tags, **data.tags}
         data.metadata = {**found.metadata, **data.metadata}
@@ -257,6 +260,7 @@ def upsert(
                 data=UpdateDatasetRequest(tags=dataset.tags, metadata=dataset.metadata),
                 user=user,
                 workspace=workspace,
+                task=task,
             )
         except EntityNotFoundError:
             return self.create(

tests/functional_tests/test_log_for_text_classification.py

@@ -92,3 +92,20 @@ def test_logging_with_metadata_limits_exceeded(monkeypatch):
     expected_record.metadata["new_key"] = "value"
     with pytest.raises(BadRequestApiError):
         rubrix.log(expected_record, name=dataset)
+
+
+def test_log_with_other_task(monkeypatch):
+    mocking_client(monkeypatch, client)
+    dataset = "test_log_with_other_task"
+
+    rubrix.delete(dataset)
+    record = TextClassificationRecord(
+        inputs="The input text",
+    )
+    rubrix.log(record, name=dataset)
+
+    with pytest.raises(BadRequestApiError):
+        rubrix.log(
+            TokenClassificationRecord(text="The text", tokens=["The", "text"]),
+            name=dataset,
+        )