NexusDashboard/migrations/versions/1164e037907f_compressss_reports.py

"""compressss reports

Revision ID: 1164e037907f
Revises: a6e42ef03da7
Create Date: 2023-11-18 01:38:00.127472

"""
from alembic import op
from sqlalchemy.orm.session import Session
import sqlalchemy as sa
from sqlalchemy.dialects import mysql
from sqlalchemy.types import JSON
from sqlalchemy.ext.declarative import declarative_base
import gzip
import json


Base = declarative_base()


# revision identifiers, used by Alembic.
revision = '1164e037907f'
down_revision = 'a6e42ef03da7'
branch_labels = None
depends_on = None

class ReportsUpgradeNew(Base):
    __tablename__ = 'reports'
    __table_args__ = {'extend_existing': True}

    data = sa.Column(
        mysql.MEDIUMBLOB(),
        nullable=False
    )

    report_type = sa.Column(
        sa.VARCHAR(35),
        nullable=False,
        primary_key=True,
        autoincrement=False
    )

    date = sa.Column(
        sa.Date(),
        primary_key=True,
        autoincrement=False
    )

    def save(self):
        sa.session.add(self)
        sa.session.commit()
        sa.session.refresh(self)


class ReportsUpgradeOld(Base):
    __tablename__ = 'reports_old'
    __table_args__ = {'extend_existing': True}

    data = sa.Column(
        JSON(),
        nullable=False
    )

    report_type = sa.Column(
        sa.VARCHAR(35),
        nullable=False,
        primary_key=True,
        autoincrement=False
    )

    date = sa.Column(
        sa.Date(),
        primary_key=True,
        autoincrement=False
    )

class ReportsDowngradeOld(Base):
    __tablename__ = 'reports'
    __table_args__ = {'extend_existing': True}

    data = sa.Column(
        mysql.MEDIUMBLOB(),
        nullable=False
    )

    report_type = sa.Column(
        sa.VARCHAR(35),
        nullable=False,
        primary_key=True,
        autoincrement=False
    )

    date = sa.Column(
        sa.Date(),
        primary_key=True,
        autoincrement=False
    )

    def save(self):
        sa.session.add(self)
        sa.session.commit()
        sa.session.refresh(self)


class ReportsDowngradeNew(Base):
    __tablename__ = 'reports_old'
    __table_args__ = {'extend_existing': True}

    data = sa.Column(
        JSON(),
        nullable=False
    )

    report_type = sa.Column(
        sa.VARCHAR(35),
        nullable=False,
        primary_key=True,
        autoincrement=False
    )

    date = sa.Column(
        sa.Date(),
        primary_key=True,
        autoincrement=False
    )


def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###

    op.rename_table('reports', 'reports_old')
    bind = op.get_bind()
    session = Session(bind=bind)
    reports = session.query(ReportsUpgradeOld)
    op.create_table('reports',
        sa.Column('data', mysql.MEDIUMBLOB(), nullable=False),
        sa.Column('report_type', sa.VARCHAR(length=35), autoincrement=False, nullable=False),
        sa.Column('date', sa.Date(), autoincrement=False, nullable=False),
        sa.PrimaryKeyConstraint('report_type', 'date')
    )
    # insert records
    new_reports = []
    # insert records
    for report in reports:
        new_reports.append({
            "data":gzip.compress(json.dumps(report.data).encode('utf-8')), 
            "report_type":report.report_type, 
            "date":report.date
        })
    op.bulk_insert(ReportsUpgradeNew.__table__, new_reports)
    op.drop_table('reports_old')
    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_table('reports')
    op.create_table('reports',
        sa.Column('data', JSON(), nullable=False),
        sa.Column('report_type', sa.VARCHAR(length=35), autoincrement=False, nullable=False),
        sa.Column('date', sa.Date(), autoincrement=False, nullable=False),
        sa.PrimaryKeyConstraint('report_type', 'date')
    )
    # ### end Alembic commands ###