1+ import csv
2+
3+ from django .contrib .auth import get_user_model
4+ from django .core .management .base import BaseCommand , CommandError
5+
6+
7+ class Command (BaseCommand ):
8+ help = "Archive users (lock accounts) from a CSV file."
9+
10+ def add_arguments (self , parser ):
11+ parser .add_argument (
12+ "--filename" ,
13+ dest = "filename" ,
14+ type = str ,
15+ required = True ,
16+ help = "Filename to import data from." ,
17+ )
18+ parser .add_argument (
19+ "--column-header" ,
20+ dest = "header" ,
21+ default = "username" ,
22+ type = str ,
23+ help = "Header associated with the identifier column in the CSV." ,
24+ )
25+ parser .add_argument (
26+ "--lookup-field" ,
27+ dest = "lookup_field" ,
28+ default = "username" ,
29+ type = str ,
30+ help = "User model field used to match identifiers (ex: username, student_id, email)." ,
31+ )
32+ parser .add_argument (
33+ "--run" ,
34+ action = "store_true" ,
35+ dest = "run" ,
36+ help = "Actually run." ,
37+ )
38+ parser .add_argument (
39+ "--confirm" ,
40+ action = "store_true" ,
41+ dest = "confirm" ,
42+ help = "Skip confirmation prompt (only applies with --run)." ,
43+ )
44+
45+ def ask (self , q ) -> bool :
46+ return input (f"{ q } [y/N]: " ).strip ().lower () == "y"
47+
48+ def read_identifiers (self , filename : str , column_header : str ):
49+ identifiers = []
50+ try :
51+ with open (filename , encoding = "utf-8" ) as csvfile :
52+ reader = csv .DictReader (csvfile )
53+ for row in reader :
54+ value = row .get (column_header )
55+ if value :
56+ identifiers .append (value .strip ())
57+ except FileNotFoundError as e :
58+ raise CommandError (f"File not found: { filename } " ) from e
59+ except OSError as e :
60+ raise CommandError (f"Error reading file: { e } " ) from e
61+ return identifiers
62+
63+ def handle (self , * args , ** options ):
64+ identifiers = self .read_identifiers (options ["filename" ], options ["header" ])
65+ if not identifiers :
66+ self .stdout .write (self .style .WARNING ("No identifiers found in the CSV." ))
67+ return
68+
69+ total_identifiers = len (identifiers )
70+ unique_identifiers = list (dict .fromkeys (identifiers ))
71+ duplicate_count = total_identifiers - len (unique_identifiers )
72+
73+ user_model = get_user_model ()
74+ lookup_field = options ["lookup_field" ]
75+ valid_fields = {field .name for field in user_model ._meta .get_fields ()}
76+ if lookup_field not in valid_fields :
77+ raise CommandError (f"Invalid lookup field: { lookup_field } " )
78+ lookup_key = f"{ lookup_field } __in"
79+
80+ users = user_model .objects .filter (** {lookup_key : unique_identifiers })
81+ found_values = set (users .values_list (lookup_field , flat = True ))
82+ missing = sorted (value for value in unique_identifiers if value not in found_values )
83+
84+ self .stdout .write (f"Identifiers provided: { total_identifiers } " )
85+ self .stdout .write (f"Unique identifiers: { len (unique_identifiers )} " )
86+ if duplicate_count :
87+ self .stdout .write (self .style .WARNING (f"Duplicate identifiers: { duplicate_count } " ))
88+ self .stdout .write (f"Matched users: { users .count ()} " )
89+ self .stdout .write (f"Missing identifiers: { len (missing )} " )
90+ if missing :
91+ self .stdout .write (self .style .WARNING (str (missing )))
92+
93+ if not options ["run" ]:
94+ self .stdout .write ("Dry run mode." )
95+ return
96+
97+ if not options ["confirm" ]:
98+ if not self .ask (
99+ "This script will archive users (lock accounts). Ensure\n "
100+ "that you know what you're doing before proceeding.\n \n "
101+ "Continue?"
102+ ):
103+ self .stdout .write ("Aborted." )
104+ return
105+
106+ archived_count , already_archived_count = user_model .archive_users (users )
107+ self .stdout .write (self .style .SUCCESS (f"Archived users: { archived_count } " ))
108+ if already_archived_count :
109+ self .stdout .write (self .style .WARNING (f"Already archived users: { already_archived_count } " ))
0 commit comments