# 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))