"""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')