[themotte/rDrama#451] Passing in the db connection as an optional param makes things easier

This commit is contained in:
faul_sname 2023-01-13 19:33:56 -08:00
parent 12ca271fe7
commit c0a546d779
3 changed files with 38 additions and 29 deletions

View file

@ -9,7 +9,6 @@ from alembic import op
from sqlalchemy.sql.expression import func, text
from sqlalchemy.orm.session import Session
from sqlalchemy import update
from flask import g
from files.__main__ import db_session
from files.classes import Comment
@ -22,18 +21,10 @@ down_revision = 'f8ba0e88ddd1'
branch_labels = None
depends_on = None
class g_db_set_from_alembic():
def __enter__(self, *args, **kwargs):
g.db = Session(bind=op.get_bind())
self.old_db = getattr(g, 'db', None)
def __exit__(self, *args, **kwargs):
g.db = self.old_db
def upgrade():
with g_db_set_from_alembic():
bulk_recompute_descendant_counts()
db =Session(bind=op.get_bind())
bulk_recompute_descendant_counts(lambda q: q, db)
def downgrade():
with g_db_set_from_alembic():
g.db.execute(update(Comment).values(descendant_count=0))
db =Session(bind=op.get_bind())
db.execute(update(Comment).values(descendant_count=0))