Draft: add custom 'migrateredis' management command

cvat-ai · Jan 3, 2025 · 7a32e8b · 7a32e8b
1 parent 26f08ce
commit 7a32e8b
Show file tree

Hide file tree

Showing 7 changed files with 95 additions and 0 deletions.
diff --git a/backend_entrypoint.sh b/backend_entrypoint.sh
@@ -21,6 +21,7 @@ cmd_init() {
 
     wait-for-it "${CVAT_REDIS_INMEM_HOST}:${CVAT_REDIS_INMEM_PORT:-6379}" -t 0
     ~/manage.py syncperiodicjobs
+    ~/manage.py migrateredis
 }
 
 cmd_run() {

diff --git a/cvat/apps/engine/management/commands/migrateredis.py b/cvat/apps/engine/management/commands/migrateredis.py
@@ -0,0 +1,50 @@
+# Copyright (C) 2025 CVAT.ai Corporation
+#
+# SPDX-License-Identifier: MIT
+
+import importlib.util as importlib_util
+import sys
+from pathlib import Path
+from typing import cast
+
+from django.conf import settings
+from django.core.management.base import BaseCommand
+
+from cvat.apps.engine.models import RedisMigration
+from cvat.apps.engine.redis_migrations import BaseMigration
+
+
+def get_migration_class(module_name: str, file_path: Path) -> BaseMigration:
+    spec = importlib_util.spec_from_file_location(module_name, file_path)
+    module = importlib_util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    MigrationClass = getattr(module, "Migration", None)
+    if not MigrationClass or not issubclass(MigrationClass, BaseMigration):
+        raise Exception(f"Invalid migration: {file_path}")
+
+    return MigrationClass
+
+
+class Command(BaseCommand):
+    help = "Applies Redis migrations and records them in the database"
+
+    def handle(self, *args, **options) -> None:
+        migrations_dir = Path(settings.REDIS_MIGRATIONS_ROOT)
+        applied_migrations = RedisMigration.objects.all().values_list("name")
+
+        for migration_file in sorted(migrations_dir.glob("[0-9]*.py")):
+            migration_name = migration_file.stem
+            if migration_name in applied_migrations:
+                continue
+
+            migration = get_migration_class(module_name=migration_name, file_path=migration_file)
+            try:
+                migration.run()
+                RedisMigration.objects.create(name=migration_name)
+                self.stdout.write(
+                    self.style.SUCCESS(f"Successfully applied migration: {migration_name}")
+                )
+            except Exception as e:
+                self.stderr.write(self.style.ERROR(f"Failed to apply migration: {migration_name}"))
+                self.stderr.write(str(e))
+                break
diff --git a/cvat/apps/engine/migrations/0087_redismigration.py b/cvat/apps/engine/migrations/0087_redismigration.py
@@ -0,0 +1,26 @@
+# Generated by Django 4.2.16 on 2025-01-03 17:08
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("engine", "0086_profile_has_analytics_access"),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name="RedisMigration",
+            fields=[
+                (
+                    "id",
+                    models.AutoField(
+                        auto_created=True, primary_key=True, serialize=False, verbose_name="ID"
+                    ),
+                ),
+                ("name", models.CharField(max_length=256)),
+                ("applied_date", models.DateTimeField(auto_now_add=True)),
+            ],
+        ),
+    ]
diff --git a/cvat/apps/engine/models.py b/cvat/apps/engine/models.py
@@ -1345,3 +1345,9 @@ class RequestSubresource(TextChoices):
     ANNOTATIONS = "annotations"
     DATASET = "dataset"
     BACKUP = "backup"
+
+
+class RedisMigration(models.Model):
+    # todo: redis_inmem/redis_ondisk
+    name = models.CharField(max_length=256)
+    applied_date = models.DateTimeField(auto_now_add=True)
diff --git a/cvat/apps/engine/redis_migrations/__init__.py b/cvat/apps/engine/redis_migrations/__init__.py
@@ -0,0 +1,8 @@
+from abc import ABCMeta, abstractmethod
+
+
+class BaseMigration(metaclass=ABCMeta):
+
+    @staticmethod
+    @abstractmethod
+    def run() -> None: ...
diff --git a/cvat/settings/base.py b/cvat/settings/base.py
@@ -407,6 +407,8 @@ class CVAT_QUEUES(Enum):
 # Make sure to update other config files when updating these directories
 DATA_ROOT = os.path.join(BASE_DIR, 'data')
 
+REDIS_MIGRATIONS_ROOT = os.path.join(BASE_DIR, 'cvat', 'apps', 'engine', 'redis_migrations')
+
 MEDIA_DATA_ROOT = os.path.join(DATA_ROOT, 'data')
 os.makedirs(MEDIA_DATA_ROOT, exist_ok=True)
 

diff --git a/dev/format_python_code.sh b/dev/format_python_code.sh
@@ -34,6 +34,8 @@ for paths in \
     "cvat/apps/dataset_manager/tests/test_annotation.py" \
     "cvat/apps/dataset_manager/tests/utils.py" \
     "cvat/apps/events/signals.py" \
+    "cvat/apps/engine/management/commands/migrateredis.py" \
+    "cvat/apps/engine/redis_migrations/*.py" \
     ; do
     ${BLACK} -- ${paths}
     ${ISORT} -- ${paths}