skorotkiewicz/rDrama
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

[themotte/rDrama#451] Use bulk_recompute_descendant_counts() in migration

Browse source
This commit is contained in:
faul_sname 2023-01-13 19:14:40 -08:00
parent cdac0174f1
commit 12ca271fe7
1 changed files with 18 additions and 24 deletions
Download patch file Download diff file Expand all files Collapse all files

42
migrations/versions/2023_01_03_07_59_56_1f30a37b08a0_dml_populate_comments_descendant_count_.py
View file

@ -6,11 +6,15 @@ Create Date: 2023-01-03 07:59:56.599209+00:00
""" """
from alembic import op from alembic import op
import sqlalchemy as sa from sqlalchemy.sql.expression import func, text
from sqlalchemy.sql.expression import func from sqlalchemy.orm.session import Session
from sqlalchemy import update
from flask import g
from files.__main__ import db_session from files.__main__ import db_session
from files.classes import Comment from files.classes import Comment
from files.helpers.comments import bulk_recompute_descendant_counts
# revision identifiers, used by Alembic. # revision identifiers, used by Alembic.
revision = '1f30a37b08a0' revision = '1f30a37b08a0'
@ -18,28 +22,18 @@ down_revision = 'f8ba0e88ddd1'
branch_labels = None branch_labels = None
depends_on = 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(): def upgrade():
db = db_session() with g_db_set_from_alembic():
bulk_recompute_descendant_counts()
max_level = db.query(func.max(Comment.level)).scalar()
for level in range(max_level, 0, -1):
db.execute('''
WITH "descendant_counts" as (
SELECT
"parent_comments"."id",
COALESCE(SUM(1 + "child_comments"."descendant_count"), 0) as "descendant_count"
FROM "comments" as "parent_comments"
LEFT JOIN "comments" AS "child_comments" ON "child_comments"."parent_comment_id" = "parent_comments"."id"
GROUP BY "parent_comments"."id"
)
UPDATE "comments"
SET "descendant_count" = "descendant_counts"."descendant_count"
FROM "descendant_counts"
WHERE "comments"."id" = "descendant_counts"."id"
''', {"level": level})
db.commit()
def downgrade(): def downgrade():
db = db_session() with g_db_set_from_alembic():
db.query(Comment).update({'descendant_count': 0}) g.db.execute(update(Comment).values(descendant_count=0))