diff --git a/docs/developer/management_script_helpers.md b/docs/developer/management_script_helpers.md new file mode 100644 index 000000000..fd1a82839 --- /dev/null +++ b/docs/developer/management_script_helpers.md @@ -0,0 +1,108 @@ +# Terminal Helper Functions +`terminal_helper.py` contains utility functions to assist with common terminal and script operations. + +## TerminalColors +`TerminalColors` provides ANSI color codes as variables to style terminal output. Example usage: + +print(f"{TerminalColors.OKGREEN}Success!{TerminalColors.ENDC}") + +## ScriptDataHelper +### bulk_update_fields + +`bulk_update_fields` performs a memory-efficient bulk update on a Django model in batches using a Paginator. + +Usage: +bulk_update_fields(Domain, page.object_list, ["first_ready"]) + +## PopulateScriptTemplate + +`PopulateScriptTemplate` is an abstract base class that provides a template for creating generic populate scripts. It handles logging and bulk updating for repetitive scripts that update a few fields. + +**Disclaimer:** This template is intended as a shorthand for simple scripts. It is not recommended for complex operations. See `transfer_federal_agency.py` for a straightforward example of how to use this template. + +To use `PopulateScriptTemplate`, create a new class that inherits from it and implement the `update_record` method. This method defines how each record should be updated. + +The class provides the following optional configuration variables: +- `prompt_title`: The header displayed by `prompt_for_execution` when the script starts (default: "Do you wish to proceed?") +- `display_run_summary_items_as_str`: If True, runs `str(item)` on each item when printing the run summary for prettier output (default: False) +- `run_summary_header`: The header for the script run summary printed after the script finishes (default: None) + +The main method provided by `PopulateScriptTemplate` is `mass_update_records`. This method loops through each valid object (specified by `filter_conditions`) and updates the fields defined in `fields_to_update` using the `update_record` method. + +Before updating, `mass_update_records` prompts the user to confirm the proposed changes. If the user does not proceed, the script will exit. + +After processing the records, `mass_update_records` performs a bulk update on the specified fields using `ScriptDataHelper.bulk_update_fields` and logs a summary of the script run using `TerminalHelper.log_script_run_summary`. + +The class also provides helper methods: +- `get_class_name`: Returns a display-friendly class name for the terminal prompt +- `get_failure_message`: Returns the message to display if a record fails to update +- `should_skip_record`: Defines the condition for skipping a record (by default, no records are skipped) + +To create a script using `PopulateScriptTemplate`: +1. Create a new class that inherits from `PopulateScriptTemplate` +2. Implement the `update_record` method to define how each record should be updated +3. Optionally, override the configuration variables and helper methods as needed +4. Call `mass_update_records` within `handle` and run the script + +## TerminalHelper +### log_script_run_summary + +`log_script_run_summary` logs a summary of a script run, including counts of updated, skipped, and failed records. + +### print_conditional + +`print_conditional` conditionally logs a statement at a specified severity if a condition is met. + +### prompt_for_execution + +`prompt_for_execution` prompts the user to inspect a string and confirm if they wish to proceed. Returns True if proceeding, False if skipping, or exits the script. + +### get_file_line_count + +`get_file_line_count` returns the number of lines in a file. + +### print_to_file_conditional + +`print_to_file_conditional` conditionally writes content to a file if a condition is met. + +Refer to the source code for full function signatures and additional details. + +### query_yes_no + +`query_yes_no` prompts the user with a yes/no question and returns True for "yes" or False for "no". + +Usage: +```python +if query_yes_no("Do you want to proceed?"): + print("Proceeding...") +else: + print("Aborting.") +``` + +### query_yes_no_exit + +`query_yes_no_exit` is similar to `query_yes_no` but includes an "exit" option to terminate the script. + +Usage: +if query_yes_no_exit("Continue, abort, or exit?"): + print("Continuing...") +else: + print("Aborting.") + # Script will exit if user selected "e" for exit + +### array_as_string + +`array_as_string` converts a list of strings into a single string with each element on a new line. + +Usage: +```python +my_list = ["apple", "banana", "cherry"] +print(array_as_string(my_list)) +``` + +Output: +``` +apple +banana +cherry +``` diff --git a/src/registrar/management/commands/transfer_federal_agency.py b/src/registrar/management/commands/transfer_federal_agency.py index 3be7d452f..dd7b1e5db 100644 --- a/src/registrar/management/commands/transfer_federal_agency.py +++ b/src/registrar/management/commands/transfer_federal_agency.py @@ -1,30 +1,31 @@ import logging from django.core.management import BaseCommand -from registrar.management.commands.utility.terminal_helper import PopulateScriptTemplate, TerminalColors +from registrar.management.commands.utility.terminal_helper import PopulateScriptTemplate from registrar.models import FederalAgency, DomainRequest -from registrar.models.utility.generic_helper import convert_queryset_to_dict + logger = logging.getLogger(__name__) +# This command uses the PopulateScriptTemplate. +# 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. class Command(BaseCommand, PopulateScriptTemplate): help = "Loops through each valid User object and updates its verification_type value" prompt_title = "Do you wish to update all Federal Agencies?" + display_run_summary_items_as_str = True def handle(self, **kwargs): """Loops through each valid User object and updates its verification_type value""" # Get all existing domain requests self.all_domain_requests = DomainRequest.objects.select_related("federal_agency").distinct() - - filter_condition = { - "agency__isnull": False, - } - updated_fields = ["federal_type"] - self.mass_update_records(FederalAgency, filter_condition, updated_fields) + self.mass_update_records( + FederalAgency, filter_conditions={"agency__isnull": False}, fields_to_update=["federal_type"] + ) def update_record(self, record: FederalAgency): - """Defines how we update the federal_type field""" + """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 diff --git a/src/registrar/management/commands/utility/terminal_helper.py b/src/registrar/management/commands/utility/terminal_helper.py index bc5b0cc8c..577423ba7 100644 --- a/src/registrar/management/commands/utility/terminal_helper.py +++ b/src/registrar/management/commands/utility/terminal_helper.py @@ -59,12 +59,29 @@ 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! class PopulateScriptTemplate(ABC): """ Contains an ABC for generic populate scripts """ - prompt_title = "Do you wish to proceed?" + # Optional script-global config variables. For the most part, you can leave these untouched. + # Defines what prompt_for_execution displays as its header when you first start the script + prompt_title: str = "Do you wish to proceed?" + + # Runs str(item) over each item when printing. Use this for prettier run summaries. + display_run_summary_items_as_str: bool = False + + # The header when printing the script run summary (after the script finishes) + run_summary_header = None + + @abstractmethod + def update_record(self, record): + """Defines how we update each field. Must be defined before using mass_update_records.""" + raise NotImplementedError def mass_update_records(self, sender, filter_conditions, fields_to_update, debug=True): """Loops through each valid "sender" object - specified by filter_conditions - and @@ -108,7 +125,14 @@ class PopulateScriptTemplate(ABC): ScriptDataHelper.bulk_update_fields(sender, to_update, fields_to_update) # Log what happened - TerminalHelper.log_script_run_summary(to_update, failed_to_update, to_skip, debug, display_as_str=True) + TerminalHelper.log_script_run_summary( + to_update, + failed_to_update, + to_skip, + debug=debug, + log_header=self.run_summary_header, + display_as_str=self.display_run_summary_items_as_str, + ) def get_class_name(self, sender) -> str: """Returns the class name that we want to display for the terminal prompt. @@ -121,15 +145,10 @@ class PopulateScriptTemplate(ABC): return f"{TerminalColors.FAIL}" f"Failed to update {record}" f"{TerminalColors.ENDC}" def should_skip_record(self, record) -> bool: # noqa - """Defines the conditions in which we should skip updating a record.""" + """Defines the condition in which we should skip updating a record.""" # By default - don't skip return False - @abstractmethod - def update_record(self, record): - """Defines how we update each field. Must be defined before using mass_populate_field.""" - raise NotImplementedError - class TerminalHelper: