From ef990099bdd029990aa3ddaa83a0898cc6798723 Mon Sep 17 00:00:00 2001
From: Jo Booth <jo.m.booth@gmail.com>
Date: Fri, 1 Nov 2024 15:50:14 -0400
Subject: [PATCH] feat: LEAP-1602: prereq for updating existing locks

---
 .../migrations/0051_tasklock_created_at.py    | 18 ++++++
 .../migrations/0052_auto_20241101_1941.py     | 55 +++++++++++++++++++
 label_studio/tasks/models.py                  | 11 ++++
 3 files changed, 84 insertions(+)
 create mode 100644 label_studio/tasks/migrations/0051_tasklock_created_at.py
 create mode 100644 label_studio/tasks/migrations/0052_auto_20241101_1941.py

diff --git a/label_studio/tasks/migrations/0051_tasklock_created_at.py b/label_studio/tasks/migrations/0051_tasklock_created_at.py
new file mode 100644
index 000000000000..3199174dc129
--- /dev/null
+++ b/label_studio/tasks/migrations/0051_tasklock_created_at.py
@@ -0,0 +1,18 @@
+# Generated by Django 4.2.16 on 2024-10-31 23:50
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('tasks', '0050_alter_predictionmeta_failed_prediction_and_more'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='tasklock',
+            name='created_at',
+            field=models.DateTimeField(blank=True, default=None, help_text='Creation time', null=True, verbose_name='created_at'),
+        ),
+    ]
diff --git a/label_studio/tasks/migrations/0052_auto_20241101_1941.py b/label_studio/tasks/migrations/0052_auto_20241101_1941.py
new file mode 100644
index 000000000000..a6e6a25f8ffb
--- /dev/null
+++ b/label_studio/tasks/migrations/0052_auto_20241101_1941.py
@@ -0,0 +1,55 @@
+from django.db import migrations
+from django.db import connection
+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__)
+migration_name = '0052_auto_20241101_1941'
+
+if connection.vendor == 'sqlite':
+    sql_update_created_at = """
+    UPDATE tasks_tasklock
+    SET created_at = datetime(expire_at, %s);
+    """
+    sql_params = (f'-{settings.TASK_LOCK_TTL} seconds',)
+else:
+    sql_update_created_at = """
+    UPDATE tasks_tasklock
+    SET created_at = expire_at - INTERVAL %s;
+    """
+    sql_params = ('%s seconds' % settings.TASK_LOCK_TTL,)
+
+def forward_migration(migration_name):
+    migration = AsyncMigrationStatus.objects.create(
+        name=migration_name,
+        status=AsyncMigrationStatus.STATUS_STARTED,
+    )
+    logger.info(f'Start async migration {migration_name}')
+
+    with connection.cursor() as cursor:
+        cursor.execute(sql_update_created_at, sql_params)
+
+    migration.status = AsyncMigrationStatus.STATUS_FINISHED
+    migration.save()
+    logger.info(f'Async migration {migration_name} complete')
+
+def forwards(apps, schema_editor):
+    # Dispatch migrations to rqworkers
+    start_job_async_or_sync(forward_migration, migration_name=migration_name)
+
+def backwards(apps, schema_editor):
+    pass
+
+class Migration(migrations.Migration):
+    atomic = False
+
+    dependencies = [
+        ('tasks', '0051_tasklock_created_at'),
+    ]
+
+    operations = [
+        migrations.RunPython(forwards, backwards),
+    ]
diff --git a/label_studio/tasks/models.py b/label_studio/tasks/models.py
index 8c55350e2f6f..0b8b52cc78ea 100644
--- a/label_studio/tasks/models.py
+++ b/label_studio/tasks/models.py
@@ -790,6 +790,8 @@ class TaskLock(models.Model):
         related_name='locks',
         help_text='Locked task',
     )
+    created_at = models.DateTimeField(_('created_at'), null=True, default=None, blank=True, help_text='Creation time')
+
     expire_at = models.DateTimeField(_('expire_at'))
     unique_id = models.UUIDField(default=uuid.uuid4, null=True, blank=True, unique=True, editable=False)
     user = models.ForeignKey(
@@ -799,6 +801,15 @@ class TaskLock(models.Model):
         help_text='User who locked this task',
     )
 
+    def save(self, *args, update_fields=None, **kwargs):
+        if update_fields is not None:
+            update_fields = {'created_at'}.union(update_fields)
+
+        if not self.pk:
+            self.created_at = now()
+
+        return super().save(*args, update_fields=update_fields, **kwargs)
+
 
 class AnnotationDraft(models.Model):
     result = JSONField(_('result'), help_text='Draft result in JSON format')