Skip to content

Commit

Permalink
feat: Added management command to update UIDs
Browse files Browse the repository at this point in the history
  • Loading branch information
zamanafzal committed Dec 9, 2024
1 parent 7ca0731 commit 7e95715
Showing 1 changed file with 138 additions and 0 deletions.
138 changes: 138 additions & 0 deletions enterprise/management/commands/update_ccb_enterprise_uids.py
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

0 comments on commit 7e95715

Please sign in to comment.