diff --git a/src/registrar/management/commands/load_senior_official_table.py b/src/registrar/management/commands/load_senior_official_table.py index 0aba7f4de..647c19d84 100644 --- a/src/registrar/management/commands/load_senior_official_table.py +++ b/src/registrar/management/commands/load_senior_official_table.py @@ -3,7 +3,7 @@ import csv import logging import os from django.core.management import BaseCommand -from registrar.management.commands.utility.terminal_helper import TerminalHelper +from registrar.management.commands.utility.terminal_helper import TerminalHelper, TerminalColors from registrar.models import SeniorOfficial, FederalAgency @@ -45,7 +45,7 @@ class Command(BaseCommand): existing_agencies = FederalAgency.objects.all() # Read the CSV - added_senior_officials, skipped_rows = [], [] + added_senior_officials, skipped_rows, updated_rows = [], [], [] with open(federal_cio_csv_path, "r") as requested_file: for row in csv.DictReader(requested_file): # Note: the csv doesn't have a phone field, but we can try to pull one anyway. @@ -61,7 +61,8 @@ class Command(BaseCommand): agency_name = row.get("Agency").strip() if row.get("Agency") else None # Only first_name, last_name, and title are required - if row and all(so_kwargs[field] for field in ["first_name", "last_name", "title"]): + required_fields = ["first_name", "last_name", "title"] + if row and all(so_kwargs[field] for field in required_fields): # Get the underlying federal agency record if agency_name: @@ -71,25 +72,56 @@ class Command(BaseCommand): new_so = SeniorOfficial(**so_kwargs) # Before adding this record, check to make sure we aren't adding a duplicate. - if not existing_senior_officials.filter(**so_kwargs).exists(): + existing_field = existing_senior_officials.filter( + first_name=so_kwargs.get("first_name"), + last_name=so_kwargs.get("last_name"), + title=so_kwargs.get("title"), + ).first() + if not existing_field: added_senior_officials.append(new_so) message = f"Added record: {new_so}" TerminalHelper.colorful_logger("INFO", "OKCYAN", message) else: - skipped_rows.append(new_so) - message = f"Skipping add on duplicate record: {new_so}" - TerminalHelper.colorful_logger("WARNING", "YELLOW", message) + duplicate_field = existing_senior_officials.filter(**so_kwargs).first() + if not duplicate_field: + # If we can, just update the row instead + for field, value in so_kwargs.items(): + if getattr(existing_field, field) != value: + setattr(existing_field, field, value) + updated_rows.append(existing_field) + message = f"Updating record: {existing_field}" + TerminalHelper.colorful_logger("INFO", "OKBLUE", message) + else: + # if this field is a duplicate, don't do anything + skipped_rows.append(duplicate_field) + message = f"Skipping add on duplicate record: {duplicate_field}" + TerminalHelper.colorful_logger("WARNING", "YELLOW", message) else: skipped_rows.append(row) message = f"Skipping row (missing first_name, last_name, or title): {row}" TerminalHelper.colorful_logger("WARNING", "YELLOW", message) # Bulk create the SO fields - SeniorOfficial.objects.bulk_create(added_senior_officials) + if len(added_senior_officials) > 0: + SeniorOfficial.objects.bulk_create(added_senior_officials) - # Log what happened - added_message = f"Added {len(added_senior_officials)} records" - TerminalHelper.colorful_logger("INFO", "OKGREEN", added_message) + added_message = f"Added {len(added_senior_officials)} records" + TerminalHelper.colorful_logger("INFO", "OKGREEN", added_message) + + # Bulk update the SO fields (if any) + if len(updated_rows) > 0: + updated_fields = [ + "first_name", + "last_name", + "title", + "email", + "phone", + "federal_agency" + ] + SeniorOfficial.objects.bulk_update(updated_rows, updated_fields) + + skipped_message = f"Updated {len(updated_rows)} records" + TerminalHelper.colorful_logger("INFO", "OKBLUE", skipped_message) if len(skipped_rows) > 0: skipped_message = f"Skipped {len(skipped_rows)} records"