From 02e0f23120b1b2af8f95271e79c2a6bf15181ed3 Mon Sep 17 00:00:00 2001 From: zandercymatics <141044360+zandercymatics@users.noreply.github.com> Date: Mon, 24 Jun 2024 09:37:31 -0600 Subject: [PATCH] Cleanup --- .../management/commands/populate_verification_type.py | 9 ++++----- .../management/commands/transfer_federal_agency_type.py | 8 ++++---- .../management/commands/utility/terminal_helper.py | 8 ++++---- 3 files changed, 12 insertions(+), 13 deletions(-) diff --git a/src/registrar/management/commands/populate_verification_type.py b/src/registrar/management/commands/populate_verification_type.py index b61521977..30cccfc15 100644 --- a/src/registrar/management/commands/populate_verification_type.py +++ b/src/registrar/management/commands/populate_verification_type.py @@ -12,12 +12,11 @@ class Command(BaseCommand, PopulateScriptTemplate): def handle(self, **kwargs): """Loops through each valid User object and updates its verification_type value""" filter_condition = {"verification_type__isnull": True} - self.mass_populate_field(User, filter_condition, ["verification_type"]) + self.mass_update_records(User, filter_condition, ["verification_type"]) - def populate_field(self, field_to_update): + def update_record(self, record: User): """Defines how we update the verification_type field""" - field_to_update.set_user_verification_type() + record.set_user_verification_type() logger.info( - f"{TerminalColors.OKCYAN}Updating {field_to_update} => " - f"{field_to_update.verification_type}{TerminalColors.OKCYAN}" + f"{TerminalColors.OKCYAN}Updating {record} => " f"{record.verification_type}{TerminalColors.OKCYAN}" ) diff --git a/src/registrar/management/commands/transfer_federal_agency_type.py b/src/registrar/management/commands/transfer_federal_agency_type.py index e9ee6f3ce..c1816aeb9 100644 --- a/src/registrar/management/commands/transfer_federal_agency_type.py +++ b/src/registrar/management/commands/transfer_federal_agency_type.py @@ -1,6 +1,6 @@ import logging from django.core.management import BaseCommand -from registrar.management.commands.utility.terminal_helper import PopulateScriptTemplate +from registrar.management.commands.utility.terminal_helper import PopulateScriptTemplate, TerminalColors from registrar.models import FederalAgency, DomainRequest @@ -19,7 +19,7 @@ class Command(BaseCommand, PopulateScriptTemplate): """Loops through each valid User object and updates its verification_type value""" # Get all existing domain requests. Select_related allows us to skip doing db queries. - self.all_domain_requests = DomainRequest.objects.select_related("federal_agency").distinct() + self.all_domain_requests = DomainRequest.objects.select_related("federal_agency") self.mass_update_records( FederalAgency, filter_conditions={"agency__isnull": False}, fields_to_update=["federal_type"] ) @@ -28,11 +28,11 @@ class Command(BaseCommand, PopulateScriptTemplate): """Defines how we update the federal_type field on each record.""" request = self.all_domain_requests.filter(federal_agency__agency=record.agency).first() record.federal_type = request.federal_type + logger.info(f"{TerminalColors.OKCYAN}Updating {str(record)} => {record.federal_type}{TerminalColors.OKCYAN}") def should_skip_record(self, record) -> bool: # noqa """Defines the conditions in which we should skip updating a record.""" requests = self.all_domain_requests.filter(federal_agency__agency=record.agency, federal_type__isnull=False) # Check if all federal_type values are the same. Skip the record otherwise. - distinct_federal_types = requests.values('federal_type').distinct() + distinct_federal_types = requests.values("federal_type").distinct() return distinct_federal_types.count() != 1 - diff --git a/src/registrar/management/commands/utility/terminal_helper.py b/src/registrar/management/commands/utility/terminal_helper.py index 577423ba7..22054ac7b 100644 --- a/src/registrar/management/commands/utility/terminal_helper.py +++ b/src/registrar/management/commands/utility/terminal_helper.py @@ -59,7 +59,6 @@ class ScriptDataHelper: model_class.objects.bulk_update(page.object_list, fields_to_update) - # This template handles logging and bulk updating for you, for repetitive scripts that update a few fields. # It is the ultimate lazy mans shorthand. Don't use this for anything terribly complicated. # See the transfer_federal_agency.py file for example usage - its really quite simple! @@ -143,17 +142,18 @@ class PopulateScriptTemplate(ABC): def get_failure_message(self, record) -> str: """Returns the message that we will display if a record fails to update""" return f"{TerminalColors.FAIL}" f"Failed to update {record}" f"{TerminalColors.ENDC}" - + def should_skip_record(self, record) -> bool: # noqa """Defines the condition in which we should skip updating a record.""" # By default - don't skip return False - class TerminalHelper: @staticmethod - def log_script_run_summary(to_update, failed_to_update, skipped, debug: bool, log_header=None, display_as_str=False): + def log_script_run_summary( + to_update, failed_to_update, skipped, debug: bool, log_header=None, display_as_str=False + ): """Prints success, failed, and skipped counts, as well as all affected objects.""" update_success_count = len(to_update)