updated user login logic for one-to-many transition domains relationship, and added some error handling

src/registrar/models/user.py

@@ -6,6 +6,7 @@ from django.db import models
 from .domain_invitation import DomainInvitation
 from registrar.models import TransitionDomain
 from registrar.models import DomainInformation
+from registrar.models import Domain
 
 from phonenumber_field.modelfields import PhoneNumberField  # type: ignore
 
@@ -64,12 +65,9 @@ class User(AbstractUser):
     def is_restricted(self):
         return self.status == self.RESTRICTED
 
-    def first_login(self):
+    def check_domain_invitations_on_login(self):
-        """Callback when the user is authenticated for the very first time.
+        """When a user first arrives on the site, we need to retrieve any domain
-
+        invitations that match their email address."""
-        When a user first arrives on the site, we need to retrieve any domain
-        invitations that match their email address.
-        """
         for invitation in DomainInvitation.objects.filter(
             email=self.email, status=DomainInvitation.INVITED
         ):
@@ -83,22 +81,73 @@ class User(AbstractUser):
                 logger.warn(
                     "Failed to retrieve invitation %s", invitation, exc_info=True
                 )
-        
+
-        transition_domain_exists = TransitionDomain.objects.filter(
+    def check_transition_domains_on_login(self):
+        """When a user first arrives on the site, we need to check
+        if they are logging in with the same e-mail as a
+        transition domain and update our database accordingly."""
+
+        for transition_domain in TransitionDomain.objects.filter(
             username=self.email
-        ).exists()
+        ):
-        if transition_domain_exists:
             # Looks like the user logged in with the same e-mail as
-            # a corresponding transition domain.  Create a Domain
+            # one or more corresponding transition domains.  
-            # Information object.
+            # Create corresponding DomainInformation objects.
-            # TODO: Do we need to check for existing Domain Info objects?
+
-            # TODO: Should we add a Domain to the DomainInfo object?
+            # NOTE: adding an ADMIN user role for this user
-            # NOTE that adding a user role for this user
+            # for each domain should already be done
-            # as admin for this domain is already done
+            # in the invitation.retrieve() method. 
-            # in the incitation.retrieve() method. So
+            # However, if the migration scripts for transition
-            # we don't need to do that here.
+            # domain objects were not executed correctly, 
-           new_domain_info = DomainInformation(creator=self)
+            # there could be transition domains without
-           new_domain_info.save()
+            # any corresponding Domain & DomainInvitation objects,
+            # which means the invitation.retrieve() method might
+            # not execute.
+            # Check that there is a corresponding domain object
+            # for this transition domain.  If not, we have an error
+            # with our data and migrations need to be run again.
+            # TODO: how should we handle this?
+
+            # Get the domain that corresponds with this transition domain
+            domain_exists = Domain.objects.filter(name=transition_domain.name).exists()
+            if not domain_exists:
+                logger.warn("""There are transition domains without
+                            corresponding domain objects! 
+                            Please run migration scripts for transition domains
+                            (See data_migration.md)""")
+                # TODO: throw exception??
+                break
+            domain = Domain.objects.get(name=transition_domain.name)
+
+            # Create a domain information object, if one doesn't
+            # already exist
+            domain_info_exists = DomainInformation.objects.filter(
+                creator=self,
+                domain=domain
+            ).exists()
+            if not domain_info_exists:
+                # TODO: Should we add a Domain to the DomainInfo object?
+                new_domain_info = DomainInformation(
+                    creator=self,
+                    domain=domain) 
+                new_domain_info.save()
+
+    def first_login(self):
+        """Callback when the user is authenticated for the very first time.
+
+        When a user first arrives on the site, we need to retrieve any domain
+        invitations that match their email address.
+
+        We also need to check if they are logging in with the same e-mail
+        as a transition domain and update our database accordingly.
+        """
+
+        # PART 1: DOMAIN INVITATIONS
+        self.check_domain_invitations_on_login()
+        
+        # PART 2: TRANSITION DOMAINS
+        self.check_transition_domains_on_login()
+        
 
     class Meta:
         permissions = [