[REFACTOR] argilla server: moving all record validators (#5603)

# Description
<!-- Please include a summary of the changes and the related issue.
Please also include relevant motivation and context. List any
dependencies that are required for this change. -->

This PR moves all the record relationship validations (suggestions,
responses, and vectors) inside the record validator. This helps
centralize the validations and prevents objects from flushing into the
DB early.

**Type of change**
<!-- Please delete options that are not relevant. Remember to title the
PR according to the type of change -->

- Refactor (change restructuring the codebase without changing
functionality)

**How Has This Been Tested**
<!-- Please add some reference about how your feature has been tested.
-->

**Checklist**
<!-- Please go over the list and make sure you've taken everything into
account -->

- I added relevant documentation
- I followed the style guidelines of this project
- I did a self-review of my code
- I made corresponding changes to the documentation
- I confirm My changes generate no new warnings
- I have added tests that prove my fix is effective or that my feature
works
- I have added relevant notes to the CHANGELOG.md file (See
https://keepachangelog.com/)

Author: frascuchon

argilla-server/src/argilla_server/bulk/records_bulk.py

@@ -16,10 +16,10 @@
 from typing import Dict, List, Sequence, Tuple, Union
 from uuid import UUID
 
+from fastapi.encoders import jsonable_encoder
 from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.orm import selectinload
-from fastapi.encoders import jsonable_encoder
 
 from argilla_server.api.schemas.v1.records import RecordCreate, RecordUpsert
 from argilla_server.api.schemas.v1.records_bulk import (
@@ -31,18 +31,13 @@
 from argilla_server.api.schemas.v1.responses import UserResponseCreate
 from argilla_server.api.schemas.v1.suggestions import SuggestionCreate
 from argilla_server.contexts import distribution
-from argilla_server.contexts.accounts import fetch_users_by_ids_as_dict
 from argilla_server.contexts.records import (
     fetch_records_by_external_ids_as_dict,
     fetch_records_by_ids_as_dict,
 )
-from argilla_server.errors.future import UnprocessableEntityError
 from argilla_server.models import Dataset, Record, Response, Suggestion, Vector, VectorSettings
 from argilla_server.search_engine import SearchEngine
 from argilla_server.validators.records import RecordsBulkCreateValidator, RecordsBulkUpsertValidator
-from argilla_server.validators.responses import ResponseCreateValidator
-from argilla_server.validators.suggestions import SuggestionCreateValidator
-from argilla_server.validators.vectors import VectorValidator
 
 
 class CreateRecordsBulk:
@@ -84,30 +79,16 @@ async def _upsert_records_relationships(self, records: List[Record], records_cre
         # https://github.com/sqlalchemy/sqlalchemy/discussions/9312
 
         await self._upsert_records_suggestions(records_and_suggestions)
-        await self._upsert_records_responses(records_and_responses)
         await self._upsert_records_vectors(records_and_vectors)
+        await self._upsert_records_responses(records_and_responses)
 
     async def _upsert_records_suggestions(
         self, records_and_suggestions: List[Tuple[Record, List[SuggestionCreate]]]
     ) -> List[Suggestion]:
         upsert_many_suggestions = []
         for idx, (record, suggestions) in enumerate(records_and_suggestions):
-            try:
-                for suggestion_create in suggestions or []:
-                    question = record.dataset.question_by_id(suggestion_create.question_id)
-                    if question is None:
-                        raise ValueError(f"question with question_id={suggestion_create.question_id} does not exist")
-
-                    try:
-                        SuggestionCreateValidator.validate(suggestion_create, question.parsed_settings, record)
-                        upsert_many_suggestions.append(dict(**suggestion_create.dict(), record_id=record.id))
-                    except (UnprocessableEntityError, ValueError) as ex:
-                        raise ValueError(f"suggestion for question name={question.name} is not valid: {ex}")
-
-            except (UnprocessableEntityError, ValueError) as ex:
-                raise UnprocessableEntityError(
-                    f"Record at position {idx} does not have valid suggestions because {ex}"
-                ) from ex
+            for suggestion_create in suggestions or []:
+                upsert_many_suggestions.append(dict(**suggestion_create.dict(), record_id=record.id))
 
         if not upsert_many_suggestions:
             return []
@@ -122,22 +103,10 @@ async def _upsert_records_suggestions(
     async def _upsert_records_responses(
         self, records_and_responses: List[Tuple[Record, List[UserResponseCreate]]]
     ) -> List[Response]:
-        user_ids = [response.user_id for _, responses in records_and_responses for response in responses or []]
-        users_by_id = await fetch_users_by_ids_as_dict(self._db, user_ids)
-
         upsert_many_responses = []
         for idx, (record, responses) in enumerate(records_and_responses):
-            try:
-                for response_create in responses or []:
-                    if response_create.user_id not in users_by_id:
-                        raise ValueError(f"user with id {response_create.user_id} not found")
-
-                    ResponseCreateValidator.validate(response_create, record)
-                    upsert_many_responses.append(dict(**response_create.dict(), record_id=record.id))
-            except (UnprocessableEntityError, ValueError) as ex:
-                raise UnprocessableEntityError(
-                    f"Record at position {idx} does not have valid responses because {ex}"
-                ) from ex
+            for response_create in responses or []:
+                upsert_many_responses.append(dict(**response_create.dict(), record_id=record.id))
 
         if not upsert_many_responses:
             return []
@@ -154,18 +123,11 @@ async def _upsert_records_vectors(
     ) -> List[Vector]:
         upsert_many_vectors = []
         for idx, (record, vectors) in enumerate(records_and_vectors):
-            try:
-                for name, value in (vectors or {}).items():
-                    settings = _get_vector_settings_by_name(record.dataset, name)
-                    if not settings:
-                        raise ValueError(f"vector with name={name} does not exist for dataset_id={record.dataset.id}")
-
-                    VectorValidator.validate(value, settings)
-                    upsert_many_vectors.append(dict(value=value, record_id=record.id, vector_settings_id=settings.id))
-            except (UnprocessableEntityError, ValueError) as ex:
-                raise UnprocessableEntityError(
-                    f"Record at position {idx} does not have valid vectors because {ex}"
-                ) from ex
+            dataset = record.dataset
+
+            for name, value in (vectors or {}).items():
+                settings = dataset.vector_settings_by_name(name)
+                upsert_many_vectors.append(dict(value=value, record_id=record.id, vector_settings_id=settings.id))
 
         if not upsert_many_vectors:
             return []
@@ -186,7 +148,7 @@ async def upsert_records_bulk(self, dataset: Dataset, bulk_upsert: RecordsBulkUp
         found_records = await self._fetch_existing_dataset_records(dataset, bulk_upsert.items)
         # found_records is passed to the validator to avoid querying the database again, but ideally, it should be
         # computed inside the validator
-        RecordsBulkUpsertValidator.validate(bulk_upsert, dataset, found_records)
+        await RecordsBulkUpsertValidator.validate(bulk_upsert, dataset, found_records)
 
         records = []
         async with self._db.begin_nested():
@@ -246,9 +208,3 @@ async def _preload_records_relationships_before_index(db: "AsyncSession", record
             selectinload(Record.vectors),
         )
     )
-
-
-def _get_vector_settings_by_name(dataset: Dataset, name: str) -> Union[VectorSettings, None]:
-    for vector_settings in dataset.vectors_settings:
-        if vector_settings.name == name:
-            return vector_settings

argilla-server/src/argilla_server/contexts/accounts.py

@@ -12,7 +12,7 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 import secrets
-from typing import Dict, Iterable, List, Sequence, Union
+from typing import Iterable, List, Sequence, Union
 from uuid import UUID
 
 from passlib.context import CryptContext
@@ -193,8 +193,3 @@ def generate_user_token(user: User) -> str:
             role=user.role,
         ),
     )
-
-
-async def fetch_users_by_ids_as_dict(db: "AsyncSession", user_ids: List[UUID]) -> Dict[UUID, User]:
-    users = await list_users_by_ids(db, set(user_ids))
-    return {user.id: user for user in users}

argilla-server/src/argilla_server/models/base.py

@@ -11,10 +11,10 @@
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
-
+from typing import Optional
 from uuid import UUID, uuid4
 
-from sqlalchemy.ext.asyncio import AsyncAttrs
+from sqlalchemy.ext.asyncio import AsyncAttrs, async_object_session, AsyncSession
 from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
 
 from argilla_server.models.mixins import CRUDMixin, TimestampMixin
@@ -32,3 +32,7 @@ class DatabaseModel(DeclarativeBase, AsyncAttrs, CRUDMixin, TimestampMixin):
 
     def is_relationship_loaded(self, relationship: str) -> bool:
         return relationship in self.__dict__
+
+    @property
+    def current_async_session(self) -> Optional[AsyncSession]:
+        return async_object_session(self)

argilla-server/src/argilla_server/models/database.py

@@ -28,7 +28,6 @@
     sql,
 )
 from sqlalchemy.engine.default import DefaultExecutionContext
-from sqlalchemy.ext.asyncio import async_object_session
 from sqlalchemy.ext.mutable import MutableDict, MutableList
 from sqlalchemy.orm import Mapped, mapped_column, relationship
 
@@ -509,7 +508,7 @@ def is_annotator(self):
     async def is_member(self, workspace_id: UUID) -> bool:
         # TODO: Change query to use exists may improve performance
         return (
-            await WorkspaceUser.get_by(async_object_session(self), workspace_id=workspace_id, user_id=self.id)
+            await WorkspaceUser.get_by(self.current_async_session, workspace_id=workspace_id, user_id=self.id)
             is not None
         )