PamGram/alembic/versions/a9c8cde17cd8_tasks_player_id.py
2024-11-18 21:47:07 +08:00

83 lines
2.8 KiB
Python

"""tasks_player_id
Revision ID: a9c8cde17cd8
Revises: 1220c5c80757
Create Date: 2024-11-18 15:47:51.218077
"""
import logging
from alembic import op
import sqlalchemy as sa
from sqlalchemy import text
from sqlalchemy.dialects import mysql
from sqlalchemy.exc import NoSuchTableError
# revision identifiers, used by Alembic.
revision = "a9c8cde17cd8"
down_revision = "1220c5c80757"
branch_labels = None
depends_on = None
logger = logging.getLogger(__name__)
def update_player_id():
connection = op.get_bind()
try:
statement = "SELECT user_id FROM task;"
task_table_data = connection.execute(text(statement))
need_check_user_id = []
if task_table_data is not None:
for row in task_table_data:
need_check_user_id.append(row[0])
need_check_user_id = list(set(need_check_user_id))
except NoSuchTableError:
logger.warning("Table 'task' doesn't exist")
return # should not happen
try:
statement = "SELECT user_id, player_id, is_chosen FROM players;"
players_table_data = connection.execute(text(statement))
player_id_map = {}
if players_table_data is not None:
for row in players_table_data:
if not row[2]:
continue
uid, pid = row[0], row[1]
if uid not in player_id_map:
player_id_map[uid] = pid
except NoSuchTableError:
logger.warning("Table 'players' doesn't exist")
return # should not happen
update = "UPDATE task SET player_id=:player_id WHERE user_id=:user_id;"
for uid in need_check_user_id:
player_id = player_id_map.get(uid, None)
if player_id is None:
logger.warning("user_id %s doesn't exist player", uid)
continue
try:
with op.get_context().autocommit_block():
connection.execute(text(update), dict(player_id=player_id, user_id=uid))
except Exception as exc: # pylint: disable=W0703
logger.error("Process sign->task Exception", exc_info=exc) # pylint: disable=W0703
def upgrade() -> None:
# ### commands auto generated by Alembic - please adjust! ###
op.add_column("task", sa.Column("player_id", sa.BigInteger(), nullable=False))
op.alter_column("task", "user_id", existing_type=mysql.BIGINT(), nullable=False)
op.create_index(op.f("ix_task_player_id"), "task", ["player_id"], unique=False)
update_player_id()
# ### end Alembic commands ###
def downgrade() -> None:
# ### commands auto generated by Alembic - please adjust! ###
op.drop_index(op.f("ix_task_player_id"), table_name="task")
op.alter_column("task", "user_id", existing_type=mysql.BIGINT(), nullable=True)
op.drop_column("task", "player_id")
# ### end Alembic commands ###