From 1e540335d6e44c377ae35ca0bd2da8e70e65ec38 Mon Sep 17 00:00:00 2001 From: David Kennedy Date: Fri, 31 May 2024 21:52:15 -0400 Subject: [PATCH] added command scripts clean_tables, export_tables and import_tables --- .../management/commands/clean_tables.py | 33 +++++++ .../management/commands/export_tables.py | 56 ++++++++++++ .../management/commands/import_tables.py | 85 +++++++++++++++++++ 3 files changed, 174 insertions(+) create mode 100644 src/registrar/management/commands/clean_tables.py create mode 100644 src/registrar/management/commands/export_tables.py create mode 100644 src/registrar/management/commands/import_tables.py diff --git a/src/registrar/management/commands/clean_tables.py b/src/registrar/management/commands/clean_tables.py new file mode 100644 index 000000000..2f331d005 --- /dev/null +++ b/src/registrar/management/commands/clean_tables.py @@ -0,0 +1,33 @@ +import logging +from django.core.management import BaseCommand +from django.apps import apps +from django.db import transaction + +logger = logging.getLogger(__name__) + +class Command(BaseCommand): + help = "Clean tables in database to prepare for import." + + def handle(self, **options): + """Delete all rows from a list of tables""" + table_names = [ + "DomainInformation", "DomainRequest", "Domain", "User", "Contact", + "Website", "DraftDomain", "HostIp", "Host" + ] + + for table_name in table_names: + self.clean_table(table_name) + + def clean_table(self, table_name): + """Delete all rows in the given table""" + try: + # Get the model class dynamically + model = apps.get_model('registrar', table_name) + # Use a transaction to ensure database integrity + with transaction.atomic(): + model.objects.all().delete() + logger.info(f"Successfully cleaned table {table_name}") + except LookupError: + logger.error(f"Model for table {table_name} not found.") + except Exception as e: + logger.error(f"Error cleaning table {table_name}: {e}") diff --git a/src/registrar/management/commands/export_tables.py b/src/registrar/management/commands/export_tables.py new file mode 100644 index 000000000..742fb93a4 --- /dev/null +++ b/src/registrar/management/commands/export_tables.py @@ -0,0 +1,56 @@ +import logging +import os +import pyzipper +from django.core.management import BaseCommand +import registrar.admin + +logger = logging.getLogger(__name__) + +class Command(BaseCommand): + help = ( + "Exports tables in csv format to zip file in tmp directory." + ) + + def handle(self, **options): + """Generates CSV files for specified tables and creates a zip archive""" + table_names = [ + "User", "Contact", "Domain", "DomainRequest", "DomainInformation", + "UserDomainRole", "DraftDomain", "Website", "HostIp", "Host" + ] + + # Ensure the tmp directory exists + os.makedirs("tmp", exist_ok=True) + + for table_name in table_names: + self.export_table(table_name) + + # Create a zip file containing all the CSV files + zip_filename = "tmp/exported_tables.zip" + with pyzipper.AESZipFile(zip_filename, 'w', compression=pyzipper.ZIP_DEFLATED) as zipf: + for table_name in table_names: + csv_filename = f"tmp/{table_name}.csv" + if os.path.exists(csv_filename): + zipf.write(csv_filename, os.path.basename(csv_filename)) + logger.info(f"Added {csv_filename} to zip archive {zip_filename}") + + # Remove the CSV files after adding them to the zip file + for table_name in table_names: + csv_filename = f"tmp/{table_name}.csv" + if os.path.exists(csv_filename): + os.remove(csv_filename) + logger.info(f"Removed temporary file {csv_filename}") + + def export_table(self, table_name): + """Export a given table to a csv file in the tmp directory""" + resourcename = f"{table_name}Resource" + try: + resourceclass = getattr(registrar.admin, resourcename) + dataset = resourceclass().export() + filename = f"tmp/{table_name}.csv" + with open(filename, "w") as outputfile: + outputfile.write(dataset.csv) + logger.info(f"Successfully exported {table_name} to {filename}") + except AttributeError: + logger.error(f"Resource class {resourcename} not found in registrar.admin") + except Exception as e: + logger.error(f"Failed to export {table_name}: {e}") diff --git a/src/registrar/management/commands/import_tables.py b/src/registrar/management/commands/import_tables.py new file mode 100644 index 000000000..bad94b5a3 --- /dev/null +++ b/src/registrar/management/commands/import_tables.py @@ -0,0 +1,85 @@ +import logging +import os +import pyzipper +import tablib +from django.apps import apps +from django.db import transaction +from django.core.management import BaseCommand +import registrar.admin + +logger = logging.getLogger(__name__) + +class Command(BaseCommand): + help = "Imports tables from a zip file, exported_tables.zip, containing CSV files in the tmp directory." + + def handle(self, **options): + """Extracts CSV files from a zip archive and imports them into the respective tables""" + table_names = [ + "User", "Contact", "Domain", "Host", "HostIp", "DraftDomain", "Website", + "DomainRequest", "DomainInformation", "UserDomainRole" + ] + + # Ensure the tmp directory exists + os.makedirs("tmp", exist_ok=True) + + # Unzip the file + zip_filename = "tmp/exported_tables.zip" + if not os.path.exists(zip_filename): + logger.error(f"Zip file {zip_filename} does not exist.") + return + + with pyzipper.AESZipFile(zip_filename, 'r') as zipf: + zipf.extractall("tmp") + logger.info(f"Extracted zip file {zip_filename} into tmp directory") + + # Import each CSV file + for table_name in table_names: + self.import_table(table_name) + + def import_table(self, table_name): + """Import data from a CSV file into the given table""" + resourcename = f"{table_name}Resource" + csv_filename = f"tmp/{table_name}.csv" + try: + if not os.path.exists(csv_filename): + logger.error(f"CSV file {csv_filename} not found.") + return + + # if table_name is Contact, clean the table first + if table_name == "Contact": + self.clean_table(table_name) + + resourceclass = getattr(registrar.admin, resourcename) + resource_instance = resourceclass() + with open(csv_filename, "r") as csvfile: + #dataset = resource_instance.import_data(csvfile.read()) + dataset = tablib.Dataset().load(csvfile.read(), format='csv') + result = resource_instance.import_data(dataset, dry_run=False) + + if result.has_errors(): + logger.error(f"Errors occurred while importing {csv_filename}: {result.row_errors()}") + else: + logger.info(f"Successfully imported {csv_filename} into {table_name}") + + except AttributeError: + logger.error(f"Resource class {resourcename} not found in registrar.admin") + except Exception as e: + logger.error(f"Failed to import {csv_filename}: {e}") + finally: + if os.path.exists(csv_filename): + os.remove(csv_filename) + logger.info(f"Removed temporary file {csv_filename}") + + def clean_table(self, table_name): + """Delete all rows in the given table""" + try: + # Get the model class dynamically + model = apps.get_model('registrar', table_name) + # Use a transaction to ensure database integrity + with transaction.atomic(): + model.objects.all().delete() + logger.info(f"Successfully cleaned table {table_name}") + except LookupError: + logger.error(f"Model for table {table_name} not found.") + except Exception as e: + logger.error(f"Error cleaning table {table_name}: {e}")