"""Add file size

Revision ID: 1921a169739b
Revises: aa00db20c10a
Create Date: 2020-04-30 14:39:41.662130

"""
import os, sys

from alembic import op, context
import sqlalchemy as sa
from sqlalchemy import orm
from sqlalchemy import Column, INTEGER, TEXT
from sqlalchemy.ext.declarative import declarative_base

# revision identifiers, used by Alembic.
revision = '1921a169739b'
down_revision = 'aa00db20c10a'
branch_labels = None
depends_on = None

Base = declarative_base()


class File(Base):
    'A File (for assignment, submission file, or submission feedback)'
    __tablename__ = 'files'
    _id = Column(INTEGER, primary_key=True)
    filename = Column(TEXT)
    checksum = Column(TEXT)
    size = Column(INTEGER)
    actual_name = Column(TEXT)


def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column('files', sa.Column('size', sa.INTEGER(), nullable=True))
    # https://alembic.sqlalchemy.org/en/latest/cookbook.html
    # http://ominian.com/2019/07/11/data-migration-with-sqlalchemy-and-alembic/
    if context.get_x_argument(as_dictionary=True).get('data', None):
        # Perform data migrations
        storage = context.get_x_argument(as_dictionary=True)['storage']
        session = orm.Session(bind=op.get_bind())
        for i in session.query(File):
            pathname = os.path.join(storage, i.actual_name)
            try:
                i.size = os.path.getsize(pathname)
            except FileNotFoundError:
                print(
                    'WARNING: File %s not found' % repr(pathname),
                    file=sys.stderr,
                )
        session.commit()
    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    # The following command won't work, see:
    # https://alembic.sqlalchemy.org/en/latest/batch.html
    # op.drop_column('files', 'size')
    # Therefore, we have to do:
    with op.batch_alter_table("files") as batch_op:
        batch_op.drop_column('size')

    # ### end Alembic commands ###
