"""Add token_count to Message

Revision ID: 190d5ef26e20
Revises: a9e886c56482
Create Date: 2023-12-07 08:55:50.378439

"""

from alembic import op
import sqlalchemy as sa

from sqlalchemy.orm import Session

# revision identifiers, used by Alembic.
revision = "190d5ef26e20"
down_revision = "a9e886c56482"
branch_labels = None
depends_on = None

from app import Message


def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    with op.batch_alter_table("message", schema=None) as batch_op:
        batch_op.add_column(sa.Column("token_count", sa.Integer(), nullable=True))

    # Use this binding to connect to the database
    bind = op.get_bind()
    session = Session(bind=bind)

    # Assuming the content column is not nullable and always has a value
    messages = session.query(Message).all()
    for message in messages:
        message.count_tokens()
        session.add(message)

    session.commit()


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    with op.batch_alter_table("message", schema=None) as batch_op:
        batch_op.drop_column("token_count")
