58 lines
2.2 KiB
Python
58 lines
2.2 KiB
Python
"""Make DriveComment generic
|
|
|
|
Revision ID: 9a6db38b86a6
|
|
Revises: 41bed7eb6ed9
|
|
Create Date: 2026-02-27 14:32:05.071480
|
|
|
|
"""
|
|
from typing import Sequence, Union
|
|
|
|
from alembic import op
|
|
import sqlalchemy as sa
|
|
|
|
|
|
# revision identifiers, used by Alembic.
|
|
revision: str = '9a6db38b86a6'
|
|
down_revision: Union[str, None] = '41bed7eb6ed9'
|
|
branch_labels: Union[str, Sequence[str], None] = None
|
|
depends_on: Union[str, Sequence[str], None] = None
|
|
|
|
|
|
def upgrade() -> None:
|
|
# 1. Add columns as nullable first
|
|
op.add_column('drive_comments', sa.Column('resource_type', sa.String(length=20), nullable=True))
|
|
op.add_column('drive_comments', sa.Column('resource_id', sa.Integer(), nullable=True))
|
|
|
|
# 2. Populate data from existing file_id safely
|
|
# Use WHERE file_id IS NOT NULL to avoid unnecessary or invalid updates
|
|
op.execute("UPDATE drive_comments SET resource_type = 'file', resource_id = file_id WHERE file_id IS NOT NULL")
|
|
|
|
# 3. Set to NOT NULL now that data is populated
|
|
op.alter_column('drive_comments', 'resource_type', nullable=False)
|
|
op.alter_column('drive_comments', 'resource_id', nullable=False)
|
|
|
|
# 4. Cleanup old column and constraint
|
|
with op.batch_alter_table('drive_comments') as batch_op:
|
|
batch_op.drop_column('file_id')
|
|
|
|
|
|
def downgrade() -> None:
|
|
# 1. Add column back as nullable first
|
|
op.add_column('drive_comments', sa.Column('file_id', sa.Integer(), nullable=True))
|
|
|
|
# 2. Restore data from generic columns back to file_id
|
|
# Only migrate records that were originally or are currently 'file' types
|
|
op.execute("UPDATE drive_comments SET file_id = resource_id WHERE resource_type = 'file'")
|
|
|
|
# 3. Delete or handle non-file comments that cannot exist in the old schema
|
|
# Alternatively, keep file_id as nullable if preferred, but here we restore original state
|
|
op.execute("DELETE FROM drive_comments WHERE resource_type != 'file'")
|
|
|
|
# 4. Restore constraints
|
|
op.alter_column('drive_comments', 'file_id', nullable=False)
|
|
op.create_foreign_key('drive_comments_file_id_fkey', 'drive_comments', 'drive_files', ['file_id'], ['id'])
|
|
|
|
# 5. Cleanup generic columns
|
|
op.drop_column('drive_comments', 'resource_id')
|
|
op.drop_column('drive_comments', 'resource_type')
|