opensource/dify
2
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-09-08 02:43:49 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor: use EnumText for TenantCreditPool.pool_type (#33959)

Browse Source
This commit is contained in:
tmimmanuel 2026-03-23 18:51:10 +01:00 committed by GitHub
parent 8ca1ebb96d
commit 75c3ef82d9
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 16 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
api/core/provider_manager.py
View File

@ -918,11 +918,11 @@ class ProviderManager:
trail_pool = CreditPoolService.get_pool( trail_pool = CreditPoolService.get_pool(
tenant_id=tenant_id, tenant_id=tenant_id,
pool_type=ProviderQuotaType.TRIAL.value, pool_type=ProviderQuotaType.TRIAL,
) )
paid_pool = CreditPoolService.get_pool( paid_pool = CreditPoolService.get_pool(
tenant_id=tenant_id, tenant_id=tenant_id,
pool_type=ProviderQuotaType.PAID.value, pool_type=ProviderQuotaType.PAID,
) )
else: else:
trail_pool = None trail_pool = None

5
api/models/model.py
View File

@ -44,6 +44,7 @@ from .enums import (
MessageChainType, MessageChainType,
MessageFileBelongsTo, MessageFileBelongsTo,
MessageStatus, MessageStatus,
ProviderQuotaType,
TagType, TagType,
) )
from .provider_ids import GenericProviderID from .provider_ids import GenericProviderID
@ -2491,7 +2492,9 @@ class TenantCreditPool(TypeBase):
StringUUID, insert_default=lambda: str(uuid4()), default_factory=lambda: str(uuid4()), init=False StringUUID, insert_default=lambda: str(uuid4()), default_factory=lambda: str(uuid4()), init=False
) )
tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=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") pool_type: Mapped[ProviderQuotaType] = mapped_column(
EnumText(ProviderQuotaType, length=40), nullable=False, default=ProviderQuotaType.TRIAL, server_default="trial"
)
quota_limit: Mapped[int] = mapped_column(BigInteger, nullable=False, default=0) quota_limit: Mapped[int] = mapped_column(BigInteger, nullable=False, default=0)
quota_used: Mapped[int] = mapped_column(BigInteger, nullable=False, default=0) quota_used: Mapped[int] = mapped_column(BigInteger, nullable=False, default=0)
created_at: Mapped[datetime] = mapped_column( created_at: Mapped[datetime] = mapped_column(

6
api/services/credit_pool_service.py
View File

@ -7,6 +7,7 @@ from configs import dify_config
from core.errors.error import QuotaExceededError from core.errors.error import QuotaExceededError
from extensions.ext_database import db from extensions.ext_database import db
from models import TenantCreditPool from models import TenantCreditPool
from models.enums import ProviderQuotaType
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@ -16,7 +17,10 @@ class CreditPoolService:
def create_default_pool(cls, tenant_id: str) -> TenantCreditPool: def create_default_pool(cls, tenant_id: str) -> TenantCreditPool:
"""create default credit pool for new tenant""" """create default credit pool for new tenant"""
credit_pool = TenantCreditPool( credit_pool = TenantCreditPool(
tenant_id=tenant_id, quota_limit=dify_config.HOSTED_POOL_CREDITS, quota_used=0, pool_type="trial" tenant_id=tenant_id,
quota_limit=dify_config.HOSTED_POOL_CREDITS,
quota_used=0,
pool_type=ProviderQuotaType.TRIAL,
) )
db.session.add(credit_pool) db.session.add(credit_pool)
db.session.commit() db.session.commit()

9
api/tests/test_containers_integration_tests/services/test_credit_pool_service.py
View File

@ -6,6 +6,7 @@ import pytest
from core.errors.error import QuotaExceededError from core.errors.error import QuotaExceededError
from models import TenantCreditPool from models import TenantCreditPool
from models.enums import ProviderQuotaType
from services.credit_pool_service import CreditPoolService from services.credit_pool_service import CreditPoolService
@ -20,7 +21,7 @@ class TestCreditPoolService:
assert isinstance(pool, TenantCreditPool) assert isinstance(pool, TenantCreditPool)
assert pool.tenant_id == tenant_id assert pool.tenant_id == tenant_id
assert pool.pool_type == "trial" assert pool.pool_type == ProviderQuotaType.TRIAL
assert pool.quota_used == 0 assert pool.quota_used == 0
assert pool.quota_limit > 0 assert pool.quota_limit > 0
@ -28,14 +29,14 @@ class TestCreditPoolService:
tenant_id = self._create_tenant_id() tenant_id = self._create_tenant_id()
CreditPoolService.create_default_pool(tenant_id) CreditPoolService.create_default_pool(tenant_id)
result = CreditPoolService.get_pool(tenant_id=tenant_id, pool_type="trial") result = CreditPoolService.get_pool(tenant_id=tenant_id, pool_type=ProviderQuotaType.TRIAL)
assert result is not None assert result is not None
assert result.tenant_id == tenant_id assert result.tenant_id == tenant_id
assert result.pool_type == "trial" assert result.pool_type == ProviderQuotaType.TRIAL
def test_get_pool_returns_none_when_not_exists(self, db_session_with_containers): def test_get_pool_returns_none_when_not_exists(self, db_session_with_containers):
result = CreditPoolService.get_pool(tenant_id=self._create_tenant_id(), pool_type="trial") result = CreditPoolService.get_pool(tenant_id=self._create_tenant_id(), pool_type=ProviderQuotaType.TRIAL)
assert result is None assert result is None