117 lines
5.0 KiB
Python
117 lines
5.0 KiB
Python
"""Add chat token usage tracking
|
|
|
|
Revision ID: cd27562afa2e
|
|
Revises: 1f4d2a8c9b77
|
|
Create Date: 2026-04-16 13:05:09.572442
|
|
|
|
"""
|
|
from typing import Sequence, Union
|
|
|
|
from alembic import op
|
|
import sqlalchemy as sa
|
|
from sqlalchemy.dialects import postgresql
|
|
|
|
# revision identifiers, used by Alembic.
|
|
revision: str = 'cd27562afa2e'
|
|
down_revision: Union[str, None] = '1f4d2a8c9b77'
|
|
branch_labels: Union[str, Sequence[str], None] = None
|
|
depends_on: Union[str, Sequence[str], None] = None
|
|
|
|
|
|
def _table_exists(bind, table_name: str) -> bool:
|
|
return sa.inspect(bind).has_table(table_name)
|
|
|
|
|
|
def upgrade() -> None:
|
|
bind = op.get_bind()
|
|
|
|
if _table_exists(bind, 'chatbot_documents'):
|
|
with op.batch_alter_table('chatbot_documents', schema=None) as batch_op:
|
|
batch_op.alter_column('chunks_count',
|
|
existing_type=sa.INTEGER(),
|
|
nullable=False)
|
|
batch_op.alter_column('created_at',
|
|
existing_type=postgresql.TIMESTAMP(),
|
|
nullable=False)
|
|
|
|
if _table_exists(bind, 'vector_chunks'):
|
|
with op.batch_alter_table('vector_chunks', schema=None) as batch_op:
|
|
batch_op.alter_column('document_id',
|
|
existing_type=sa.VARCHAR(),
|
|
nullable=False)
|
|
batch_op.alter_column('chunk_text',
|
|
existing_type=sa.VARCHAR(),
|
|
type_=sa.Text(),
|
|
existing_nullable=False)
|
|
batch_op.alter_column('chunk_index',
|
|
existing_type=sa.INTEGER(),
|
|
nullable=False)
|
|
batch_op.alter_column('created_at',
|
|
existing_type=postgresql.TIMESTAMP(),
|
|
type_=sa.DateTime(timezone=True),
|
|
nullable=False)
|
|
batch_op.drop_constraint('vector_chunks_user_id_fkey', type_='foreignkey')
|
|
batch_op.create_foreign_key(None, 'users', ['user_id'], ['id'], ondelete='CASCADE')
|
|
|
|
if _table_exists(bind, 'vector_indices'):
|
|
with op.batch_alter_table('vector_indices', schema=None) as batch_op:
|
|
batch_op.alter_column('total_vectors',
|
|
existing_type=sa.INTEGER(),
|
|
nullable=False)
|
|
batch_op.alter_column('updated_at',
|
|
existing_type=postgresql.TIMESTAMP(),
|
|
type_=sa.DateTime(timezone=True),
|
|
nullable=False)
|
|
batch_op.drop_constraint('vector_indices_user_id_fkey', type_='foreignkey')
|
|
batch_op.create_foreign_key(None, 'users', ['user_id'], ['id'], ondelete='CASCADE')
|
|
|
|
|
|
def downgrade() -> None:
|
|
bind = op.get_bind()
|
|
|
|
if _table_exists(bind, 'vector_indices'):
|
|
with op.batch_alter_table('vector_indices', schema=None) as batch_op:
|
|
batch_op.drop_constraint(None, type_='foreignkey')
|
|
batch_op.create_foreign_key('vector_indices_user_id_fkey', 'users', ['user_id'], ['id'])
|
|
batch_op.drop_index(batch_op.f('ix_vector_indices_user_id'))
|
|
batch_op.create_index('ix_vector_indices_user_id', ['user_id'], unique=True)
|
|
batch_op.create_index('ix_vector_indices_id', ['id'], unique=False)
|
|
batch_op.alter_column('updated_at',
|
|
existing_type=sa.DateTime(timezone=True),
|
|
type_=postgresql.TIMESTAMP(),
|
|
nullable=True)
|
|
batch_op.alter_column('total_vectors',
|
|
existing_type=sa.INTEGER(),
|
|
nullable=True)
|
|
|
|
if _table_exists(bind, 'vector_chunks'):
|
|
with op.batch_alter_table('vector_chunks', schema=None) as batch_op:
|
|
batch_op.drop_constraint(None, type_='foreignkey')
|
|
batch_op.create_foreign_key('vector_chunks_user_id_fkey', 'users', ['user_id'], ['id'])
|
|
batch_op.create_index('ix_vector_chunks_id', ['id'], unique=False)
|
|
batch_op.alter_column('created_at',
|
|
existing_type=sa.DateTime(timezone=True),
|
|
type_=postgresql.TIMESTAMP(),
|
|
nullable=True)
|
|
batch_op.alter_column('chunk_index',
|
|
existing_type=sa.INTEGER(),
|
|
nullable=True)
|
|
batch_op.alter_column('chunk_text',
|
|
existing_type=sa.Text(),
|
|
type_=sa.VARCHAR(),
|
|
existing_nullable=False)
|
|
batch_op.alter_column('document_id',
|
|
existing_type=sa.VARCHAR(),
|
|
nullable=True)
|
|
|
|
if _table_exists(bind, 'chatbot_documents'):
|
|
with op.batch_alter_table('chatbot_documents', schema=None) as batch_op:
|
|
batch_op.create_index('ix_chatbot_documents_id', ['id'], unique=False)
|
|
batch_op.create_index('ix_chatbot_documents_drive_file_id', ['drive_file_id'], unique=False)
|
|
batch_op.alter_column('created_at',
|
|
existing_type=postgresql.TIMESTAMP(),
|
|
nullable=True)
|
|
batch_op.alter_column('chunks_count',
|
|
existing_type=sa.INTEGER(),
|
|
nullable=True)
|