1  """Per-API-key monthly token quotas.
 2  
 3  Adds three columns to ``api_keys``:
 4  * ``token_quota_monthly`` (nullable int) — monthly cap; NULL = unlimited
 5  * ``tokens_used_this_month`` (int, default 0) — rolling counter
 6  * ``quota_reset_at`` (nullable datetime) — first-of-next-month rollover
 7  
 8  Lets an SMB mint one key per customer and isolate quota per customer.
 9  Enforced by ``check_api_key_quota`` in ``restai/budget.py`` alongside
10  the existing project rate limit + team budget checks.
11  """
12  import sqlalchemy as sa
13  from alembic import op
14  
15  
16  revision = '042'
17  down_revision = '041'
18  branch_labels = None
19  depends_on = None
20  
21  
22  def upgrade():
23      for col_name, col in (
24          ('token_quota_monthly', sa.Column('token_quota_monthly', sa.Integer(), nullable=True)),
25          ('tokens_used_this_month', sa.Column('tokens_used_this_month', sa.Integer(), nullable=False, server_default='0')),
26          ('quota_reset_at', sa.Column('quota_reset_at', sa.DateTime(), nullable=True)),
27      ):
28          try:
29              op.add_column('api_keys', col)
30          except Exception:
31              # Column already present (partial re-run) — safe to ignore.
32              pass
33  
34  
35  def downgrade():
36      for col_name in ('token_quota_monthly', 'tokens_used_this_month', 'quota_reset_at'):
37          try:
38              op.drop_column('api_keys', col_name)
39          except Exception:
40              pass