Simplify logic

src/registrar/management/commands/create_federal_portfolio.py

@@ -14,7 +14,11 @@ class Command(BaseCommand):
     help = "Creates a federal portfolio given a FederalAgency name"
 
     def add_arguments(self, parser):
-        """Add our arguments."""
+        """Add three arguments:
+        1. agency_name => the value of FederalAgency.agency
+        2. --parse_requests => if true, adds the given portfolio to each related DomainRequest
+        3. --parse_domains => if true, adds the given portfolio to each related DomainInformation
+        """
         parser.add_argument(
             "agency_name",
             help="The name of the FederalAgency to add",
@@ -37,19 +41,14 @@ class Command(BaseCommand):
         if not parse_requests and not parse_domains:
             raise CommandError("You must specify at least one of --parse_requests or --parse_domains.")
 
-        agencies = FederalAgency.objects.filter(agency__iexact=agency_name)
+        federal_agency = FederalAgency.objects.filter(agency__iexact=agency_name).first()
-
+        if not federal_agency:
-        # TODO - maybe we can add an option here to add this if it doesn't exist?
-        if not agencies.exists():
             raise ValueError(
                 f"Cannot find the federal agency '{agency_name}' in our database. "
                 "The value you enter for `agency_name` must be "
                 "prepopulated in the FederalAgency table before proceeding."
             )
 
-        # There should be a one-to-one relationship between the name and the agency.
-        federal_agency = agencies.get()
-
         portfolio = self.create_or_modify_portfolio(federal_agency)
         self.create_suborganizations(portfolio, federal_agency)
 
@@ -60,9 +59,7 @@ class Command(BaseCommand):
             self.handle_portfolio_domains(portfolio, federal_agency)
 
     def create_or_modify_portfolio(self, federal_agency):
-        """Tries to create a portfolio record based off of a federal agency.
+        """Creates or modifies a portfolio record based on a federal agency."""
-        If the record already exists, we prompt the user to proceed then
-        update the record."""
         portfolio_args = {
             "federal_agency": federal_agency,
             "organization_name": federal_agency.agency,
@@ -71,19 +68,15 @@ class Command(BaseCommand):
             "notes": "Auto-generated record",
         }
 
-        senior_official = federal_agency.so_federal_agency
+        if federal_agency.so_federal_agency.exists():
-        if senior_official.exists():
+            portfolio_args["senior_official"] = federal_agency.so_federal_agency.first()
-            portfolio_args["senior_official"] = senior_official.first()
 
-        # Create the Portfolio value if it doesn't exist
+        portfolio, created = Portfolio.objects.get_or_create(**portfolio_args)
-        existing_portfolio = Portfolio.objects.filter(organization_name=federal_agency.agency)
+
-        if not existing_portfolio.exists():
+        if created:
-            portfolio = Portfolio.objects.create(**portfolio_args)
             message = f"Created portfolio '{portfolio}'"
             TerminalHelper.colorful_logger(logger.info, TerminalColors.OKGREEN, message)
-            return portfolio
         else:
-
             proceed = TerminalHelper.prompt_for_execution(
                 system_exit_on_terminate=False,
                 info_to_inspect=f"""The given portfolio '{federal_agency.agency}' already exists in our DB.
@@ -91,62 +84,44 @@ class Command(BaseCommand):
                 """,
                 prompt_title="Do you wish to modify this record?",
             )
-            if not proceed:
+            if proceed:
-                if len(existing_portfolio) > 1:
+                for key, value in portfolio_args.items():
-                    raise ValueError(f"Could not use portfolio '{federal_agency.agency}': multiple records exist.")
+                    setattr(portfolio, key, value)
-                else:
+                portfolio.save()
-                    # Just return the portfolio object without modifying it
+                message = f"Modified portfolio '{portfolio}'"
-                    return existing_portfolio.get()
+                TerminalHelper.colorful_logger(logger.info, TerminalColors.MAGENTA, message)
 
-            if len(existing_portfolio) > 1:
+        return portfolio
-                raise ValueError(f"Could not update portfolio '{federal_agency.agency}': multiple records exist.")
-
-            existing_portfolio.update(**portfolio_args)
-            message = f"Modified portfolio '{existing_portfolio.first()}'"
-            TerminalHelper.colorful_logger(logger.info, TerminalColors.MAGENTA, message)
-            return existing_portfolio.get()
 
     def create_suborganizations(self, portfolio: Portfolio, federal_agency: FederalAgency):
-        """Given a list of organization_names on DomainInformation objects (filtered by agency),
+        """Create Suborganizations tied to the given portfolio based on DomainInformation objects"""
-        create multiple Suborganizations tied to the given portfolio"""
+        valid_agencies = DomainInformation.objects.filter(federal_agency=federal_agency, organization_name__isnull=False)
-        valid_agencies = DomainInformation.objects.filter(federal_agency=federal_agency)
+        org_names = set(valid_agencies.values_list("organization_name", flat=True))
-        org_names = valid_agencies.values_list("organization_name", flat=True)
+
-        if len(org_names) < 1:
+        if not org_names:
-            message = f"No suborganizations found for {federal_agency}"
+            TerminalHelper.colorful_logger(logger.warning, TerminalColors.YELLOW, f"No suborganizations found for {federal_agency}")
-            TerminalHelper.colorful_logger(logger.warning, TerminalColors.YELLOW, message)
             return
 
-        # Check if we need to update any existing suborgs first.
+        # Check if we need to update any existing suborgs first. This step is optional.
-        # This step is optional.
         existing_suborgs = Suborganization.objects.filter(name__in=org_names)
-        if len(existing_suborgs) > 0:
+        if existing_suborgs.exists():
             self._update_existing_suborganizations(portfolio, existing_suborgs)
 
-        # Add any suborgs that don't presently exist
+        # Create new suborgs, as long as they don't exist in the db already
-        excluded_org_names = existing_suborgs.values_list("name", flat=True)
+        new_suborgs = []
-        suborgs = []
+        for name in org_names - set(existing_suborgs.values_list("name", flat=True)):
-        for name in org_names:
+            if name.lower() == portfolio.organization_name.lower():
-            if name and name not in excluded_org_names:
+                # If the suborg name is a portfolio name that currently exists, thats not a suborg - thats the portfolio itself!
-                if portfolio.organization_name and name.lower() == portfolio.organization_name.lower():
+                # In this case, we can use this as an opportunity to update address information.
-                    # If the suborg name is the name that currently exists,
+                self._update_portfolio_location_details(portfolio, valid_agencies.filter(organization_name=name).first())
-                    # thats not a suborg - thats the portfolio itself!
+            else:
-                    # In this case, we can use this as an opportunity to update
+                new_suborgs.append(Suborganization(name=name, portfolio=portfolio))
-                    # address information and the like
-                    self._update_portfolio_location_details(portfolio, valid_agencies.get(organization_name=name))
-                else:
-                    suborg = Suborganization(
-                        name=name,
-                        portfolio=portfolio,
-                    )
-                    suborgs.append(suborg)
 
-        if len(org_names) > 1:
+        if new_suborgs:
-            Suborganization.objects.bulk_create(suborgs)
+            Suborganization.objects.bulk_create(new_suborgs)
-            message = f"Added {len(suborgs)} suborganizations"
+            TerminalHelper.colorful_logger(logger.info, TerminalColors.OKGREEN, f"Added {len(new_suborgs)} suborganizations")
-            TerminalHelper.colorful_logger(logger.info, TerminalColors.OKGREEN, message)
         else:
-            message = "No suborganizations added"
+            TerminalHelper.colorful_logger(logger.warning, TerminalColors.YELLOW, "No suborganizations added")
-            TerminalHelper.colorful_logger(logger.warning, TerminalColors.YELLOW, message)
 
     def _update_existing_suborganizations(self, portfolio, orgs_to_update):
         """
@@ -163,15 +138,13 @@ class Command(BaseCommand):
             """,
             prompt_title="Do you wish to modify existing suborganizations?",
         )
-        if not proceed:
+        if proceed:
-            return
+            for org in orgs_to_update:
+                org.portfolio = portfolio
 
-        for org in orgs_to_update:
+            Suborganization.objects.bulk_update(orgs_to_update, ["portfolio"])
-            org.portfolio = portfolio
+            message = f"Updated {len(orgs_to_update)} suborganizations"
-
+            TerminalHelper.colorful_logger(logger.info, TerminalColors.MAGENTA, message)
-        Suborganization.objects.bulk_update(orgs_to_update, ["portfolio"])
-        message = f"Updated {len(orgs_to_update)} suborganizations"
-        TerminalHelper.colorful_logger(logger.info, TerminalColors.MAGENTA, message)
 
     def _update_portfolio_location_details(self, portfolio: Portfolio, domain_info: DomainInformation):
         """
@@ -191,8 +164,8 @@ class Command(BaseCommand):
             # Copy the value from the domain info object to the portfolio object
             value = getattr(domain_info, prop_name)
             setattr(portfolio, prop_name, value)
-        portfolio.save()
 
+        portfolio.save()
         message = f"Updated location details on portfolio '{portfolio}'"
         TerminalHelper.colorful_logger(logger.info, TerminalColors.OKGREEN, message)
 
@@ -230,6 +203,5 @@ class Command(BaseCommand):
             domain_info.portfolio = portfolio
 
         DomainInformation.objects.bulk_update(domain_infos, ["portfolio"])
-
         message = f"Added portfolio '{portfolio}' to {len(domain_infos)} domains"
         TerminalHelper.colorful_logger(logger.info, TerminalColors.OKGREEN, message)