Add script template

src/registrar/management/commands/populate_verification_type.py

@@ -1,46 +1,22 @@
-import argparse
 import logging
 from typing import List
 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
 
 logger = logging.getLogger(__name__)
 
 
-class Command(BaseCommand):
+class Command(ScriptTemplate):
     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"""
-
-        users = User.objects.filter(verification_type__isnull=True)
-
-        # 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 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)
+        filter_condition = {
+            "verification_type__isnull": True
+        }
+        ScriptTemplate.mass_populate_field(User, filter_condition, ["verification_type"])
+    
+    def populate_field(self, field_to_update):
+        """Defines how we update the verification_type field"""
+        field_to_update.set_user_verification_type()

src/registrar/management/commands/utility/terminal_helper.py

@@ -2,6 +2,7 @@ import logging
 import sys
 from django.core.paginator import Paginator
 from typing import List
+from django.core.management import BaseCommand
 from registrar.utility.enums import LogCode
 
 logger = logging.getLogger(__name__)
@@ -58,6 +59,56 @@ class ScriptDataHelper:
             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:
     @staticmethod
     def log_script_run_summary(to_update, failed_to_update, skipped, debug: bool, log_header=None):