diff --git a/changes/2116.feature.md b/changes/2116.feature.md new file mode 100644 index 0000000000..755eaedf36 --- /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 0000000000..c39b0b6b66 --- /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 1b5f87b7a7..d203ea1bf0 100644 --- a/src/ai/backend/manager/models/vfolder.py +++ b/src/ai/backend/manager/models/vfolder.py @@ -104,7 +104,11 @@ from .rbac.exceptions import InvalidScope, 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, +) if TYPE_CHECKING: from ..api.context import BackgroundTaskManager @@ -374,14 +378,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), ) @@ -1467,12 +1471,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)