zydronium/manage.get.gov
1
0
Fork 0
mirror of https://github.com/cisagov/manage.get.gov.git synced 2025-05-19 19:09:22 +02:00
Code Issues Projects Releases Packages Wiki Activity

formatted for code readability

Browse source
This commit is contained in:
David Kennedy 2024-06-01 19:57:04 -04:00
parent 736189e22c
commit 16fa86edc9
No known key found for this signature in database
GPG key ID: 6528A5386E66B96B
4 changed files with 56 additions and 39 deletions
Download patch file Download diff file Expand all files Collapse all files

22
src/registrar/admin.py
View file

@ -2244,6 +2244,7 @@ class DraftDomainAdmin(ListHeaderAdmin, ImportExportModelAdmin):
# If no redirection is needed, return the original response # If no redirection is needed, return the original response
return response return response
class PublicContactResource(resources.ModelResource): class PublicContactResource(resources.ModelResource):
"""defines how each field in the referenced model should be mapped to the corresponding fields in the """defines how each field in the referenced model should be mapped to the corresponding fields in the
import/export file""" import/export file"""
@ -2251,31 +2252,20 @@ class PublicContactResource(resources.ModelResource):
class Meta: class Meta:
model = models.PublicContact model = models.PublicContact
def import_row( def import_row(self, row, instance_loader, using_transactions=True, dry_run=False, raise_errors=None, **kwargs):
self,
row,
instance_loader,
using_transactions=True,
dry_run=False,
raise_errors=None,
**kwargs
):
"""Override kwargs skip_epp_save and set to True""" """Override kwargs skip_epp_save and set to True"""
kwargs['skip_epp_save'] = True kwargs["skip_epp_save"] = True
return super().import_row( return super().import_row(
row, row,
instance_loader, instance_loader,
using_transactions=using_transactions, using_transactions=using_transactions,
dry_run=dry_run, dry_run=dry_run,
raise_errors=raise_errors, raise_errors=raise_errors,
**kwargs **kwargs,
) )
def save_instance( def save_instance(self, instance, is_create, using_transactions=True, dry_run=False):
self, instance, is_create, using_transactions=True, dry_run=False """Override save_instance setting skip_epp_save to True"""
):
"""Override save_instance setting skip_epp_save to True
"""
self.before_save_instance(instance, using_transactions, dry_run) self.before_save_instance(instance, using_transactions, dry_run)
if self._meta.use_bulk: if self._meta.use_bulk:
if is_create: if is_create:

15
src/registrar/management/commands/clean_tables.py
View file

@ -6,6 +6,7 @@ from django.db import transaction
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
class Command(BaseCommand): class Command(BaseCommand):
help = "Clean tables in database to prepare for import." help = "Clean tables in database to prepare for import."
@ -17,8 +18,16 @@ class Command(BaseCommand):
return return
table_names = [ table_names = [
"DomainInformation", "DomainRequest", "PublicContact", "Domain", "User", "DomainInformation",
"Contact", "Website", "DraftDomain", "HostIp", "Host" "DomainRequest",
"PublicContact",
"Domain",
"User",
"Contact",
"Website",
"DraftDomain",
"HostIp",
"Host",
] ]
for table_name in table_names: for table_name in table_names:
@ -28,7 +37,7 @@ class Command(BaseCommand):
"""Delete all rows in the given table""" """Delete all rows in the given table"""
try: try:
# Get the model class dynamically # Get the model class dynamically
model = apps.get_model('registrar', table_name) model = apps.get_model("registrar", table_name)
# Use a transaction to ensure database integrity # Use a transaction to ensure database integrity
with transaction.atomic(): with transaction.atomic():
model.objects.all().delete() model.objects.all().delete()

20
src/registrar/management/commands/export_tables.py
View file

@ -6,16 +6,24 @@ import registrar.admin
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
class Command(BaseCommand): class Command(BaseCommand):
help = ( help = "Exports tables in csv format to zip file in tmp directory."
"Exports tables in csv format to zip file in tmp directory."
)
def handle(self, **options): def handle(self, **options):
"""Generates CSV files for specified tables and creates a zip archive""" """Generates CSV files for specified tables and creates a zip archive"""
table_names = [ table_names = [
"User", "Contact", "Domain", "DomainRequest", "DomainInformation", "User",
"UserDomainRole", "DraftDomain", "Website", "HostIp", "Host", "PublicContact" "Contact",
"Domain",
"DomainRequest",
"DomainInformation",
"UserDomainRole",
"DraftDomain",
"Website",
"HostIp",
"Host",
"PublicContact",
] ]
# Ensure the tmp directory exists # Ensure the tmp directory exists
@ -26,7 +34,7 @@ class Command(BaseCommand):
# Create a zip file containing all the CSV files # Create a zip file containing all the CSV files
zip_filename = "tmp/exported_tables.zip" zip_filename = "tmp/exported_tables.zip"
with pyzipper.AESZipFile(zip_filename, 'w', compression=pyzipper.ZIP_DEFLATED) as zipf: with pyzipper.AESZipFile(zip_filename, "w", compression=pyzipper.ZIP_DEFLATED) as zipf:
for table_name in table_names: for table_name in table_names:
csv_filename = f"tmp/{table_name}.csv" csv_filename = f"tmp/{table_name}.csv"
if os.path.exists(csv_filename): if os.path.exists(csv_filename):

22
src/registrar/management/commands/import_tables.py
View file

@ -10,6 +10,7 @@ import registrar.admin
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
class Command(BaseCommand): class Command(BaseCommand):
help = "Imports tables from a zip file, exported_tables.zip, containing CSV files in the tmp directory." help = "Imports tables from a zip file, exported_tables.zip, containing CSV files in the tmp directory."
@ -21,8 +22,17 @@ class Command(BaseCommand):
return return
table_names = [ table_names = [
"User", "Contact", "Domain", "Host", "HostIp", "DraftDomain", "Website", "User",
"DomainRequest", "DomainInformation", "UserDomainRole", "PublicContact" "Contact",
"Domain",
"Host",
"HostIp",
"DraftDomain",
"Website",
"DomainRequest",
"DomainInformation",
"UserDomainRole",
"PublicContact",
] ]
# Ensure the tmp directory exists # Ensure the tmp directory exists
@ -34,7 +44,7 @@ class Command(BaseCommand):
logger.error(f"Zip file {zip_filename} does not exist.") logger.error(f"Zip file {zip_filename} does not exist.")
return return
with pyzipper.AESZipFile(zip_filename, 'r') as zipf: with pyzipper.AESZipFile(zip_filename, "r") as zipf:
zipf.extractall("tmp") zipf.extractall("tmp")
logger.info(f"Extracted zip file {zip_filename} into tmp directory") logger.info(f"Extracted zip file {zip_filename} into tmp directory")
@ -58,8 +68,8 @@ class Command(BaseCommand):
resourceclass = getattr(registrar.admin, resourcename) resourceclass = getattr(registrar.admin, resourcename)
resource_instance = resourceclass() resource_instance = resourceclass()
with open(csv_filename, "r") as csvfile: with open(csv_filename, "r") as csvfile:
#dataset = resource_instance.import_data(csvfile.read()) # dataset = resource_instance.import_data(csvfile.read())
dataset = tablib.Dataset().load(csvfile.read(), format='csv') dataset = tablib.Dataset().load(csvfile.read(), format="csv")
result = resource_instance.import_data(dataset, dry_run=False, skip_epp_save=True) result = resource_instance.import_data(dataset, dry_run=False, skip_epp_save=True)
if result.has_errors(): if result.has_errors():
@ -80,7 +90,7 @@ class Command(BaseCommand):
"""Delete all rows in the given table""" """Delete all rows in the given table"""
try: try:
# Get the model class dynamically # Get the model class dynamically
model = apps.get_model('registrar', table_name) model = apps.get_model("registrar", table_name)
# Use a transaction to ensure database integrity # Use a transaction to ensure database integrity
with transaction.atomic(): with transaction.atomic():
model.objects.all().delete() model.objects.all().delete()