resolve conflict with staging

atst/models/application.py

@@ -1,4 +1,4 @@
-from sqlalchemy import and_, Column, ForeignKey, String, UniqueConstraint, TIMESTAMP
+from sqlalchemy import and_, Column, ForeignKey, String, UniqueConstraint
 from sqlalchemy.orm import relationship, synonym
 
 from atst.models.base import Base
@@ -9,7 +9,11 @@ from atst.models.types import Id
 
 
 class Application(
-    Base, mixins.TimestampsMixin, mixins.AuditableMixin, mixins.DeletableMixin
+    Base,
+    mixins.TimestampsMixin,
+    mixins.AuditableMixin,
+    mixins.DeletableMixin,
+    mixins.ClaimableMixin,
 ):
     __tablename__ = "applications"
 
@@ -41,7 +45,6 @@ class Application(
     )
 
     cloud_id = Column(String)
-    claimed_until = Column(TIMESTAMP(timezone=True))
 
     @property
     def users(self):

atst/models/application_role.py

@@ -1,5 +1,5 @@
 from enum import Enum
-from sqlalchemy import Index, ForeignKey, Column, Enum as SQLAEnum, Table
+from sqlalchemy import Index, ForeignKey, Column, Enum as SQLAEnum, Table, String
 from sqlalchemy.dialects.postgresql import UUID
 from sqlalchemy.orm import relationship
 from sqlalchemy.event import listen
@@ -33,6 +33,7 @@ class ApplicationRole(
     mixins.AuditableMixin,
     mixins.PermissionsMixin,
     mixins.DeletableMixin,
+    mixins.ClaimableMixin,
 ):
     __tablename__ = "application_roles"
 
@@ -59,6 +60,8 @@ class ApplicationRole(
         primaryjoin="and_(EnvironmentRole.application_role_id == ApplicationRole.id, EnvironmentRole.deleted == False)",
     )
 
+    cloud_id = Column(String)
+
     @property
     def latest_invitation(self):
         if self.invitations:

atst/models/environment.py

@@ -1,4 +1,4 @@
-from sqlalchemy import Column, ForeignKey, String, TIMESTAMP, UniqueConstraint
+from sqlalchemy import Column, ForeignKey, String, UniqueConstraint
 from sqlalchemy.orm import relationship
 from sqlalchemy.dialects.postgresql import JSONB
 from enum import Enum
@@ -9,7 +9,11 @@ import atst.models.types as types
 
 
 class Environment(
-    Base, mixins.TimestampsMixin, mixins.AuditableMixin, mixins.DeletableMixin
+    Base,
+    mixins.TimestampsMixin,
+    mixins.AuditableMixin,
+    mixins.DeletableMixin,
+    mixins.ClaimableMixin,
 ):
     __tablename__ = "environments"
 
@@ -28,8 +32,6 @@ class Environment(
     cloud_id = Column(String)
     root_user_info = Column(JSONB(none_as_null=True))
 
-    claimed_until = Column(TIMESTAMP(timezone=True))
-
     roles = relationship(
         "EnvironmentRole",
         back_populates="environment",

atst/models/environment_role.py

@@ -1,5 +1,5 @@
 from enum import Enum
-from sqlalchemy import Index, ForeignKey, Column, String, TIMESTAMP, Enum as SQLAEnum
+from sqlalchemy import Index, ForeignKey, Column, String, Enum as SQLAEnum
 from sqlalchemy.dialects.postgresql import UUID
 from sqlalchemy.orm import relationship
 
@@ -15,7 +15,11 @@ class CSPRole(Enum):
 
 
 class EnvironmentRole(
-    Base, mixins.TimestampsMixin, mixins.AuditableMixin, mixins.DeletableMixin
+    Base,
+    mixins.TimestampsMixin,
+    mixins.AuditableMixin,
+    mixins.DeletableMixin,
+    mixins.ClaimableMixin,
 ):
     __tablename__ = "environment_roles"
 
@@ -33,7 +37,6 @@ class EnvironmentRole(
     application_role = relationship("ApplicationRole")
 
     csp_user_id = Column(String())
-    claimed_until = Column(TIMESTAMP(timezone=True))
 
     class Status(Enum):
         PENDING = "pending"

atst/models/mixins/__init__.py

@@ -4,3 +4,4 @@ from .permissions import PermissionsMixin
 from .deletable import DeletableMixin
 from .invites import InvitesMixin
 from .state_machines import FSMMixin
+from .claimable import ClaimableMixin

atst/models/mixins/claimable.py

@@ -0,0 +1,5 @@
+from sqlalchemy import Column, TIMESTAMP
+
+
+class ClaimableMixin(object):
+    claimed_until = Column(TIMESTAMP(timezone=True))

atst/models/portfolio_state_machine.py

@@ -175,11 +175,14 @@ class PortfolioStateMachine(
             app.logger.info(exc.json())
             print(exc.json())
             app.logger.info(payload_data)
+            # TODO: Ensure that failing the stage does not preclude a Celery retry
             self.fail_stage(stage)
+        # TODO: catch and handle general CSP exception here
         except (ConnectionException, UnknownServerException) as exc:
             app.logger.error(
                 f"CSP api call. Caught exception for {self.__repr__()}.", exc_info=1,
             )
+            # TODO: Ensure that failing the stage does not preclude a Celery retry
             self.fail_stage(stage)
 
         self.finish_stage(stage)

atst/models/utils.py

@@ -1,3 +1,5 @@
+from typing import List
+
 from sqlalchemy import func, sql, Interval, and_, or_
 from contextlib import contextmanager
 
@@ -28,7 +30,7 @@ def claim_for_update(resource, minutes=30):
         .filter(
             and_(
                 Model.id == resource.id,
-                or_(Model.claimed_until == None, Model.claimed_until <= func.now()),
+                or_(Model.claimed_until.is_(None), Model.claimed_until <= func.now()),
             )
         )
         .update({"claimed_until": claim_until}, synchronize_session="fetch")
@@ -48,3 +50,51 @@ def claim_for_update(resource, minutes=30):
             Model.claimed_until != None
         ).update({"claimed_until": None}, synchronize_session="fetch")
         db.session.commit()
+
+
+@contextmanager
+def claim_many_for_update(resources: List, minutes=30):
+    """
+    Claim a mutually exclusive expiring hold on a group of resources.
+    Uses the database as a central source of time in case the server clocks have drifted.
+
+    Args:
+        resources:   A list of SQLAlchemy model instances with a `claimed_until` attribute.
+        minutes:    The maximum amount of time, in minutes, to hold the claim.
+    """
+    Model = resources[0].__class__
+
+    claim_until = func.now() + func.cast(
+        sql.functions.concat(minutes, " MINUTES"), Interval
+    )
+
+    ids = tuple(r.id for r in resources)
+
+    # Optimistically query for and update the resources in question. If they're
+    # already claimed, `rows_updated` will be 0 and we can give up.
+    rows_updated = (
+        db.session.query(Model)
+        .filter(
+            and_(
+                Model.id.in_(ids),
+                or_(Model.claimed_until.is_(None), Model.claimed_until <= func.now()),
+            )
+        )
+        .update({"claimed_until": claim_until}, synchronize_session="fetch")
+    )
+    if rows_updated < 1:
+        # TODO: Generalize this exception class so it can take multiple resources
+        raise ClaimFailedException(resources[0])
+
+    # Fetch the claimed resources
+    claimed = db.session.query(Model).filter(Model.id.in_(ids)).all()
+
+    try:
+        # Give the resource to the caller.
+        yield claimed
+    finally:
+        # Release the claim.
+        db.session.query(Model).filter(Model.id.in_(ids)).filter(
+            Model.claimed_until != None
+        ).update({"claimed_until": None}, synchronize_session="fetch")
+        db.session.commit()