zydronium/manage.get.gov
1
0
Fork 0
mirror of https://github.com/cisagov/manage.get.gov.git synced 2025-05-19 10:59:21 +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
return response
class PublicContactResource(resources.ModelResource):
"""defines how each field in the referenced model should be mapped to the corresponding fields in the
import/export file"""
@ -2251,31 +2252,20 @@ class PublicContactResource(resources.ModelResource):
class Meta:
model = models.PublicContact
def import_row(
self,
row,
instance_loader,
using_transactions=True,
dry_run=False,
raise_errors=None,
**kwargs
):
def import_row(self, row, instance_loader, using_transactions=True, dry_run=False, raise_errors=None, **kwargs):
"""Override kwargs skip_epp_save and set to True"""
kwargs['skip_epp_save'] = True
kwargs["skip_epp_save"] = True
return super().import_row(
row,
instance_loader,
using_transactions=using_transactions,
dry_run=dry_run,
raise_errors=raise_errors,
**kwargs
**kwargs,
)
def save_instance(
self, instance, is_create, using_transactions=True, dry_run=False
):
"""Override save_instance setting skip_epp_save to True
"""
def save_instance(self, instance, is_create, using_transactions=True, dry_run=False):
"""Override save_instance setting skip_epp_save to True"""
self.before_save_instance(instance, using_transactions, dry_run)
if self._meta.use_bulk:
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__)
class Command(BaseCommand):
help = "Clean tables in database to prepare for import."
@ -17,8 +18,16 @@ class Command(BaseCommand):
return
table_names = [
"DomainInformation", "DomainRequest", "PublicContact", "Domain", "User",
"Contact", "Website", "DraftDomain", "HostIp", "Host"
"DomainInformation",
"DomainRequest",
"PublicContact",
"Domain",
"User",
"Contact",
"Website",
"DraftDomain",
"HostIp",
"Host",
]
for table_name in table_names:
@ -28,7 +37,7 @@ class Command(BaseCommand):
"""Delete all rows in the given table"""
try:
# 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
with transaction.atomic():
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__)
class Command(BaseCommand):
help = (
"Exports tables in csv format to zip file in tmp directory."
)
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", "PublicContact"
"User",
"Contact",
"Domain",
"DomainRequest",
"DomainInformation",
"UserDomainRole",
"DraftDomain",
"Website",
"HostIp",
"Host",
"PublicContact",
]
# Ensure the tmp directory exists
@ -26,7 +34,7 @@ class Command(BaseCommand):
# 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:
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):

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

@ -10,6 +10,7 @@ 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."
@ -21,8 +22,17 @@ class Command(BaseCommand):
return
table_names = [
"User", "Contact", "Domain", "Host", "HostIp", "DraftDomain", "Website",
"DomainRequest", "DomainInformation", "UserDomainRole", "PublicContact"
"User",
"Contact",
"Domain",
"Host",
"HostIp",
"DraftDomain",
"Website",
"DomainRequest",
"DomainInformation",
"UserDomainRole",
"PublicContact",
]
# Ensure the tmp directory exists
@ -34,7 +44,7 @@ class Command(BaseCommand):
logger.error(f"Zip file {zip_filename} does not exist.")
return
with pyzipper.AESZipFile(zip_filename, 'r') as zipf:
with pyzipper.AESZipFile(zip_filename, "r") as zipf:
zipf.extractall("tmp")
logger.info(f"Extracted zip file {zip_filename} into tmp directory")
@ -58,8 +68,8 @@ class Command(BaseCommand):
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')
# dataset = resource_instance.import_data(csvfile.read())
dataset = tablib.Dataset().load(csvfile.read(), format="csv")
result = resource_instance.import_data(dataset, dry_run=False, skip_epp_save=True)
if result.has_errors():
@ -80,7 +90,7 @@ class Command(BaseCommand):
"""Delete all rows in the given table"""
try:
# 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
with transaction.atomic():
model.objects.all().delete()