Migration script

1 files changed, 113 insertions, 0 deletions

diff --git a/migrations/versions/283aad3fa60b_.py b/migrations/versions/283aad3fa60b_.py
new file mode 100644
index 0000000..ff9547d
--- /dev/null
+++ b/migrations/versions/283aad3fa60b_.py
@@ -0,0 +1,113 @@
+"""Create ConstExpenseGroup and migrate ConstExpense
+
+Revision ID: 283aad3fa60b
+Revises: 27880ab1f050
+Create Date: 2016-01-10 15:20:48.661353
+
+"""
+
+# revision identifiers, used by Alembic.
+revision = '283aad3fa60b'
+down_revision = '27880ab1f050'
+
+from alembic import op
+import sqlalchemy as sa
+from sqlalchemy import Column, Integer, Unicode, Date, ForeignKey
+from sqlalchemy.ext.declarative import declarative_base, declared_attr
+from sqlalchemy.orm import sessionmaker, relationship, backref
+
+class Base (declarative_base()):
+    __abstract__ = True
+    id = Column(Integer, primary_key = True)
+
+    @declared_attr
+    def __tablename__ (cls):
+        return cls.__name__.lower()
+
+class User(Base):
+    pass
+
+class Category(Base):
+    pass
+
+class ConstExpenseGroup (Base):
+    description = Column(Unicode(50), nullable=False)
+    start       = Column(Date, nullable=False)
+    end         = Column(Date, nullable=False)
+    user_id     = Column(Integer, ForeignKey(User.id), nullable=False)
+    category_id = Column(Integer, ForeignKey(Category.id), nullable=False)
+
+class ConstExpense (Base):
+    description = Column(Unicode(50))
+    start       = Column(Date)
+    end         = Column(Date)
+    user_id     = Column(Integer)
+    category_id = Column(Integer)
+
+    group_id    = Column(Integer, ForeignKey(ConstExpenseGroup.id))
+    group       = relationship(ConstExpenseGroup, uselist = False)
+
+    prev_id     = Column(Integer, ForeignKey('constexpense.id'))
+    prev        = relationship('ConstExpense',
+            remote_side = 'ConstExpense.id', uselist = False,
+            backref=backref('next', uselist = False, lazy = 'immediate'))
+
+def migrate_data(session):
+    expenses = []
+    for exp in session.query(ConstExpense).filter(ConstExpense.prev_id == None).order_by(ConstExpense.id):
+        group = ConstExpenseGroup(
+            description = exp.description,
+            start = exp.start,
+            end = exp.end,
+            user_id = exp.user_id,
+            category_id = exp.category_id)
+
+        exp.description = None
+        exp.group = group
+        session.add(group)
+        session.add(exp)
+
+        if exp.next is not None:
+            expenses.append(exp)
+
+    session.flush()
+
+    for exp in expenses:
+        group = exp.group
+        while exp.next is not None:
+            exp = exp.next
+            if exp.description == group.description:
+                exp.description = None
+
+            exp.group = group
+            if exp.start < group.start:
+                group.start = exp.start
+            if exp.end > group.end:
+                group.end = exp.end
+
+            session.add(exp)
+            session.add(group)
+
+
+def upgrade():
+    bind = op.get_bind()
+    session = sessionmaker()(bind=bind)
+
+    # new table
+    ConstExpenseGroup.__table__.create(bind)
+
+    with op.batch_alter_table('constexpense') as batch_op:
+        batch_op.add_column(sa.Column('group_id', sa.Integer(), nullable=True))
+        batch_op.create_foreign_key('fk_constexpense_group', 'constexpensegroup', ['group_id'], ['id'])
+
+    # migrate data
+    migrate_data(session)
+    session.commit()
+
+    # new indexes
+    op.create_index(op.f('ix_constexpensegroup_end'), 'constexpensegroup', ['end'], unique=False)
+    op.create_index(op.f('ix_constexpensegroup_start'), 'constexpensegroup', ['start'], unique=False)
+    op.create_index(op.f('ix_constexpense_group_id'), 'constexpense', ['group_id'], unique=False)
+
+def downgrade():
+    raise Exception("Irreversible migration")