Work in progress: Require invite acceptance to join an org

data/model/legacy.py

@@ -43,6 +43,9 @@ class InvalidRobotException(DataModelException):
 class InvalidTeamException(DataModelException):
   pass
 
+class InvalidTeamMemberException(DataModelException):
+  pass
+
 
 class InvalidPasswordException(DataModelException):
   pass
@@ -291,11 +294,46 @@ def remove_team(org_name, team_name, removed_by_username):
   team.delete_instance(recursive=True, delete_nullable=True)
 
 
+def add_or_invite_to_team(team, user=None, email=None, adder=None):
+  # If the user is a member of the organization, then we simply add the
+  # user directly to the team. Otherwise, an invite is created for the user/email.
+  # We return None if the user was directly added and the invite object if the user was invited.
+  if email:
+    try:
+      user = User.get(email=email)
+    except User.DoesNotExist:
+      pass
+
+  requires_invite = True
+  if user:
+    orgname = team.organization.username
+
+    # If the user is part of the organization (or a robot), then no invite is required.
+    if user.robot:
+      requires_invite = False
+      if not user.username.startswith(orgname + '+'):
+        raise InvalidTeamMemberException('Cannot add the specified robot to this team, ' +
+                                         'as it is not a member of the organization')
+    else:      
+      Org = User.alias()
+      found = User.select(User.username)
+      found = found.where(User.username == user.username).join(TeamMember).join(Team)
+      found = found.join(Org, on=(Org.username == orgname)).limit(1)
+      requires_invite = not any(found)
+
+  # If we have a valid user and no invite is required, simply add the user to the team.
+  if user and not requires_invite:
+    add_user_to_team(user, team)
+    return None
+
+  return TeamMemberInvite.create(user=user, email=email if not user else None, team=team)
+
+
 def add_user_to_team(user, team):
   try:
     return TeamMember.create(user=user, team=team)
   except Exception:
-    raise DataModelException('Unable to add user \'%s\' to team: \'%s\'' %
+    raise DataModelException('User \'%s\' is already a member of team \'%s\'' %
                              (user.username, team.name))
 
 
@@ -570,6 +608,10 @@ def get_organization_team_members(teamid):
   query = joined.where(Team.id == teamid)
   return query
 
+def get_organization_team_member_invites(teamid):
+  joined = TeamMemberInvite.select().join(Team).join(User)
+  query = joined.where(Team.id == teamid)
+  return query
 
 def get_organization_member_set(orgname):
   Org = User.alias()