|
| 1 | +from django.db import migrations, connection |
| 2 | +from copy import deepcopy |
| 3 | +from django.apps import apps as django_apps |
| 4 | +from django.conf import settings |
| 5 | +from core.models import AsyncMigrationStatus |
| 6 | +from core.redis import start_job_async_or_sync |
| 7 | +from core.utils.iterators import iterate_queryset |
| 8 | +import logging |
| 9 | + |
| 10 | +migration_name = '0017_update_agreement_selected_to_nested_structure' |
| 11 | + |
| 12 | +logger = logging.getLogger(__name__) |
| 13 | + |
| 14 | + |
| 15 | +def forward_migration(): |
| 16 | + """ |
| 17 | + Migrates views that have agreement_selected populated to the new structure |
| 18 | +
|
| 19 | + Old structure: |
| 20 | + 'agreement_selected': { |
| 21 | + 'annotators': List[int] |
| 22 | + 'models': List[str] |
| 23 | + 'ground_truth': bool |
| 24 | + } |
| 25 | +
|
| 26 | + New structure: |
| 27 | + 'agreement_selected': { |
| 28 | + 'annotators': { |
| 29 | + 'all': bool |
| 30 | + 'ids': List[int] |
| 31 | + }, |
| 32 | + 'models': { |
| 33 | + 'all': bool |
| 34 | + 'ids': List[str] |
| 35 | + }, |
| 36 | + 'ground_truth': bool |
| 37 | + } |
| 38 | + """ |
| 39 | + migration, created = AsyncMigrationStatus.objects.get_or_create( |
| 40 | + name=migration_name, |
| 41 | + defaults={'status': AsyncMigrationStatus.STATUS_STARTED} |
| 42 | + ) |
| 43 | + if not created: |
| 44 | + return # already in progress or done |
| 45 | + |
| 46 | + # Look up models at runtime inside the worker process |
| 47 | + View = django_apps.get_model('data_manager', 'View') |
| 48 | + |
| 49 | + # Iterate using values() to avoid loading full model instances |
| 50 | + # Fetch only the fields we need, filtering to views that have 'agreement_selected' in data |
| 51 | + qs = ( |
| 52 | + View.objects |
| 53 | + .filter(data__has_key='agreement_selected') |
| 54 | + .filter(data__agreement_selected__isnull=False) |
| 55 | + .values('id', 'data') |
| 56 | + ) |
| 57 | + |
| 58 | + updated = 0 |
| 59 | + for row in qs: |
| 60 | + view_id = row['id'] |
| 61 | + data = row.get('data') or {} |
| 62 | + |
| 63 | + new_data = deepcopy(data) |
| 64 | + # Always use the new nested structure |
| 65 | + new_data['agreement_selected'] = { |
| 66 | + 'annotators': {'all': True, 'ids': []}, |
| 67 | + 'models': {'all': True, 'ids': []}, |
| 68 | + 'ground_truth': False |
| 69 | + } |
| 70 | + |
| 71 | + # Update only the JSON field via update(); do not load model instance or call save() |
| 72 | + View.objects.filter(id=view_id).update(data=new_data) |
| 73 | + logger.info(f'Updated View {view_id} agreement selected to default all annotators + all models') |
| 74 | + updated += 1 |
| 75 | + |
| 76 | + if updated: |
| 77 | + logger.info(f'{migration_name} Updated {updated} View rows') |
| 78 | + |
| 79 | + migration.status = AsyncMigrationStatus.STATUS_FINISHED |
| 80 | + migration.save(update_fields=['status']) |
| 81 | + |
| 82 | +def forwards(apps, schema_editor): |
| 83 | + start_job_async_or_sync(forward_migration, queue_name=settings.SERVICE_QUEUE_NAME) |
| 84 | + |
| 85 | + |
| 86 | +def backwards(apps, schema_editor): |
| 87 | + # Irreversible: we cannot reconstruct the previous annotator lists safely |
| 88 | + pass |
| 89 | + |
| 90 | + |
| 91 | +class Migration(migrations.Migration): |
| 92 | + atomic = False |
| 93 | + |
| 94 | + dependencies = [ |
| 95 | + ('data_manager', '0016_migrate_agreement_selected_annotators_to_unique') |
| 96 | + ] |
| 97 | + |
| 98 | + operations = [ |
| 99 | + migrations.RunPython(forwards, backwards), |
| 100 | + ] |
| 101 | + |
| 102 | + |
| 103 | + |
0 commit comments