add credit pool sys

2026-01-08 07:14:14 +00:00 · 2025-09-26 10:43:34 +08:00
parent e682749d03
commit 2b51fc23d9
11 changed files with 493 additions and 41 deletions
--- a/api/migrations/versions/2025_09_25_1520-58a70d22fdbd_add_table_credit_pool.py
+++ b/api/migrations/versions/2025_09_25_1520-58a70d22fdbd_add_table_credit_pool.py
@@ -0,0 +1,96 @@
+"""add table credit pool
+
+Revision ID: 58a70d22fdbd
+Revises: 68519ad5cd18
+Create Date: 2025-09-25 15:20:40.367078
+
+"""
+from alembic import op
+import models as models
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = '58a70d22fdbd'
+down_revision = '68519ad5cd18'
+branch_labels = None
+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), 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')
+    )
+    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)
+    # Data migration: Move trial quota data from providers to tenant_credit_pools
+    migrate_trial_quota_data()
+    
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table('tenant_credit_pools', schema=None) as batch_op:
+        batch_op.drop_index('tenant_credit_pool_tenant_id_idx')
+        batch_op.drop_index('tenant_credit_pool_pool_type_idx')
+
+    op.drop_table('tenant_credit_pools')
+    # ### end Alembic commands ###
+
+
+def migrate_trial_quota_data():
+    """
+    Migrate quota data from providers table to tenant_credit_pools table
+    for providers with quota_type='trial', provider_name='openai', provider_type='system'
+    """
+    # Create connection
+    bind = op.get_bind()
+    
+    # Query providers that match the criteria
+    select_sql = sa.text("""
+        SELECT tenant_id, quota_limit, quota_used
+        FROM providers 
+        WHERE quota_type = 'trial' 
+        AND provider_name = 'openai' 
+        AND provider_type = 'system'
+        AND quota_limit IS NOT NULL
+    """)
+    
+    result = bind.execute(select_sql)
+    providers_data = result.fetchall()
+    
+    # Insert data into tenant_credit_pools
+    for provider_data in providers_data:
+        tenant_id, quota_limit, quota_used = provider_data
+        
+        # Check if credit pool already exists for this tenant
+        check_sql = sa.text("""
+            SELECT COUNT(*) 
+            FROM tenant_credit_pools 
+            WHERE tenant_id = :tenant_id AND pool_type = 'trial'
+        """)
+        
+        existing_count = bind.execute(check_sql, {"tenant_id": tenant_id}).scalar()
+        
+        if existing_count == 0:
+            # Insert new credit pool record
+            insert_sql = sa.text("""
+                INSERT INTO tenant_credit_pools (tenant_id, pool_type, quota_limit, quota_used, created_at, updated_at)
+                VALUES (:tenant_id, 'trial', :quota_limit, :quota_used, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)
+            """)
+            
+            bind.execute(insert_sql, {
+                "tenant_id": tenant_id,
+                "quota_limit": quota_limit or 0,
+                "quota_used": quota_used or 0
+            })