diff --git a/api/migrations/versions/2025_12_25_1039-7df29de0f6be_add_credit_pool.py b/api/migrations/versions/2025_12_25_1039-7df29de0f6be_add_credit_pool.py index e89fcee7e5..6a9bfd2be0 100644 --- a/api/migrations/versions/2025_12_25_1039-7df29de0f6be_add_credit_pool.py +++ b/api/migrations/versions/2025_12_25_1039-7df29de0f6be_add_credit_pool.py @@ -10,6 +10,10 @@ import models as models import sqlalchemy as sa from sqlalchemy.dialects import postgresql + +def _is_pg(conn): + return conn.dialect.name == "postgresql" + # revision identifiers, used by Alembic. revision = '7df29de0f6be' down_revision = '03ea244985ce' @@ -19,16 +23,31 @@ depends_on = None def upgrade(): # ### commands auto generated by Alembic - please adjust! ### - op.create_table('tenant_credit_pools', - sa.Column('id', models.types.StringUUID(), server_default=sa.text('uuid_generate_v4()'), nullable=False), - sa.Column('tenant_id', models.types.StringUUID(), nullable=False), - sa.Column('pool_type', sa.String(length=40), server_default='trial', nullable=False), - sa.Column('quota_limit', sa.BigInteger(), nullable=False), - sa.Column('quota_used', sa.BigInteger(), nullable=False), - sa.Column('created_at', sa.DateTime(), server_default=sa.text('CURRENT_TIMESTAMP'), nullable=False), - sa.Column('updated_at', sa.DateTime(), server_default=sa.text('CURRENT_TIMESTAMP'), nullable=False), - sa.PrimaryKeyConstraint('id', name='tenant_credit_pool_pkey') - ) + conn = op.get_bind() + + if _is_pg(conn): + op.create_table('tenant_credit_pools', + sa.Column('id', models.types.StringUUID(), server_default=sa.text('uuid_generate_v4()'), nullable=False), + sa.Column('tenant_id', models.types.StringUUID(), nullable=False), + sa.Column('pool_type', sa.String(length=40), server_default='trial', nullable=False), + sa.Column('quota_limit', sa.BigInteger(), nullable=False), + sa.Column('quota_used', sa.BigInteger(), nullable=False), + sa.Column('created_at', sa.DateTime(), server_default=sa.text('CURRENT_TIMESTAMP'), nullable=False), + sa.Column('updated_at', sa.DateTime(), server_default=sa.text('CURRENT_TIMESTAMP'), nullable=False), + sa.PrimaryKeyConstraint('id', name='tenant_credit_pool_pkey') + ) + else: + # For MySQL and other databases, UUID should be generated at application level + op.create_table('tenant_credit_pools', + sa.Column('id', models.types.StringUUID(), nullable=False), + sa.Column('tenant_id', models.types.StringUUID(), nullable=False), + sa.Column('pool_type', sa.String(length=40), server_default='trial', nullable=False), + sa.Column('quota_limit', sa.BigInteger(), nullable=False), + sa.Column('quota_used', sa.BigInteger(), nullable=False), + sa.Column('created_at', sa.DateTime(), server_default=sa.func.current_timestamp(), nullable=False), + sa.Column('updated_at', sa.DateTime(), server_default=sa.func.current_timestamp(), nullable=False), + sa.PrimaryKeyConstraint('id', name='tenant_credit_pool_pkey') + ) with op.batch_alter_table('tenant_credit_pools', schema=None) as batch_op: batch_op.create_index('tenant_credit_pool_pool_type_idx', ['pool_type'], unique=False) batch_op.create_index('tenant_credit_pool_tenant_id_idx', ['tenant_id'], unique=False) diff --git a/api/models/model.py b/api/models/model.py index c1c6e04ce9..5a274c29cd 100644 --- a/api/models/model.py +++ b/api/models/model.py @@ -2166,7 +2166,9 @@ class TenantCreditPool(TypeBase): sa.Index("tenant_credit_pool_pool_type_idx", "pool_type"), ) - id: Mapped[str] = mapped_column(StringUUID, primary_key=True, server_default=text("uuid_generate_v4()"), init=False) + id: Mapped[str] = mapped_column( + StringUUID, insert_default=lambda: str(uuid4()), default_factory=lambda: str(uuid4()), init=False + ) tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False) pool_type: Mapped[str] = mapped_column(String(40), nullable=False, default="trial", server_default="trial") quota_limit: Mapped[int] = mapped_column(BigInteger, nullable=False, default=0)