tracim/tracim/migration/versions/f3852e1349c4_all_files_only...

"""all files only on disk.

Revision ID: f3852e1349c4
Revises: 913efdf409e5
Create Date: 2017-07-24 17:15:54.278141

"""

from alembic import context
from alembic import op

from depot.fields.sqlalchemy import UploadedFileField
from depot.manager import DepotManager

import sqlalchemy as sa

# revision identifiers, used by Alembic.
revision = 'f3852e1349c4'
down_revision = '913efdf409e5'


def configure_depot():
    """Configure Depot."""
    depot_storage_name = context.config.get_main_option('depot_storage_name')
    depot_storage_path = context.config.get_main_option('depot_storage_dir')
    depot_storage_settings = {'depot.storage_path': depot_storage_path}
    DepotManager.configure(
        depot_storage_name,
        depot_storage_settings,
    )


revision_helper = sa.Table(
    'content_revisions',
    sa.MetaData(),
    sa.Column('revision_id', sa.Integer, primary_key=True),
    sa.Column('file_content', sa.LargeBinary),
    sa.Column('depot_file', UploadedFileField, nullable=True),
    sa.Column('type', sa.String(32), nullable=False),
)


def upgrade():
    """Drop the file content from revision."""
    with op.batch_alter_table('content_revisions') as batch_op:
        batch_op.drop_column('file_content')


def downgrade():
    """Add the file content in revision."""
    with op.batch_alter_table('content_revisions') as batch_op:
        batch_op.add_column(sa.Column('file_content', sa.LargeBinary))

    configure_depot()
    depot = DepotManager.get()
    connection = op.get_bind()
    select_query = revision_helper.select() \
        .where(revision_helper.c.type == 'file') \
        .where(revision_helper.c.depot_file.isnot(None))
    files = connection.execute(select_query).fetchall()
    for file in files:
        depot_file_content = depot.get(file.depot_file).read()
        update_query = revision_helper.update() \
            .where(revision_helper.c.revision_id == file.revision_id) \
            .values(file_content=depot_file_content)
        connection.execute(update_query)