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

migrations/versions/2023_01_03_07_59_56_1f30a37b08a0_dml_populate_comments_descendant_count_.py

@@ -6,11 +6,15 @@ Create Date: 2023-01-03 07:59:56.599209+00:00
 
 """
 from alembic import op
-import sqlalchemy as sa
-from sqlalchemy.sql.expression import func
+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
+from files.helpers.comments import bulk_recompute_descendant_counts
+
 
 # revision identifiers, used by Alembic.
 revision = '1f30a37b08a0'
@@ -18,28 +22,18 @@ 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():
-	db = db_session()
-
-	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()
-
+	with g_db_set_from_alembic():
+		bulk_recompute_descendant_counts()
 
 def downgrade():
-	db = db_session()
-	db.query(Comment).update({'descendant_count': 0})
+	with g_db_set_from_alembic():
+		g.db.execute(update(Comment).values(descendant_count=0))