43 lines
1.4 KiB
Python
43 lines
1.4 KiB
Python
"""add_tenant_chat_token_quota
|
|
|
|
Revision ID: 540c970ff890
|
|
Revises: 7fc8dd61bb95
|
|
Create Date: 2026-05-13 11:21:14.192033
|
|
|
|
"""
|
|
from typing import Sequence, Union
|
|
|
|
from alembic import op
|
|
import sqlalchemy as sa
|
|
|
|
|
|
# revision identifiers, used by Alembic.
|
|
revision: str = '540c970ff890'
|
|
down_revision: Union[str, None] = '7fc8dd61bb95'
|
|
branch_labels: Union[str, Sequence[str], None] = None
|
|
depends_on: Union[str, Sequence[str], None] = None
|
|
|
|
|
|
def upgrade() -> None:
|
|
# Get current connection and inspector
|
|
conn = op.get_bind()
|
|
inspector = sa.inspect(conn)
|
|
|
|
# Check if columns already exist (defensive migration)
|
|
columns = [c['name'] for c in inspector.get_columns('tenants')]
|
|
|
|
if 'chat_token_daily_limit' not in columns:
|
|
op.add_column('tenants', sa.Column('chat_token_daily_limit', sa.BigInteger(), nullable=False, server_default='1000000'))
|
|
|
|
if 'chat_tokens_used_today' not in columns:
|
|
op.add_column('tenants', sa.Column('chat_tokens_used_today', sa.BigInteger(), nullable=False, server_default='0'))
|
|
|
|
if 'chat_tokens_last_reset_at' not in columns:
|
|
op.add_column('tenants', sa.Column('chat_tokens_last_reset_at', sa.DateTime(timezone=True), server_default=sa.text('now()'), nullable=False))
|
|
|
|
|
|
def downgrade() -> None:
|
|
op.drop_column('tenants', 'chat_tokens_last_reset_at')
|
|
op.drop_column('tenants', 'chat_tokens_used_today')
|
|
op.drop_column('tenants', 'chat_token_daily_limit')
|