From b50a496d3f4b986caa4b7c6a1ffccc0fdb73bc45 Mon Sep 17 00:00:00 2001
From: Gyubong Lee <jopemachine@naver.com>
Date: Tue, 7 May 2024 08:11:48 +0000
Subject: [PATCH] refactor: Replace vfolder's status_history's type map with
 list

---
 changes/2116.feature.md                       |  1 +
 ...replace_vfolders_status_history_s_type_.py | 67 +++++++++++++++++++
 src/ai/backend/manager/models/vfolder.py      | 30 +++++----
 3 files changed, 84 insertions(+), 14 deletions(-)
 create mode 100644 changes/2116.feature.md
 create mode 100644 src/ai/backend/manager/models/alembic/versions/786be66ef4e5_replace_vfolders_status_history_s_type_.py

diff --git a/changes/2116.feature.md b/changes/2116.feature.md
new file mode 100644
index 00000000000..755eaedf36c
--- /dev/null
+++ b/changes/2116.feature.md
@@ -0,0 +1 @@
+Change the type of `vfolders.status_history` from a mapping of status and timestamps to a list of log entries containing status and timestamps, to preserve the log entries.
diff --git a/src/ai/backend/manager/models/alembic/versions/786be66ef4e5_replace_vfolders_status_history_s_type_.py b/src/ai/backend/manager/models/alembic/versions/786be66ef4e5_replace_vfolders_status_history_s_type_.py
new file mode 100644
index 00000000000..c39b0b6b662
--- /dev/null
+++ b/src/ai/backend/manager/models/alembic/versions/786be66ef4e5_replace_vfolders_status_history_s_type_.py
@@ -0,0 +1,67 @@
+"""Replace vfolders status_history's type map with list
+
+Revision ID: 786be66ef4e5
+Revises: 8c8e90aebacd
+Create Date: 2024-05-07 05:10:23.799723
+
+"""
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "786be66ef4e5"
+down_revision = "8c8e90aebacd"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    op.execute(
+        """
+        WITH data AS (
+            SELECT id,
+                (jsonb_each(status_history)).key AS status,
+                (jsonb_each(status_history)).value AS timestamp
+            FROM vfolders
+        )
+        UPDATE vfolders
+        SET status_history = (
+            SELECT jsonb_agg(
+                jsonb_build_object('status', status, 'timestamp', timestamp)
+            )
+            FROM data
+            WHERE data.id = vfolders.id
+        );
+    """
+    )
+
+    op.execute("UPDATE vfolders SET status_history = '[]'::jsonb WHERE status_history IS NULL;")
+    op.alter_column(
+        "vfolders",
+        "status_history",
+        nullable=False,
+        default=[],
+    )
+
+
+def downgrade():
+    op.execute(
+        """
+        WITH data AS (
+            SELECT id,
+                jsonb_object_agg(
+                    elem->>'status', elem->>'timestamp'
+                ) AS new_status_history
+            FROM vfolders,
+            jsonb_array_elements(status_history) AS elem
+            GROUP BY id
+        )
+        UPDATE vfolders
+        SET status_history = data.new_status_history
+        FROM data
+        WHERE data.id = vfolders.id;
+    """
+    )
+
+    op.alter_column("vfolders", "status_history", nullable=True, default=None)
+    op.execute("UPDATE vfolders SET status_history = NULL WHERE status_history = '[]'::jsonb;")
diff --git a/src/ai/backend/manager/models/vfolder.py b/src/ai/backend/manager/models/vfolder.py
index 35e0a40e202..b54aaea5c41 100644
--- a/src/ai/backend/manager/models/vfolder.py
+++ b/src/ai/backend/manager/models/vfolder.py
@@ -97,7 +97,12 @@
 from .rbac.exceptions import NotEnoughPermission
 from .session import DEAD_SESSION_STATUSES, SessionRow
 from .user import UserRole, UserRow
-from .utils import ExtendedAsyncSAEngine, execute_with_retry, sql_json_merge
+from .utils import (
+    ExtendedAsyncSAEngine,
+    execute_with_retry,
+    sql_append_dict_to_list,
+    sql_json_merge,
+)
 
 if TYPE_CHECKING:
     from ..api.context import BackgroundTaskManager
@@ -367,14 +372,14 @@ class VFolderCloneInfo(NamedTuple):
         nullable=False,
         index=True,
     ),
-    # status_history records the most recent status changes for each status
+    # status_history records the status changes of the vfolder.
     # e.g)
-    # {
-    #   "ready": "2022-10-22T10:22:30",
-    #   "delete-pending": "2022-10-22T11:40:30",
-    #   "delete-ongoing": "2022-10-25T10:22:30"
-    # }
-    sa.Column("status_history", pgsql.JSONB(), nullable=True, default=sa.null()),
+    # [
+    #   {"status": "ready", "timestamp": "2022-10-22T10:22:30"},
+    #   {"status": "delete-pending", "timestamp": "2022-10-22T11:40:30"},
+    #   {"status": "delete-ongoing", "timestamp": "2022-10-25T10:22:30"}
+    # ]
+    sa.Column("status_history", pgsql.JSONB(), nullable=False, default=[]),
     sa.Column("status_changed", sa.DateTime(timezone=True), nullable=True, index=True),
 )
 
@@ -1525,12 +1530,9 @@ async def _update() -> None:
                 .values(
                     status=update_status,
                     status_changed=now,
-                    status_history=sql_json_merge(
-                        vfolders.c.status_history,
-                        (),
-                        {
-                            update_status.name: now.isoformat(),
-                        },
+                    status_history=sql_append_dict_to_list(
+                        VFolderRow.status_history,
+                        {"status": update_status.name, "timestamp": now.isoformat()},
                     ),
                 )
                 .where(cond)