refactor max_dm_count

This commit is contained in:
Rachid Mrad 2024-02-29 20:50:50 -05:00
parent c17f2481f1
commit 7e271fc273
No known key found for this signature in database

View file

@ -25,7 +25,7 @@ def write_header(writer, columns):
def get_domain_infos(filter_condition, sort_fields): def get_domain_infos(filter_condition, sort_fields):
domain_infos = ( domain_infos = (
DomainInformation.objects.select_related("domain", "authorizing_official") DomainInformation.objects.select_related("domain", "domain__permissions", "authorizing_official")
.filter(**filter_condition) .filter(**filter_condition)
.order_by(*sort_fields) .order_by(*sort_fields)
) )
@ -161,19 +161,20 @@ def write_csv(
# Reduce the memory overhead when performing the write operation # Reduce the memory overhead when performing the write operation
paginator = Paginator(all_domain_infos, 1000) paginator = Paginator(all_domain_infos, 1000)
if get_domain_managers and len(all_domain_infos) > 0: # The maximum amount of domain managers an account has
# We want to get the max amont of domain managers an # We get the max so we can set the column header accurately
# account has to set the column header dynamically max_dm_count = 0
max_dm_count = max(len(domain_info.domain.permissions.all()) for domain_info in all_domain_infos) total_body_rows = []
update_columns_with_domain_managers(columns, max_dm_count)
if should_write_header:
write_header(writer, columns)
for page_num in paginator.page_range: for page_num in paginator.page_range:
rows = [] rows = []
page = paginator.page(page_num) page = paginator.page(page_num)
for domain_info in page.object_list: for domain_info in page.object_list:
# Get count of all the domain managers for an account
dm_count = len(domain_info.domain.permissions)
if dm_count > max_dm_count:
max_dm_count = dm_count
try: try:
row = parse_row(columns, domain_info, security_emails_dict, get_domain_managers) row = parse_row(columns, domain_info, security_emails_dict, get_domain_managers)
rows.append(row) rows.append(row)
@ -182,8 +183,11 @@ def write_csv(
# It indicates that DomainInformation.domain is None. # It indicates that DomainInformation.domain is None.
logger.error("csv_export -> Error when parsing row, domain was None") logger.error("csv_export -> Error when parsing row, domain was None")
continue continue
total_body_rows.append(rows)
writer.writerows(rows) update_columns_with_domain_managers(columns, max_dm_count)
write_header(writer, columns)
writer.writerows(total_body_rows)
def export_data_type_to_csv(csv_file): def export_data_type_to_csv(csv_file):