[PERF] argilla server: Reduce general transaction time (#5609)

# 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. -->

- Remove unnecessary extra SAVEPOINTS
- Move out of the commit all search index operations
- ~~the refresh parameter for es can be set to `wait_for` instead of
`true` since index operations are out of transactions. This will help
with the es performance.~~

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

- Bug fix (non-breaking change which fixes an issue)
- New feature (non-breaking change which adds functionality)
- Breaking change (fix or feature that would cause existing
functionality to not work as expected)
- Refactor (change restructuring the codebase without changing
functionality)
- Improvement (change adding some improvement to an existing
functionality)
- Documentation update

**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/pdm.lock

@@ -30,8 +30,8 @@ dependencies = [
     "backoff>=1.11.1",
     # Database dependencies
     "alembic ~= 1.9.0",
-    "SQLAlchemy == 2.0.31",
-    "greenlet >= 2.0.0",
+    "SQLAlchemy == 2.0.35",
+    "greenlet ~= 3.1.0",
     # Async SQLite
     "aiosqlite == 0.20.0",
     # metrics

argilla-server/pyproject.toml

@@ -48,27 +48,26 @@ def __init__(self, db: AsyncSession, search_engine: SearchEngine):
     async def create_records_bulk(self, dataset: Dataset, bulk_create: RecordsBulkCreate) -> RecordsBulk:
         await RecordsBulkCreateValidator.validate(self._db, bulk_create, dataset)
 
-        async with self._db.begin_nested():
-            records = [
-                Record(
-                    fields=jsonable_encoder(record_create.fields),
-                    metadata_=record_create.metadata,
-                    external_id=record_create.external_id,
-                    dataset_id=dataset.id,
-                )
-                for record_create in bulk_create.items
-            ]
-
-            self._db.add_all(records)
-            await self._db.flush(records)
-
-            await self._upsert_records_relationships(records, bulk_create.items)
-            await _preload_records_relationships_before_index(self._db, records)
-            await distribution.unsafe_update_records_status(self._db, records)
-            await self._search_engine.index_records(dataset, records)
+        records = [
+            Record(
+                fields=jsonable_encoder(record_create.fields),
+                metadata_=record_create.metadata,
+                external_id=record_create.external_id,
+                dataset_id=dataset.id,
+            )
+            for record_create in bulk_create.items
+        ]
+
+        self._db.add_all(records)
+        await self._db.flush(records)
+        await self._upsert_records_relationships(records, bulk_create.items)
+        await distribution.unsafe_update_records_status(self._db, records)
 
         await self._db.commit()
 
+        await _preload_records_relationships_before_index(self._db, records)
+        await self._search_engine.index_records(dataset, records)
+
         return RecordsBulk(items=records)
 
     async def _upsert_records_relationships(self, records: List[Record], records_create: List[RecordCreate]) -> None:
@@ -139,7 +138,8 @@ async def _upsert_records_vectors(
             autocommit=False,
         )
 
-    def _metadata_is_set(self, record_create: RecordCreate) -> bool:
+    @classmethod
+    def _metadata_is_set(cls, record_create: RecordCreate) -> bool:
         return "metadata" in record_create.__fields_set__
 
 
@@ -151,32 +151,32 @@ async def upsert_records_bulk(self, dataset: Dataset, bulk_upsert: RecordsBulkUp
         await RecordsBulkUpsertValidator.validate(bulk_upsert, dataset, found_records)
 
         records = []
-        async with self._db.begin_nested():
-            for record_upsert in bulk_upsert.items:
-                record = found_records.get(record_upsert.id) or found_records.get(record_upsert.external_id)
-                if not record:
-                    record = Record(
-                        fields=jsonable_encoder(record_upsert.fields),
-                        metadata_=record_upsert.metadata,
-                        external_id=record_upsert.external_id,
-                        dataset_id=dataset.id,
-                    )
-                elif self._metadata_is_set(record_upsert):
-                    record.metadata_ = record_upsert.metadata
-                    record.updated_at = datetime.utcnow()
-
-                records.append(record)
-
-            self._db.add_all(records)
-            await self._db.flush(records)
-
-            await self._upsert_records_relationships(records, bulk_upsert.items)
-            await _preload_records_relationships_before_index(self._db, records)
-            await distribution.unsafe_update_records_status(self._db, records)
-            await self._search_engine.index_records(dataset, records)
+
+        for record_upsert in bulk_upsert.items:
+            record = found_records.get(record_upsert.id) or found_records.get(record_upsert.external_id)
+            if not record:
+                record = Record(
+                    fields=jsonable_encoder(record_upsert.fields),
+                    metadata_=record_upsert.metadata,
+                    external_id=record_upsert.external_id,
+                    dataset_id=dataset.id,
+                )
+            elif self._metadata_is_set(record_upsert):
+                record.metadata_ = record_upsert.metadata
+                record.updated_at = datetime.utcnow()
+
+            records.append(record)
+
+        self._db.add_all(records)
+        await self._db.flush(records)
+        await self._upsert_records_relationships(records, bulk_upsert.items)
+        await distribution.unsafe_update_records_status(self._db, records)
 
         await self._db.commit()
 
+        await _preload_records_relationships_before_index(self._db, records)
+        await self._search_engine.index_records(dataset, records)
+
         return RecordsBulkWithUpdateInfo(
             items=records,
             updated_item_ids=[record.id for record in found_records.values()],

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

@@ -110,29 +110,28 @@ async def create_user(db: AsyncSession, user_attrs: dict, workspaces: Union[List
     if await get_user_by_username(db, user_attrs["username"]) is not None:
         raise NotUniqueError(f"User username `{user_attrs['username']}` is not unique")
 
-    async with db.begin_nested():
-        user = await User.create(
-            db,
-            first_name=user_attrs["first_name"],
-            last_name=user_attrs["last_name"],
-            username=user_attrs["username"],
-            role=user_attrs["role"],
-            password_hash=hash_password(user_attrs["password"]),
-            autocommit=False,
-        )
-
-        if workspaces is not None:
-            for workspace_name in workspaces:
-                workspace = await Workspace.get_by(db, name=workspace_name)
-                if not workspace:
-                    raise UnprocessableEntityError(f"Workspace '{workspace_name}' does not exist")
-
-                await WorkspaceUser.create(
-                    db,
-                    workspace_id=workspace.id,
-                    user_id=user.id,
-                    autocommit=False,
-                )
+    user = await User.create(
+        db,
+        first_name=user_attrs["first_name"],
+        last_name=user_attrs["last_name"],
+        username=user_attrs["username"],
+        role=user_attrs["role"],
+        password_hash=hash_password(user_attrs["password"]),
+        autocommit=False,
+    )
+
+    if workspaces is not None:
+        for workspace_name in workspaces:
+            workspace = await Workspace.get_by(db, name=workspace_name)
+            if not workspace:
+                raise UnprocessableEntityError(f"Workspace '{workspace_name}' does not exist")
+
+            await WorkspaceUser.create(
+                db,
+                workspace_id=workspace.id,
+                user_id=user.id,
+                autocommit=False,
+            )
 
     await db.commit()