added command scripts clean_tables, export_tables and import_tables

This commit is contained in:
David Kennedy 2024-05-31 21:52:15 -04:00
parent cb1c0dece0
commit 1e540335d6
No known key found for this signature in database
GPG key ID: 6528A5386E66B96B
3 changed files with 174 additions and 0 deletions

View file

@ -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}")

View file

@ -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}")

View file

@ -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}")