Files
docqube_backend/alembic/versions/540c970ff890_add_tenant_chat_token_quota.py
T
2026-09-08 11:00:05 +05:30

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