zydronium/manage.get.gov
1
0
Fork 0
mirror of https://github.com/cisagov/manage.get.gov.git synced 2025-07-21 18:25:58 +02:00
Code Issues Projects Releases Packages Wiki Activity

Add logic to update instead of add

Browse source
This commit is contained in:
zandercymatics 2024-07-29 14:10:29 -06:00
parent 67f519693a
commit a7aa765566
No known key found for this signature in database
GPG key ID: FF4636ABEC9682B7
1 changed files with 43 additions and 11 deletions
Download patch file Download diff file Expand all files Collapse all files

54
src/registrar/management/commands/load_senior_official_table.py
View file

@ -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"