[CVAT] Fix possible no assignments error when there are (#3346)

* Fix possible no assignments error when there are

* Update readme

Author: zhiltsov-max

packages/examples/cvat/exchange-oracle/README.md

@@ -73,5 +73,6 @@ Available at `/docs` route
 
 To run tests
 ```
-docker compose -f docker-compose.test.yml up --build test --attach test --exit-code-from test
+docker compose -f docker-compose.test.yml up --build test --attach test --exit-code-from test && \
+   docker compose -f docker-compose.test.yml down
 ```

packages/examples/cvat/exchange-oracle/src/services/cvat.py

@@ -794,7 +794,8 @@ def count_jobs_by_escrow_address(
 
 def get_free_job(
     session: Session,
-    cvat_projects: list[int],
+    escrow_address: str,
+    chain_id: int,
     *,
     user_wallet_address: str,
     for_update: bool | ForUpdateParams = False,
@@ -805,7 +806,11 @@ def get_free_job(
     return (
         _maybe_for_update(session.query(Job), enable=for_update)
         .where(
-            Job.cvat_project_id.in_(cvat_projects),
+            Job.project.has(
+                (Project.escrow_address == escrow_address)
+                & (Project.chain_id == chain_id)
+                & (Project.status == ProjectStatuses.annotation)
+            ),
             Job.status == JobStatuses.new,
             ~Job.assignments.any(
                 (
@@ -984,22 +989,27 @@ def get_user_assignments_in_cvat_projects(
     )
 
 
-def count_active_user_assignments(
+def has_active_user_assignments(
     session: Session,
     wallet_address: int,
-    cvat_projects: list[int],
-) -> int:
-    return (
+    escrow_address: str,
+    chain_id: int,
+) -> bool:
+    return session.query(
         session.query(Assignment)
         .where(
-            Assignment.job.has(Job.cvat_project_id.in_(cvat_projects)),
+            Assignment.job.has(
+                Job.project.has(
+                    (Project.escrow_address == escrow_address) & (Project.chain_id == chain_id)
+                )
+            ),
             Assignment.user_wallet_address == wallet_address,
             Assignment.status == AssignmentStatuses.created.value,
             Assignment.completed_at == None,
             utcnow() < Assignment.expires_at,
         )
-        .count()
-    )
+        .exists()
+    ).scalar()
 
 
 # Image

packages/examples/cvat/exchange-oracle/src/services/exchange.py

@@ -18,52 +18,45 @@ def __str__(self) -> str:
         )
 
 
-def create_assignment(escrow_address: str, chain_id: Networks, wallet_address: str) -> str | None:  # noqa: ARG001 (don't we want to use chain_id for filter?)
+def create_assignment(escrow_address: str, chain_id: Networks, wallet_address: str) -> str | None:
     with SessionLocal.begin() as session:
         user = get_or_404(
             cvat_service.get_user_by_id(session, wallet_address, for_update=True),
             wallet_address,
             "user",
         )
 
-        # There can be several projects under one escrow, we need any
-        project = cvat_service.get_project_by_escrow_address(
+        if cvat_service.has_active_user_assignments(
             session,
-            escrow_address,
-            status_in=[
-                ProjectStatuses.annotation
-            ],  # avoid unnecessary locking on completed projects
-            for_update=True,
-        )
-
-        if not project:
-            # Retry without a lock to check if the project doesn't exist
-            get_or_404(
-                cvat_service.get_project_by_escrow_address(
-                    session, escrow_address, status_in=[ProjectStatuses.annotation]
-                ),
-                escrow_address,
-                "job",
-            )
-            return None
-
-        has_active_assignments = (
-            cvat_service.count_active_user_assignments(
-                session, wallet_address=wallet_address, cvat_projects=[project.cvat_id]
-            )
-            > 0
-        )
-        if has_active_assignments:
+            wallet_address=wallet_address,
+            escrow_address=escrow_address,
+            chain_id=chain_id.value,
+        ):
             raise UserHasUnfinishedAssignmentError(
                 "The user already has an unfinished assignment in this project"
             )
 
+        # TODO: Try to put into 1 request. SQLAlchemy generates 2 queries with simple
+        # .options(selectinload(Job.project))
+        project = get_or_404(
+            cvat_service.get_project_by_escrow_address(
+                session, escrow_address, status_in=[ProjectStatuses.annotation]
+            ),
+            escrow_address,
+            "job",
+        )
+
         unassigned_job = cvat_service.get_free_job(
             session,
-            cvat_projects=[project.cvat_id],
+            escrow_address=escrow_address,
+            chain_id=chain_id.value,
             user_wallet_address=wallet_address,
             for_update=True,
+            # lock the job to be able to make a rollback if CVAT requests fail
+            # can potentially be optimized to make less DB requests
+            # and rely only on assignment expiration
         )
+
         if not unassigned_job:
             return None
 
@@ -72,7 +65,12 @@ def create_assignment(escrow_address: str, chain_id: Networks, wallet_address: s
             wallet_address=user.wallet_address,
             cvat_job_id=unassigned_job.cvat_id,
             expires_at=utcnow()
-            + timedelta(seconds=get_default_assignment_timeout(TaskTypes(project.job_type))),
+            + timedelta(
+                seconds=get_default_assignment_timeout(
+                    TaskTypes(project.job_type)
+                    # TODO: need to update this if we have multiple job types per escrow
+                )
+            ),
         )
 
         cvat_service.touch(session, Job, [unassigned_job.id])