add officers to task order and redirect to TO when they accept a workspace invite

atst/domain/task_orders.py

@@ -4,6 +4,9 @@ from atst.database import db
 from atst.models.task_order import TaskOrder
 from .exceptions import NotFoundError
 
+class TaskOrderError(Exception):
+    pass
+
 
 class TaskOrders(object):
     SECTIONS = {
@@ -89,3 +92,17 @@ class TaskOrders(object):
                 return False
 
         return True
+
+    OFFICERS = ["contracting_officer", "contracting_officer_representative", "security_officer"]
+
+    @classmethod
+    def add_officer(cls, task_order, user, role):
+        if role in TaskOrders.OFFICERS:
+            setattr(task_order, role, user)
+
+            db.session.add(task_order)
+            db.session.commit()
+
+            return task_order
+        else:
+            raise TaskOrderError("{} is not an officer role on task orders".format(role))

atst/models/task_order.py

@@ -14,7 +14,18 @@ class TaskOrder(Base, mixins.TimestampsMixin):
     workspace = relationship("Workspace")
 
     user_id = Column(ForeignKey("users.id"))
-    creator = relationship("User")
+    creator = relationship("User", foreign_keys="TaskOrder.user_id")
+
+    ko_id = Column(ForeignKey("users.id"))
+    contracting_officer = relationship("User", foreign_keys="TaskOrder.ko_id")
+
+    cor_id = Column(ForeignKey("users.id"))
+    contracting_officer_representative = relationship(
+        "User", foreign_keys="TaskOrder.cor_id"
+    )
+
+    so_id = Column(ForeignKey("users.id"))
+    security_officer = relationship("User", foreign_keys="TaskOrder.so_id")
 
     scope = Column(String)  # Cloud Project Scope
     defense_component = Column(String)  # Department of Defense Component

atst/routes/task_orders/new.py

@@ -154,7 +154,7 @@ class UpdateTaskOrderWorkflow(ShowTaskOrderWorkflow):
         },
     ]
 
-    def _update_invitations(self):
+    def _update_officer_invitations(self):
         for officer_type in self.OFFICER_INVITATIONS:
             field = officer_type["field"]
             if (
@@ -169,16 +169,19 @@ class UpdateTaskOrderWorkflow(ShowTaskOrderWorkflow):
                 }
                 invite_service = InvitationService(
                     self.user,
-                    self.task_order.workspace,
+                    self.workspace,
                     {**officer_data, "workspace_role": officer_type["role"]},
                     subject=officer_type["subject"],
                     email_template=officer_type["template"],
                 )
-                invite_service.invite()
+                invite = invite_service.invite()
+                TaskOrders.add_officer(
+                    self.task_order, invite.user, officer_type["role"]
+                )
 
     def update(self):
         self._update_task_order()
-        self._update_invitations()
+        self._update_officer_invitations()
         return self.task_order
 
 

atst/routes/workspaces/invitations.py

@@ -20,6 +20,25 @@ def send_invite_email(owner_name, token, new_member_email):
 def accept_invitation(token):
     invite = Invitations.accept(g.current_user, token)
 
+    # TODO: this will eventually redirect to different places depending on
+    # whether the user is an officer for the TO and what kind of officer they
+    # are. It will also have to manage cases like:
+    #   - the logged-in user has multiple roles on the TO (e.g., KO and COR)
+    #   - the logged-in user has officer roles on multiple unsigned TOs
+    for task_order in invite.workspace.task_orders:
+        if g.current_user == task_order.contracting_officer:
+            return redirect(
+                url_for("task_orders.new", screen=4, task_order_id=task_order.id)
+            )
+        elif g.current_user == task_order.contracting_officer_representative:
+            return redirect(
+                url_for("task_orders.new", screen=4, task_order_id=task_order.id)
+            )
+        elif g.current_user == task_order.security_officer:
+            return redirect(
+                url_for("task_orders.new", screen=4, task_order_id=task_order.id)
+            )
+
     return redirect(
         url_for("workspaces.show_workspace", workspace_id=invite.workspace.id)
     )