chenzhaoyang
2025-12-17 d3e5a4b7658ece4f845bbc0c4f95acf3fbdf8a61
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
# Generated by Django 5.1.9 on 2025-08-18 12:00
 
from django.db import migrations
from django.conf import settings
from core.models import AsyncMigrationStatus
from core.redis import start_job_async_or_sync
import logging
logger = logging.getLogger(__name__)
 
IS_SQLITE = settings.DJANGO_DB == settings.DJANGO_DB_SQLITE
 
migration_name = '0057_annotation_proj_result_octlen_idx_async'
 
sql_create_index = (
    'CREATE INDEX CONCURRENTLY IF NOT EXISTS annotation_proj_result_octlen_idx '
    'ON task_completion (project_id, octet_length(result::text) DESC) '
    'INCLUDE (id);'
)
sql_drop_index = (
    'DROP INDEX CONCURRENTLY IF EXISTS annotation_proj_result_octlen_idx;'
)
 
def forward_migration(migration_name, db_alias):
    migration, created = AsyncMigrationStatus.objects.using(db_alias).get_or_create(
        name=migration_name,
        defaults={'status': AsyncMigrationStatus.STATUS_STARTED},
    )
    if not created:
        return
    
    logger.info(f'Start async migration {migration_name}')
    from django.db import connections
    cursor = connections[db_alias].cursor()
    cursor.execute(sql_create_index)
    migration.status = AsyncMigrationStatus.STATUS_FINISHED
    migration.save(using=db_alias)
    logger.info(f'Async migration {migration_name} complete')
 
def backward_migration(migration_name, db_alias):
    migration = AsyncMigrationStatus.objects.using(db_alias).create(
        name=migration_name,
        status=AsyncMigrationStatus.STATUS_STARTED,
    )
    logger.info(f'Start revert of async migration {migration_name}')
    from django.db import connections
    cursor = connections[db_alias].cursor()
    cursor.execute(sql_drop_index)
    migration.status = AsyncMigrationStatus.STATUS_FINISHED
    migration.save(using=db_alias)
    logger.info(f'Async migration {migration_name} revert complete')
 
def forwards(apps, schema_editor):
    if IS_SQLITE:
        logger.info('SQLite execution')
        logger.info('Skipping async index creation for non-PostgreSQL databases')
        return
 
    db_alias = schema_editor.connection.alias
    start_job_async_or_sync(forward_migration, migration_name=migration_name, db_alias=db_alias)
 
def backwards(apps, schema_editor):
    db_alias = schema_editor.connection.alias
    start_job_async_or_sync(backward_migration, migration_name=migration_name, db_alias=db_alias)
 
class Migration(migrations.Migration):
    atomic = False
 
    dependencies = [
        ("tasks", "0056_task_prediction_result_proj_gin_idx_async"),
    ]
    operations = [
        migrations.RunPython(forwards, backwards),
    ]