Add script template

This commit is contained in:
zandercymatics 2024-04-22 14:32:42 -06:00
parent c24a235d8e
commit 5aa4e8f484
No known key found for this signature in database
GPG key ID: FF4636ABEC9682B7
2 changed files with 62 additions and 35 deletions

View file

@ -1,46 +1,22 @@
import argparse
import logging import logging
from typing import List from typing import List
from django.core.management import BaseCommand from django.core.management import BaseCommand
from registrar.management.commands.utility.terminal_helper import TerminalColors, TerminalHelper, ScriptDataHelper from registrar.management.commands.utility.terminal_helper import ScriptTemplate, TerminalColors
from registrar.models import User from registrar.models import User
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
class Command(BaseCommand): class Command(ScriptTemplate):
help = "Loops through each valid User object and updates its verification_type value" help = "Loops through each valid User object and updates its verification_type value"
def handle(self, **kwargs): def handle(self):
"""Loops through each valid User object and updates its verification_type value""" """Loops through each valid User object and updates its verification_type value"""
filter_condition = {
users = User.objects.filter(verification_type__isnull=True) "verification_type__isnull": True
}
# Code execution will stop here if the user prompts "N" ScriptTemplate.mass_populate_field(User, filter_condition, ["verification_type"])
TerminalHelper.prompt_for_execution(
system_exit_on_terminate=True, def populate_field(self, field_to_update):
info_to_inspect=f""" """Defines how we update the verification_type field"""
==Proposed Changes== field_to_update.set_user_verification_type()
Number of User objects to change: {len(users)}
This field will be updated on each record: verification_type
""",
prompt_title="Do you wish to patch verification_type data?",
)
logger.info("Updating...")
user_to_update: List[User] = []
user_failed_to_update: List[User] = []
for user in users:
try:
user.set_user_verification_type()
user_to_update.append(user)
except Exception as err:
user_failed_to_update.append(user)
logger.error(err)
logger.error(f"{TerminalColors.FAIL}" f"Failed to update {user}" f"{TerminalColors.ENDC}")
# Do a bulk update on the first_ready field
ScriptDataHelper.bulk_update_fields(User, user_to_update, ["verification_type"])
# Log what happened
TerminalHelper.log_script_run_summary(user_to_update, user_failed_to_update, skipped=[], debug=True)

View file

@ -2,6 +2,7 @@ import logging
import sys import sys
from django.core.paginator import Paginator from django.core.paginator import Paginator
from typing import List from typing import List
from django.core.management import BaseCommand
from registrar.utility.enums import LogCode from registrar.utility.enums import LogCode
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@ -58,6 +59,56 @@ class ScriptDataHelper:
model_class.objects.bulk_update(page.object_list, fields_to_update) model_class.objects.bulk_update(page.object_list, fields_to_update)
class ScriptTemplate(BaseCommand):
"""
Contains common script actions for our scripts which can be prefilled as templates.
"""
@staticmethod
def mass_populate_field(sender, filter_conditions, fields_to_update):
"""Loops through each valid "sender" object - specified by filter_conditions - and
updates fields defined by fields_to_update using populate_function.
You must define populate_field before you can use this function.
"""
objects = sender.objects.filter(**filter_conditions)
# Code execution will stop here if the user prompts "N"
TerminalHelper.prompt_for_execution(
system_exit_on_terminate=True,
info_to_inspect=f"""
==Proposed Changes==
Number of {sender} objects to change: {len(objects)}
These fields will be updated on each record: {fields_to_update}
""",
prompt_title="Do you wish to patch this data?",
)
logger.info("Updating...")
to_update: List[sender] = []
failed_to_update: List[sender] = []
for updated_object in objects:
try:
ScriptTemplate.populate_field(updated_object)
to_update.append(updated_object)
except Exception as err:
to_update.append(updated_object)
logger.error(err)
logger.error(f"{TerminalColors.FAIL}" f"Failed to update {updated_object}" f"{TerminalColors.ENDC}")
# Do a bulk update on the first_ready field
ScriptDataHelper.bulk_update_fields(sender, to_update, fields_to_update)
# Log what happened
TerminalHelper.log_script_run_summary(to_update, failed_to_update, skipped=[], debug=True)
@staticmethod
def populate_field(field_to_update):
"""Defines how we update each field. Must be defined before using mass_populate_field."""
raise NotImplementedError("This method should be implemented by the child class.")
class TerminalHelper: class TerminalHelper:
@staticmethod @staticmethod
def log_script_run_summary(to_update, failed_to_update, skipped, debug: bool, log_header=None): def log_script_run_summary(to_update, failed_to_update, skipped, debug: bool, log_header=None):