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

67 lines
1.6 KiB
Python

"""add tenant storage quota bytes
Revision ID: c7d4e6f1a9b0
Revises: e81a9d22ca5f
Create Date: 2026-03-07 09:25:00.000000
"""
from typing import Sequence, Union
from alembic import op
import sqlalchemy as sa
# revision identifiers, used by Alembic.
revision: str = "c7d4e6f1a9b0"
down_revision: Union[str, None] = "e81a9d22ca5f"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
op.add_column(
"tenants",
sa.Column("storage_quota_bytes", sa.BigInteger(), nullable=True),
)
# Preserve existing effective quota behavior by summing each tenant's
# current per-user quotas. For tenants with no users, default to 1 GB.
op.execute(
"""
UPDATE tenants t
SET storage_quota_bytes = COALESCE(
(
SELECT SUM(COALESCE(usu.max_bytes_quota, 1073741824))
FROM users u
LEFT JOIN user_storage_usage usu ON usu.user_id = u.id
WHERE u.tenant_id = t.id
),
1073741824
)
"""
)
op.alter_column(
"tenants",
"storage_quota_bytes",
existing_type=sa.BigInteger(),
nullable=False,
server_default=sa.text("1073741824"),
)
op.create_check_constraint(
"ck_tenants_storage_quota_bytes_positive",
"tenants",
"storage_quota_bytes > 0",
)
def downgrade() -> None:
op.drop_constraint(
"ck_tenants_storage_quota_bytes_positive",
"tenants",
type_="check",
)
op.drop_column("tenants", "storage_quota_bytes")