39 lines
1.3 KiB
Python
39 lines
1.3 KiB
Python
"""Add daily_token_usage table.
|
|
|
|
Revision ID: add_daily_token_usage
|
|
"""
|
|
|
|
from alembic import op
|
|
import sqlalchemy as sa
|
|
from sqlalchemy.dialects import postgresql
|
|
|
|
revision = "add_daily_token_usage"
|
|
down_revision = "add_agentbay_enum_value"
|
|
branch_labels = None
|
|
depends_on = None
|
|
|
|
def upgrade() -> None:
|
|
# Use raw SQL with IF NOT EXISTS to safely handle instances where the table already exists
|
|
op.execute("""
|
|
CREATE TABLE IF NOT EXISTS daily_token_usage (
|
|
id UUID PRIMARY KEY,
|
|
tenant_id UUID NOT NULL REFERENCES tenants(id),
|
|
agent_id UUID NOT NULL REFERENCES agents(id),
|
|
date TIMESTAMPTZ NOT NULL,
|
|
tokens_used INTEGER NOT NULL DEFAULT 0,
|
|
created_at TIMESTAMPTZ DEFAULT now(),
|
|
updated_at TIMESTAMPTZ DEFAULT now()
|
|
)
|
|
""")
|
|
op.execute("""
|
|
CREATE UNIQUE INDEX IF NOT EXISTS uq_daily_token_usage_agent_date
|
|
ON daily_token_usage (agent_id, date)
|
|
""")
|
|
op.execute("CREATE INDEX IF NOT EXISTS ix_daily_token_usage_tenant_id ON daily_token_usage (tenant_id)")
|
|
op.execute("CREATE INDEX IF NOT EXISTS ix_daily_token_usage_agent_id ON daily_token_usage (agent_id)")
|
|
op.execute("CREATE INDEX IF NOT EXISTS ix_daily_token_usage_date ON daily_token_usage (date)")
|
|
|
|
|
|
def downgrade() -> None:
|
|
op.execute("DROP TABLE IF EXISTS daily_token_usage")
|