move_cap_accounts_to_last_section.py

import io
import csv
import time

from dataclasses import dataclass

from fund_admin.identity.cache import get_user
from fund_admin.reporting.models.templates import FinancialPackageTemplateSettings
from fund_admin.reporting.models.package import FinancialPackageReport
from fund_admin.documents.services import DocumentService


@dataclass
class CustomTemplateReportRow:
    setting_id: int
    custom_template_name: str
    firm_name: str | None
    firm_id: int | None
    created_by_user_id: int | None
    created_by_user_name: str | None
    created_by_user_email: str | None
    num_reports_using_setting: int
    report_id: int
    report_period: str
    report_fund_id: int | None
    report_fund_name: str | None
    report_status: str
    report_preparer_user_id: int | None
    report_preparer_name: str | None

def convert_setting_to_domain_class(setting):
    setting_id = setting.id
    firm = setting.firm
    created_by_user_id = setting.created_by_user
    user = None
    data = []

    try:
        user = get_user(created_by_user_id)
    except Exception as e:
        pass

    reports_using_template = FinancialPackageReport.objects.filter(template_settings__id=setting_id)
    num_reports_using_template = reports_using_template.count()
    for report in reports_using_template:
        report_fund = report.fund
        report_preparer = report.preparer

        row = CustomTemplateReportRow(
            setting_id=setting_id,
            custom_template_name=setting.label,
            firm_id=firm.id if firm else None,
            firm_name=firm.name if firm else None,
            created_by_user_id=created_by_user_id if created_by_user_id else None,
            created_by_user_name=user.full_name if user else None,
            created_by_user_email=user.email if user else None,
            num_reports_using_setting=num_reports_using_template,
            report_id=report.id,
            report_period=report.period_label,
            report_fund_id=report_fund.id if report_fund else None,
            report_fund_name=report_fund.name if report_fund else None,
            report_status=report.status,
            report_preparer_user_id=report_preparer.id if report_preparer else None,
            report_preparer_name=report_preparer.name if report_preparer else None,
        )

        data.append(row)
    return data


def save_data_to_csv(data: list[CustomTemplateReportRow]):
    output = io.StringIO()
    writer = csv.writer(output, quoting=csv.QUOTE_NONNUMERIC)
    
    # Write CSV Headers
    writer.writerow([
        'Report ID',
        'Report Period',
        'Firm ID',
        'Firm Name',
        'Report Fund ID',
        'Report Fund Name',
        'Report Status',
        'Report Preparer',
        'Custom Template ID',
        'Custom Template Name',
        'Custom Template Created By User ID',
        'Custom Template Created By User Name',
        'Num Reports Using This Custom Template'
    ])
    
    for row in data:
        writer.writerow([
            row.report_id,
            row.report_period,
            row.firm_id,
            row.firm_name,
            row.report_fund_id,
            row.report_fund_name,
            row.report_status,
            row.report_preparer_name,
            row.setting_id,
            row.custom_template_name,
            row.created_by_user_id if row.created_by_user_id else None,
            row.created_by_user_name if row.created_by_user_name else None,
            row.num_reports_using_setting,
        ])
        
    service = DocumentService()
    doc = service.upload_and_create("affected_rows.csv", output.getvalue(), 0)
    url = service.get_url(doc.id, ttl_seconds=60*60)
    print("Document URL:", url)
    

def main():
    start_time = time.time()
    data = []
    settings = FinancialPackageTemplateSettings.objects.all()
    num_settings = settings.count()
    print(f"Num Settings: {num_settings}")

    for setting in settings:
        sections_to_include = setting.financial_package_statements_to_include.order_by('index')
        num_sections_to_include = sections_to_include.count()
        cap_accounts_section = sections_to_include.filter(type='partner_capital_account_statements')

        # Skip setting without cap_accounts_section
        if not cap_accounts_section.count():
            continue

        cap_accounts_index = cap_accounts_section[0].index

        # Skip settings where cap_accounts_section is already the last one
        if cap_accounts_index == num_sections_to_include - 1:
            continue

        # Update indices for remaining settings
        print(f"Working on Setting Label {setting.label}")
        rows_for_custom_template = convert_setting_to_domain_class(setting)
        data.append(rows_for_custom_template)


    flat_data = [item for sublist in data for item in sublist]
    save_data_to_csv(flat_data)
    print(f'Job completed in: {time.time() - start_time}')
    
main()