41 lines
1.1 KiB
Python
41 lines
1.1 KiB
Python
"""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')
|