init

Bin

2025-12-16 9e0b2ba2c317b1a86212f24cbae3195ad1f3dbfa

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
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
# Generated by Django 5.1.10 on 2025-08-07 16:14
 
import logging
 
from django.conf import settings
from django.db import migrations
 
from core.redis import start_job_async_or_sync
from core.models import AsyncMigrationStatus
from core.utils.common import btree_gin_migration_operations
 
logger = logging.getLogger(__name__)
 
IS_SQLITE = settings.DJANGO_DB == settings.DJANGO_DB_SQLITE
 
migration_name = "0056_prediction_result_proj_gin_idx_async"
 
SQL_CREATE_INDEX = (
    "CREATE INDEX CONCURRENTLY IF NOT EXISTS tasks_predictions_result_proj_gin "
    "ON prediction USING GIN (project_id, CAST(result AS text) gin_trgm_ops);"
)
 
SQL_DROP_INDEX = "DROP INDEX CONCURRENTLY IF EXISTS tasks_predictions_result_proj_gin;"
 
def _forward(migration_name: str, db_alias: str):
    """Create the GIN index inside a dedicated job."""
    # If the migration has already been executed, do nothing
    migration, created = AsyncMigrationStatus.objects.using(db_alias).get_or_create(
        name=migration_name,
        defaults={"status": AsyncMigrationStatus.STATUS_STARTED},
    )
    if not created:
        logger.info("Migration %s already executed", migration_name)
        return
 
    logger.info("Starting async migration %s", migration_name)
    from django.db import connections
 
    with connections[db_alias].cursor() as cursor:
        cursor.execute(SQL_CREATE_INDEX)
    migration.status = AsyncMigrationStatus.STATUS_FINISHED
    migration.save(using=db_alias)
    logger.info("Async migration %s complete", migration_name)
 
 
def _backward(migration_name: str, db_alias: str):
    """Revert the GIN index creation."""
    migration = AsyncMigrationStatus.objects.using(db_alias).create(
        name=migration_name,
        status=AsyncMigrationStatus.STATUS_STARTED,
    )
    logger.info("Reverting async migration %s", migration_name)
    from django.db import connections
 
    with connections[db_alias].cursor() as cursor:
        cursor.execute(SQL_DROP_INDEX)
    migration.status = AsyncMigrationStatus.STATUS_FINISHED
    migration.save(using=db_alias)
    logger.info("Revert of async migration %s complete", migration_name)
 
 
def forwards(apps, schema_editor):
    if IS_SQLITE:
        logger.info("SQLite detected; skipping GIN index creation")
        return
 
    # Only run on PostgreSQL
    if not schema_editor.connection.vendor.startswith("postgres"):
        logger.info("Database vendor: %s. Skipping index creation", schema_editor.connection.vendor)
        return
 
    db_alias = schema_editor.connection.alias
    start_job_async_or_sync(_forward, migration_name=migration_name, db_alias=db_alias)
 
 
def backwards(apps, schema_editor):
    if IS_SQLITE:
        logger.info("SQLite detected; skipping GIN index drop")
        return
 
    if not schema_editor.connection.vendor.startswith("postgres"):
        logger.info("Database vendor: %s. Skipping index drop", schema_editor.connection.vendor)
        return
 
    db_alias = schema_editor.connection.alias
    start_job_async_or_sync(_backward, migration_name=migration_name, db_alias=db_alias)
 
 
 
class Migration(migrations.Migration):
    atomic = False
 
    dependencies = [
        ("tasks", "0055_task_proj_octlen_idx_async"),
    ]
 
    operations = btree_gin_migration_operations(migrations.RunPython(forwards, backwards))