models: add table lock and translate old db usage (bug 1886854)

Author: zzzeid

src/lando/api/legacy/api/diff_warnings.py

@@ -15,7 +15,6 @@
 
 from lando.api.legacy.decorators import require_phabricator_api_key
 from lando.main.models.revision import DiffWarning, DiffWarningStatus
-from lando.api.legacy.storage import db
 
 logger = logging.getLogger(__name__)
 
@@ -41,8 +40,7 @@ def post(data: dict):
             type="https://developer.mozilla.org/en-US/docs/Web/HTTP/Status/400",
         )
     warning = DiffWarning(**data)
-    db.session.add(warning)
-    db.session.commit()
+    warning.save()
     return warning.serialize(), 201
 
 
@@ -58,7 +56,7 @@ def delete(pk: str):
             type="https://developer.mozilla.org/en-US/docs/Web/HTTP/Status/400",
         )
     warning.status = DiffWarningStatus.ARCHIVED
-    db.session.commit()
+    warning.save()
     return warning.serialize(), 200
 
 

src/lando/api/legacy/api/landing_jobs.py

@@ -8,7 +8,6 @@
 
 from lando.api import auth
 from lando.main.models.landing_job import LandingJob, LandingJobAction, LandingJobStatus
-from lando.api.legacy.storage import db
 
 logger = logging.getLogger(__name__)
 
@@ -61,7 +60,7 @@ def put(landing_job_id: str, data: dict):
 
     if landing_job.status in (LandingJobStatus.SUBMITTED, LandingJobStatus.DEFERRED):
         landing_job.transition_status(LandingJobAction.CANCEL)
-        db.session.commit()
+        landing_job.save()
         return {"id": landing_job.id}, 200
     else:
         raise ProblemException(

src/lando/api/legacy/api/transplants.py

@@ -55,7 +55,6 @@
     get_landable_repos_for_revision_data,
     request_extended_revision_data,
 )
-from lando.api.legacy.storage import db
 from lando.api.legacy.tasks import admin_remove_phab_project
 from lando.api.legacy.transplants import (
     TransplantAssessment,
@@ -354,10 +353,8 @@ def post(phab: PhabricatorClient, data: dict):
         lando_revision = Revision.get_from_revision_id(revision_id)
         if not lando_revision:
             lando_revision = Revision(revision_id=revision_id)
-            db.session.add(lando_revision)
-
         lando_revision.diff_id = diff_id
-        db.session.commit()
+        lando_revision.save()
 
         revision_reviewers[lando_revision.id] = get_approved_by_ids(
             phab,
@@ -373,7 +370,7 @@ def post(phab: PhabricatorClient, data: dict):
 
         raw_diff = phab.call_conduit("differential.getrawdiff", diffID=diff["id"])
         lando_revision.set_patch(raw_diff, patch_data)
-        db.session.commit()
+        lando_revision.save()
         lando_revisions.append(lando_revision)
 
     ldap_username = g.auth0_user.email
@@ -384,8 +381,7 @@ def post(phab: PhabricatorClient, data: dict):
         )
     )
     stack_ids = [revision.revision_id for revision in lando_revisions]
-    with db.session.begin_nested():
-        LandingJob.lock_table()
+    with LandingJob.lock_table():
         if (
             LandingJob.revisions_query(stack_ids)
             .filter(
@@ -412,7 +408,7 @@ def post(phab: PhabricatorClient, data: dict):
     # Submit landing job.
     job.status = LandingJobStatus.SUBMITTED
     job.set_landed_revision_diffs()
-    db.session.commit()
+    job.save()
 
     logger.info(f"New landing job {job.id} created for {landing_repo.tree} repo.")
 

src/lando/main/models/base.py

@@ -1,12 +1,14 @@
 from __future__ import annotations
 
+from contextlib import ContextDecorator
 import logging
 import os
 import subprocess
 import tempfile
 from pathlib import Path
 
 from django.db import models
+from django.db import connection
 
 from lando import settings
 from lando.utils import GitPatchHelper
@@ -23,6 +25,22 @@ class BaseModel(models.Model):
     class Meta:
         abstract = True
 
+    class lock_table(ContextDecorator):
+        """Decorator to lock table for current model."""
+
+        def __init__(self, model, lock):
+            self.lock = lock
+
+            if lock not in ("SHARE ROW EXCLUSIVE", ):
+                raise ValueError(f"{lock} not valid.")
+
+        def __enter__(self):
+            cursor = connection.cursor()
+            cursor.execute(f"LOCK TABLE {self._meta.db_table} IN {self.lock} MODE")
+
+        def __exit__(self, exc_type, exc_value, traceback):
+            pass
+
 
 class Repo(BaseModel):
     name = models.CharField(max_length=255, unique=True)