Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

refactor: Replace vfolder's status_history's type map with list #2111

Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions changes/2111.feature.md
Original file line number Diff line number Diff line change
@@ -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.
Original file line number Diff line number Diff line change
@@ -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;")
29 changes: 15 additions & 14 deletions src/ai/backend/manager/models/vfolder.py
Original file line number Diff line number Diff line change
Expand Up @@ -73,7 +73,11 @@
from .minilang.ordering import OrderSpecItem, QueryOrderParser
from .minilang.queryfilter import FieldSpecItem, QueryFilterParser, enum_field_getter
from .user import UserRole
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
Expand Down Expand Up @@ -332,14 +336,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),
sa.CheckConstraint(
"(ownership_type = 'user' AND \"user\" IS NOT NULL) OR "
Expand Down Expand Up @@ -942,12 +946,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)
Expand Down
Loading