Rework concurrent upload using upsert strategy

Replace the old try/except IntegrityError + cleanup loop pattern with an atomic upsert in Upload.create_upload().
Decouple the upload directory name from the DB primary key via transaction_id. Upload directory is created at this stage, no need to take care for it later.

With the upsert logic, the ObjectDeletedError scenario which happened because a concurrent request could delete a stale upload row mid-operation is eliminated:
  - During push_finish, the heartbeat context manager continuously updates last_ping, keeping the upload fresh
  throughout the operation
  - A concurrent request can only take over an upload whose last_ping has expired
  - Since heartbeat prevents expiry, no other request can claim the row while push_finish is running
  - The upload object therefore stays valid for the full lifetime of the request — ObjectDeletedError becomes
  impossible

Author: varmar05

server/mergin/sync/models.py

@@ -19,7 +19,7 @@
 from flask_login import current_user
 from pygeodiff import GeoDiff
 from sqlalchemy import text, null, desc, nullslast, tuple_
-from sqlalchemy.dialects.postgresql import ARRAY, BIGINT, UUID, JSONB, ENUM
+from sqlalchemy.dialects.postgresql import ARRAY, BIGINT, UUID, JSONB, ENUM, insert
 from sqlalchemy.types import String
 from sqlalchemy.ext.hybrid import hybrid_property
 from pygeodiff.geodifflib import GeoDiffLibError, GeoDiffLibConflictError
@@ -1808,6 +1808,7 @@ class Upload(db.Model):
     created = db.Column(db.DateTime, default=datetime.utcnow)
     # last ping time to determine if upload is still active
     last_ping = db.Column(db.DateTime, nullable=False, default=datetime.utcnow)
+    transaction_id = db.Column(db.String, unique=True, nullable=False, index=True)
 
     user = db.relationship("User")
     project = db.relationship(
@@ -1825,10 +1826,98 @@ def __init__(self, project: Project, version: int, changes: dict, user_id: int):
         self.version = version
         self.changes = ChangesSchema().dump(changes)
         self.user_id = user_id
+        self.transaction_id = str(uuid.uuid4())
+
+    @classmethod
+    def create_upload(
+        cls, project_id: str, version: int, changes: dict, user_id: int
+    ) -> Upload | None:
+        """Create upload session, it can either create a new record or handover an existing one but with new transaction id
+        Old transaction folder is removed and new one is created.
+        """
+        now = datetime.now(timezone.utc)
+        expiration = current_app.config["LOCKFILE_EXPIRATION"]
+        new_tx_id = str(uuid.uuid4())
+
+        # CTE captures the existing row's transaction_id BEFORE the upsert (pre-statement snapshot)
+        # NULL in RETURNING means fresh INSERT, non-NULL means we took over a stale upload
+        existing_cte = (
+            db.select(Upload.transaction_id)
+            .where(
+                Upload.project_id == project_id,
+                Upload.version == version,
+            )
+            .cte("existing")
+        )
+
+        stmt = (
+            insert(Upload)
+            .values(
+                id=str(uuid.uuid4()),
+                transaction_id=new_tx_id,
+                project_id=project_id,
+                version=version,
+                user_id=user_id,
+                last_ping=now,
+                changes=ChangesSchema().dump(changes),
+            )
+            .add_cte(existing_cte)
+        )
+
+        upsert_stmt = stmt.on_conflict_do_update(
+            constraint="uq_upload_project_id",
+            set_={
+                "transaction_id": new_tx_id,
+                "user_id": user_id,
+                "last_ping": now,
+                "changes": ChangesSchema().dump(changes),
+            },
+            # ONLY update if the existing row is stale
+            where=(Upload.last_ping < (now - timedelta(seconds=expiration))),
+        )
+
+        upsert_stmt = upsert_stmt.returning(
+            Upload,
+            db.select(existing_cte.c.transaction_id)
+            .scalar_subquery()
+            .label("old_transaction_id"),
+        )
+
+        result = db.session.execute(upsert_stmt).fetchone()
+        db.session.commit()
+
+        # if nothing returned, it means the WHERE clause failed (active upload)
+        if not result:
+            return
+
+        upload = result.Upload
+        old_transaction_id = result.old_transaction_id
+        os.makedirs(upload.upload_dir)
+
+        # old_transaction_id is NULL on fresh INSERT, set to old UUID when taking over a stale upload
+        if old_transaction_id:
+            upload.project.sync_failed(
+                "", "push_lost", "Push artefact removed by subsequent push", user_id
+            )
+            if os.path.exists(
+                os.path.join(
+                    upload.project.storage.project_dir, "tmp", old_transaction_id
+                )
+            ):
+                move_to_tmp(
+                    os.path.join(
+                        upload.project.storage.project_dir, "tmp", old_transaction_id
+                    ),
+                    old_transaction_id,
+                )
+
+        return upload
 
     @property
     def upload_dir(self):
-        return os.path.join(self.project.storage.project_dir, "tmp", self.id)
+        return os.path.join(
+            self.project.storage.project_dir, "tmp", self.transaction_id
+        )
 
     def is_active(self):
         """Check if upload is still active because there was a ping from underlying process"""
@@ -1896,7 +1985,7 @@ def clear(self):
         Uploaded files and table records are removed, and another upload can start.
         """
         try:
-            move_to_tmp(self.upload_dir, self.id)
+            move_to_tmp(self.upload_dir, self.transaction_id)
             db.session.delete(self)
             db.session.commit()
         except Exception:

server/mergin/sync/permissions.py

@@ -271,17 +271,16 @@ def check_project_permissions(
     return None
 
 
-def get_upload(transaction_id):
-    upload = Upload.query.get_or_404(transaction_id)
+def get_upload_or_fail(transaction_id: str) -> Upload:
+    upload = Upload.query.filter_by(transaction_id=transaction_id).first_or_404()
     # upload to 'removed' projects is forbidden
     if upload.project.removed_at:
         abort(404)
 
     if upload.user_id != current_user.id:
         abort(403, "You do not have permissions for ongoing upload")
 
-    upload_dir = os.path.join(upload.project.storage.project_dir, "tmp", transaction_id)
-    return upload, upload_dir
+    return upload
 
 
 def projects_query(permission, as_admin=True, public=True):

server/mergin/sync/public_api_controller.py

@@ -25,7 +25,7 @@
 from pygeodiff import GeoDiffLibError
 from flask_login import current_user
 from sqlalchemy import and_, desc, asc
-from sqlalchemy.exc import IntegrityError
+from sqlalchemy.exc import IntegrityError, SQLAlchemyError
 from gevent import sleep
 import base64
 from werkzeug.exceptions import HTTPException, Conflict
@@ -70,7 +70,7 @@
     require_project,
     projects_query,
     ProjectPermissions,
-    get_upload,
+    get_upload_or_fail,
     require_project_by_uuid,
 )
 from .utils import (
@@ -774,13 +774,6 @@ def project_push(namespace, project_name):
     if all(len(changes[key]) == 0 for key in changes.keys()):
         abort(400, "No changes")
 
-    # reject upload early if there is another one already running
-    pending_upload = Upload.query.filter_by(
-        project_id=project.id, version=version
-    ).first()
-    if pending_upload and pending_upload.is_active():
-        abort(400, "Another process is running. Please try later.")
-
     try:
         ChangesSchema().validate(changes)
         upload_changes = ChangesSchema().dump(changes)
@@ -812,44 +805,20 @@ def project_push(namespace, project_name):
     if requested_storage > ws.storage:
         return StorageLimitHit(current_usage, ws.storage).response(422)
 
-    upload = Upload(project, version, upload_changes, current_user.id)
-    db.session.add(upload)
     try:
-        # Creating upload transaction with different project's version is possible.
-        db.session.commit()
+        upload = Upload.create_upload(
+            project.id, version, upload_changes, current_user.id
+        )
+        if not upload:
+            abort(400, "Another process is running. Please try later.")
+
         logging.info(
-            f"Upload transaction {upload.id} created for project: {project.id}, version: {version}"
+            f"Upload transaction {upload.transaction_id} created for project: {project.id}, version: {version}"
         )
-    except IntegrityError:
+    except (IntegrityError, SQLAlchemyError) as err:
         db.session.rollback()
-        # check and clean dangling uploads or abort
-        for current_upload in project.uploads.all():
-            if current_upload.is_active():
-                abort(400, "Another process is running. Please try later.")
-            db.session.delete(current_upload)
-            db.session.commit()
-            # previous push attempt is definitely lost
-            project.sync_failed(
-                "",
-                "push_lost",
-                "Push artefact removed by subsequent push",
-                current_user.id,
-            )
-
-        # Try again after cleanup
-        db.session.add(upload)
-        try:
-            db.session.commit()
-            logging.info(
-                f"Upload transaction {upload.id} created for project: {project.id}, version: {version}"
-            )
-            move_to_tmp(upload.upload_dir)
-        except IntegrityError as err:
-            logging.error(f"Failed to create upload session: {str(err)}")
-            abort(422, "Failed to create upload session. Please try later.")
-
-    # Create transaction folder
-    os.makedirs(upload.upload_dir)
+        logging.exception(f"Failed to create upload: {str(err)}")
+        abort(422, "Failed to create upload session. Please try later.")
 
     # Update immediately without uploading of new/modified files and remove transaction after successful commit
     if not (changes["added"] or changes["updated"]):
@@ -874,15 +843,15 @@ def project_push(namespace, project_name):
             db.session.commit()
             logging.info(
                 f"A project version {ProjectVersion.to_v_name(next_version)} for project: {project.id} created. "
-                f"Transaction id: {upload.id}. No upload."
+                f"Transaction id: {upload.transaction_id}. No upload."
             )
             project_version_created.send(pv)
             push_finished.send(pv)
             return jsonify(ProjectSchema().dump(project)), 200
         except IntegrityError as err:
             db.session.rollback()
             logging.exception(
-                f"Failed to upload a new project version using transaction id: {upload.id}: {str(err)}"
+                f"Failed to upload a new project version using transaction id: {upload.transaction_id}: {str(err)}"
             )
             abort(422, "Failed to upload a new project version. Please try later.")
         except gevent.timeout.Timeout:
@@ -891,7 +860,7 @@ def project_push(namespace, project_name):
         finally:
             upload.clear()
 
-    return {"transaction": upload.id}, 200
+    return {"transaction": upload.transaction_id}, 200
 
 
 @auth_required
@@ -908,7 +877,7 @@ def chunk_upload(transaction_id, chunk_id):
 
     :rtype: Dict
     """
-    upload, upload_dir = get_upload(transaction_id)
+    upload = get_upload_or_fail(transaction_id)
     request.view_args["project"] = upload.project
     chunks = []
     for file in upload.changes["added"] + upload.changes["updated"]:
@@ -917,7 +886,7 @@ def chunk_upload(transaction_id, chunk_id):
     if chunk_id not in chunks:
         abort(404)
 
-    dest = os.path.join(upload_dir, "chunks", chunk_id)
+    dest = os.path.join(upload.upload_dir, "chunks", chunk_id)
     with upload.heartbeat(30):
         try:
             # we could have used request.data here, but it could eventually cause OOM issue
@@ -950,7 +919,7 @@ def push_finish(transaction_id):
 
     :rtype: None
     """
-    upload, upload_dir = get_upload(transaction_id)
+    upload = get_upload_or_fail(transaction_id)
     request.view_args["project"] = upload.project
     project = upload.project
     next_version = project.next_version()
@@ -989,7 +958,7 @@ def push_finish(transaction_id):
 
             abort(422, f"Failed to create new version: {msg}")
 
-    files_dir = os.path.join(upload_dir, "files", v_next_version)
+    files_dir = os.path.join(upload.upload_dir, "files", v_next_version)
     target_dir = os.path.join(project.storage.project_dir, v_next_version)
     if os.path.exists(target_dir):
         pv = ProjectVersion.query.filter_by(
@@ -1007,29 +976,31 @@ def push_finish(transaction_id):
         move_to_tmp(target_dir)
 
     try:
-        user_agent = get_user_agent(request)
-        device_id = get_device_id(request)
-        pv = ProjectVersion(
-            project,
-            next_version,
-            current_user.id,
-            file_changes,
-            get_ip(request),
-            user_agent,
-            device_id,
-        )
-        db.session.add(pv)
-        db.session.add(project)
-        db.session.commit()
+        # let's keep upload alive until all work is done so no one else can claim it
+        with upload.heartbeat(5):
+            user_agent = get_user_agent(request)
+            device_id = get_device_id(request)
+            pv = ProjectVersion(
+                project,
+                next_version,
+                current_user.id,
+                file_changes,
+                get_ip(request),
+                user_agent,
+                device_id,
+            )
+            db.session.add(pv)
+            db.session.add(project)
+            db.session.commit()
 
-        # let's move uploaded files where they are expected to be
-        os.renames(files_dir, version_dir)
+            # let's move uploaded files where they are expected to be
+            os.renames(files_dir, version_dir)
 
-        logging.info(
-            f"Push finished for project: {project.id}, project version: {v_next_version}, transaction id: {transaction_id}."
-        )
-        project_version_created.send(pv)
-        push_finished.send(pv)
+            logging.info(
+                f"Push finished for project: {project.id}, project version: {v_next_version}, transaction id: {transaction_id}."
+            )
+            project_version_created.send(pv)
+            push_finished.send(pv)
     except (psycopg2.Error, FileNotFoundError, IntegrityError) as err:
         db.session.rollback()
         logging.exception(
@@ -1059,10 +1030,8 @@ def push_cancel(transaction_id):
 
     :rtype: None
     """
-    upload, upload_dir = get_upload(transaction_id)
-    db.session.delete(upload)
-    db.session.commit()
-    move_to_tmp(upload_dir)
+    upload = get_upload_or_fail(transaction_id)
+    upload.clear()
     return NoContent, 200