-
Notifications
You must be signed in to change notification settings - Fork 48
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat: Added management command to update UIDs
- Loading branch information
1 parent
7ca0731
commit 7e95715
Showing
1 changed file
with
138 additions
and
0 deletions.
There are no files selected for viewing
138 changes: 138 additions & 0 deletions
138
enterprise/management/commands/update_ccb_enterprise_uids.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,138 @@ | ||
import logging | ||
import csv | ||
from django.db import transaction | ||
from django.core.management.base import BaseCommand | ||
from django.core.exceptions import ValidationError | ||
from social_django.models import UserSocialAuth | ||
|
||
logger = logging.getLogger(__name__) | ||
|
||
|
||
class CSVUpdateError(Exception): | ||
"""Custom exception for CSV update process.""" | ||
pass | ||
|
||
|
||
class Command(BaseCommand): | ||
""" | ||
Update the CCB enterprise related social auth records UID to the new one. | ||
Example usage: | ||
./manage.py lms update_ccb_enterprise_uids csv_file_path | ||
""" | ||
|
||
help = 'Robust record update from CSV with console logging' | ||
|
||
def add_arguments(self, parser): | ||
parser.add_argument('csv_file', type=str, help='Path to the CSV file') | ||
parser.add_argument( | ||
'--no-dry-run', | ||
action='store_false', | ||
dest='dry_run', | ||
default=True, | ||
help='Actually save changes instead of simulating' | ||
) | ||
|
||
def handle(self, *args, **options): | ||
logger.info("Command has started...") | ||
csv_path = options['csv_file'] | ||
dry_run = options['dry_run'] | ||
|
||
total_processed = 0 | ||
total_updated = 0 | ||
total_errors = 0 | ||
|
||
try: | ||
with open(csv_path, 'r') as csvfile: | ||
reader = csv.DictReader(csvfile) | ||
|
||
for row_num, row in enumerate(reader, start=1): | ||
total_processed += 1 | ||
|
||
try: | ||
# Use a separate transaction for each record | ||
with transaction.atomic(): | ||
if self.update_record(row, dry_run): | ||
total_updated += 1 | ||
|
||
except Exception as row_error: | ||
total_errors += 1 | ||
error_msg = f"Row {row_num} update failed: {row} - Error: {str(row_error)}" | ||
logger.error(error_msg, exc_info=True) | ||
|
||
# Summary logging | ||
summary_msg = ( | ||
f"CSV Update Summary:\n" | ||
f"Total Records Processed: {total_processed}\n" | ||
f"Records Successfully Updated: {total_updated}\n" | ||
f"Errors Encountered: {total_errors}\n" | ||
f"Dry Run Mode: {'Enabled' if dry_run else 'Disabled'}" | ||
) | ||
logger.info(summary_msg) | ||
except IOError as io_error: | ||
logger.critical(f"File I/O error: {str(io_error)}") | ||
|
||
except Exception as e: | ||
logger.critical(f"Critical error in CSV processing: {str(e)}") | ||
|
||
def update_record(self, row, dry_run=True): | ||
""" | ||
Update a single record with detailed error handling. | ||
Args: | ||
row (dict): CSV row data | ||
dry_run (bool): Whether to simulate or actually save changes | ||
Returns: | ||
bool: Whether the update was successful | ||
""" | ||
try: | ||
# Validate required fields are present | ||
old_uid = row.get('old-uid') | ||
new_uid = row.get('new-uid') | ||
|
||
if not old_uid or not new_uid: | ||
raise CSVUpdateError("Missing required UID fields") | ||
|
||
# Construct full UIDs with proper prefixes | ||
old_uid_with_prefix = f'sf-ccb:{old_uid}' | ||
new_uid_with_prefix = f'ccb:samlp|ccb-c52b6697-57da-492b-8d17-7d5f0c0290f2|{new_uid}@ccb.org.co' | ||
|
||
# Retrieve instances | ||
instance_with_old_uid = UserSocialAuth.objects.filter(uid=old_uid_with_prefix).first() | ||
|
||
if not instance_with_old_uid: | ||
raise CSVUpdateError(f"No record found with old UID {old_uid_with_prefix}") | ||
|
||
# Optional: Check for conflicting new UID | ||
instance_with_new_uid = UserSocialAuth.objects.filter(uid=new_uid_with_prefix).first() | ||
if instance_with_new_uid: | ||
log_entry = f"Warning: Existing record with new UID {new_uid_with_prefix} is deleting." | ||
logger.info(log_entry) | ||
if not dry_run: | ||
instance_with_new_uid.delete() | ||
|
||
if not dry_run: | ||
instance_with_old_uid.uid = new_uid_with_prefix | ||
instance_with_old_uid.save() | ||
|
||
# Logging | ||
log_entry = f"Successfully updated record: Old UID {old_uid_with_prefix} → New UID {new_uid_with_prefix}" | ||
logger.info(log_entry) | ||
|
||
return True | ||
|
||
except ValidationError as ve: | ||
error_msg = f"Validation error: {ve}" | ||
logger.error(error_msg) | ||
raise | ||
|
||
except CSVUpdateError as update_error: | ||
error_msg = f"Update processing error: {update_error}" | ||
logger.error(error_msg) | ||
raise | ||
|
||
except Exception as e: | ||
error_msg = f"Unexpected error during record update: {e}" | ||
logger.error(error_msg, exc_info=True) | ||
raise |