from __future__ import annotations from dataclasses import dataclass import html import re from pathlib import Path import zipfile from django.core.management.base import BaseCommand, CommandError from django.db import transaction from xml.etree import ElementTree as ET from apps.accounts.models import User @dataclass class BackfillResult: updated: int = 0 skipped_missing_user: int = 0 skipped_has_referrer: int = 0 skipped_missing_sponsor: int = 0 skipped_missing_referral_code: int = 0 skipped_missing_mobile: int = 0 skipped_self_referral: int = 0 skipped_conflicting_sponsor: int = 0 class Command(BaseCommand): help = ( "Backfill missing referred_by relationships for existing users by matching each Excel row to a system user " "and resolving the sponsor legacy ID to a sponsor referral code." ) def add_arguments(self, parser): parser.add_argument( "--file", required=True, help="Absolute path to the legacy Excel file (.xlsx).", ) parser.add_argument( "--sheet", default=None, help="Optional sheet name. Defaults to the active sheet.", ) parser.add_argument( "--dry-run", action="store_true", default=False, help="Report the changes without saving them.", ) @staticmethod def _clean_string(value) -> str: if value is None: return "" return str(value).strip() @staticmethod def _clean_phone(value) -> str: raw = Command._clean_string(value) if not raw: return "" if raw.endswith(".0") and raw.replace(".", "", 1).isdigit(): raw = raw[:-2] digits = "".join(ch for ch in raw if ch.isdigit()) return digits if digits else raw @staticmethod def _read_zip_text(zf: zipfile.ZipFile, name: str) -> str: original_update_crc = zipfile.ZipExtFile._update_crc zipfile.ZipExtFile._update_crc = lambda self, newdata: None try: return zf.read(name).decode("utf-8", errors="ignore") finally: zipfile.ZipExtFile._update_crc = original_update_crc @staticmethod def _attr_value(attrs: str, name: str) -> str: match = re.search(rf'{name}="([^"]*)"', attrs) return match.group(1) if match else "" def _resolve_sheet_path(self, zf: zipfile.ZipFile, sheet_name: str | None) -> str: workbook_xml = self._read_zip_text(zf, "xl/workbook.xml") workbook_rels = self._read_zip_text(zf, "xl/_rels/workbook.xml.rels") workbook_root = ET.fromstring(workbook_xml) rels_root = ET.fromstring(workbook_rels) namespaces = { "main": "http://schemas.openxmlformats.org/spreadsheetml/2006/main", "r": "http://schemas.openxmlformats.org/officeDocument/2006/relationships", "rel": "http://schemas.openxmlformats.org/package/2006/relationships", } requested_rid = None first_rid = None for sheet in workbook_root.findall("main:sheets/main:sheet", namespaces): rid = sheet.attrib.get("{http://schemas.openxmlformats.org/officeDocument/2006/relationships}id") if first_rid is None: first_rid = rid if sheet_name and sheet.attrib.get("name") == sheet_name: requested_rid = rid break rid = requested_rid or first_rid if not rid: raise CommandError("Could not resolve a worksheet in the workbook.") target = None for rel in rels_root.findall("rel:Relationship", namespaces): if rel.attrib.get("Id") == rid: target = rel.attrib.get("Target") break if not target: raise CommandError("Could not resolve the worksheet relationship target.") if target.startswith("/"): target = target.lstrip("/") if not target.startswith("xl/"): target = f"xl/{target}" return target def _load_shared_strings(self, zf: zipfile.ZipFile) -> list[str]: if "xl/sharedStrings.xml" not in zf.namelist(): return [] shared_xml = self._read_zip_text(zf, "xl/sharedStrings.xml") try: root = ET.fromstring(shared_xml) except ET.ParseError: return [] namespace = {"main": "http://schemas.openxmlformats.org/spreadsheetml/2006/main"} values: list[str] = [] for item in root.findall("main:si", namespace): texts = [node.text or "" for node in item.findall(".//main:t", namespace)] values.append("".join(texts)) return values def _extract_cell_value(self, cell_xml: str, shared_strings: list[str]) -> tuple[str, str]: attrs_match = re.match(r"]*)>", cell_xml) attrs = attrs_match.group(1) if attrs_match else "" ref = self._attr_value(attrs, "r") cell_type = self._attr_value(attrs, "t") if cell_type == "inlineStr": match = re.search(r"]*>(.*?)", cell_xml, flags=re.S) return ref, html.unescape(match.group(1)) if match else "" if cell_type == "s": match = re.search(r"(.*?)", cell_xml, flags=re.S) if not match: return ref, "" index_text = html.unescape(match.group(1)).strip() if index_text.isdigit(): index = int(index_text) if 0 <= index < len(shared_strings): return ref, shared_strings[index] return ref, index_text match = re.search(r"(.*?)", cell_xml, flags=re.S) return ref, html.unescape(match.group(1)) if match else "" def _iter_sheet_rows(self, sheet_xml: str, shared_strings: list[str]): for row_match in re.finditer(r']*r="(\d+)"[^>]*>(.*?)', sheet_xml, flags=re.S): row_num = int(row_match.group(1)) row_xml = row_match.group(2) row_values: dict[str, str] = {} for cell_match in re.finditer(r']*>.*?', row_xml, flags=re.S): ref, value = self._extract_cell_value(cell_match.group(0), shared_strings) if ref: column = re.match(r"([A-Z]+)", ref) row_values[column.group(1) if column else ref] = value yield row_num, row_values def _resolve_sponsor_user( self, rows_by_number: dict[int, dict[str, str]], sponsor_legacy_id: str, id_to_row: dict[str, int], ) -> User | None: if not sponsor_legacy_id: return None sponsor_row = id_to_row.get(sponsor_legacy_id) sponsor_referral = "" sponsor_mobile = "" if sponsor_row: sponsor_values = rows_by_number.get(sponsor_row, {}) sponsor_referral = self._clean_string(sponsor_values.get("D")) sponsor_mobile = self._clean_phone(sponsor_values.get("F")) sponsor_user = None if sponsor_referral: sponsor_user = User.objects.filter(referral_code=sponsor_referral).first() if sponsor_user is None and sponsor_mobile: sponsor_user = User.objects.filter(mobile=sponsor_mobile).first() return sponsor_user def handle(self, *args, **options): file_path = Path(options["file"]) sheet_name = options.get("sheet") dry_run = bool(options.get("dry_run")) if not file_path.exists(): raise CommandError(f"File not found: {file_path}") if file_path.suffix.lower() != ".xlsx": raise CommandError("Only .xlsx files are supported.") with zipfile.ZipFile(file_path) as zf: sheet_path = self._resolve_sheet_path(zf, sheet_name) sheet_xml = self._read_zip_text(zf, sheet_path) shared_strings = self._load_shared_strings(zf) complete_rows = list(self._iter_sheet_rows(sheet_xml, shared_strings)) if not complete_rows: raise CommandError("No readable rows were found in the workbook sheet.") # Expected columns: # A: S.No, B: ID, C: Name, D: Referral code, E: Join Date, F: Phone, G: Sponsor, H: Status rows_by_number = {row_num: row_values for row_num, row_values in complete_rows} id_to_row: dict[str, int] = {} for row_num, row_values in complete_rows: legacy_id = self._clean_string(row_values.get("B")) if legacy_id: id_to_row[legacy_id] = row_num result = BackfillResult() updates: dict[int, tuple[User, User]] = {} for row_num, row_values in complete_rows: if row_num == 1: continue mobile = self._clean_phone(row_values.get("F")) sponsor_legacy_id = self._clean_string(row_values.get("G")) if not mobile: result.skipped_missing_mobile += 1 continue user = User.objects.filter(mobile=mobile).first() if user is None: result.skipped_missing_user += 1 continue if user.referred_by_id: result.skipped_has_referrer += 1 continue sponsor_user = self._resolve_sponsor_user(rows_by_number, sponsor_legacy_id, id_to_row) if sponsor_user is None: if sponsor_legacy_id: result.skipped_missing_sponsor += 1 else: result.skipped_missing_referral_code += 1 continue if sponsor_user.pk == user.pk: result.skipped_self_referral += 1 continue existing = updates.get(user.pk) if existing: existing_sponsor = existing[1] if existing_sponsor.pk != sponsor_user.pk: result.skipped_conflicting_sponsor += 1 continue updates[user.pk] = (user, sponsor_user) if not updates: self.stdout.write("No users needed referred_by backfill.") else: for user, sponsor_user in updates.values(): self.stdout.write( f"Will set {user.mobile} -> referred_by {sponsor_user.mobile or sponsor_user.email} " f"(referral_code={sponsor_user.referral_code})" ) if not dry_run: with transaction.atomic(): for user, sponsor_user in updates.values(): user.referred_by = sponsor_user user.save(update_fields=["referred_by", "updated_at"]) result.updated = len(updates) self.stdout.write(self.style.SUCCESS("Backfill complete.")) self.stdout.write(f"Updated: {result.updated if not dry_run else len(updates)}") self.stdout.write(f"Skipped missing user: {result.skipped_missing_user}") self.stdout.write(f"Skipped already linked: {result.skipped_has_referrer}") self.stdout.write(f"Skipped missing sponsor: {result.skipped_missing_sponsor}") self.stdout.write(f"Skipped missing mobile: {result.skipped_missing_mobile}") self.stdout.write(f"Skipped self referral: {result.skipped_self_referral}") self.stdout.write(f"Skipped conflicting sponsor: {result.skipped_conflicting_sponsor}") self.stdout.write(f"Dry run: {dry_run}")