HouseWatch/housewatch/api/async_migration.py at 9b126c0d588c1dc33045dfea3dd76f8abf888d29 · PostHog/HouseWatch · GitHub

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
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
import structlog
from rest_framework import serializers, viewsets
from rest_framework.decorators import action
from housewatch.models.async_migration import AsyncMigration, MigrationStatus
from housewatch.celery import run_async_migration
from rest_framework.response import Response


logger = structlog.get_logger(__name__)


class AsyncMigrationSerializer(serializers.ModelSerializer):
    class Meta:
        model = AsyncMigration
        fields = [
            "id",
            "name",
            "description",
            "progress",
            "status",
            "current_operation_index",
            "current_query_id",
            "task_id",
            "started_at",
            "finished_at",
            "operations",
            "rollback_operations",
            "last_error",
        ]
        read_only_fields = [
            "id",
            "progress",
            "status",
            "current_operation_index",
            "current_query_id",
            "task_id",
            "started_at",
            "finished_at",
            "last_error",
        ]

    def create(self, validated_data):
        validated_data["progress"] = 0
        validated_data["current_operation_index"] = 0
        validated_data["status"] = MigrationStatus.NotStarted
        return super().create(validated_data)


class AsyncMigrationsViewset(viewsets.ModelViewSet):
    queryset = AsyncMigration.objects.all().order_by("name")
    serializer_class = AsyncMigrationSerializer

    @action(methods=["POST"], detail=True)
    def trigger(self, request, **kwargs):

        migration = self.get_object()

        migration.status = MigrationStatus.Starting
        migration.save()

        run_async_migration.delay(migration.name)
        return Response({"success": True}, status=200)

    # @action(methods=["GET"], detail=False)
    # def test(self, request, **kwargs):
    #     simple.delay()
    #     return Response()

    # def _force_stop(self, rollback: bool):
    #     migration = self.get_object()
    #     if migration.status not in [MigrationStatus.Running, MigrationStatus.Starting]:
    #         return response.Response(
    #             {"success": False, "error": "Can't stop a migration that isn't running."}, status=400
    #         )
    #     force_stop_migration(migration, rollback=rollback)
    #     return response.Response({"success": True}, status=200)

    # # DANGEROUS! Can cause another task to be lost
    # @action(methods=["POST"], detail=True)
    # def force_stop(self, request, **kwargs):
    #     return self._force_stop(rollback=True)

    # # DANGEROUS! Can cause another task to be lost
    # @action(methods=["POST"], detail=True)
    # def force_stop_without_rollback(self, request, **kwargs):
    #     return self._force_stop(rollback=False)

    # @action(methods=["POST"], detail=True)
    # def rollback(self, request, **kwargs):
    #     migration = self.get_object()
    #     if migration.status != MigrationStatus.Errored:
    #         return response.Response(
    #             {"success": False, "error": "Can't rollback a migration that isn't in errored state."}, status=400
    #         )

    #     rollback_migration(migration)
    #     return response.Response({"success": True}, status=200)

    # @action(methods=["POST"], detail=True)
    # def force_rollback(self, request, **kwargs):
    #     migration = self.get_object()
    #     if migration.status != MigrationStatus.CompletedSuccessfully:
    #         return response.Response(
    #             {"success": False, "error": "Can't force rollback a migration that did not complete successfully."},
    #             status=400,
    #         )

    #     rollback_migration(migration)
    #     return response.Response({"success": True}, status=200)

    # @action(methods=["GET"], detail=True)
    # def errors(self, request, **kwargs):
    #     migration = self.get_object()
    #     return response.Response(
    #         [
    #             AsyncMigrationErrorsSerializer(e).data
    #             for e in AsyncMigrationError.objects.filter(async_migration=migration).order_by("-created_at")
    #         ]
    #     )