fix: provider create cause IntegrityError (#1866)

This commit is contained in:
takatost 2024-01-03 09:12:53 +08:00 committed by GitHub
parent 4562e83b24
commit 2381264a3f
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -3,6 +3,8 @@ from collections import defaultdict
from json import JSONDecodeError
from typing import Optional
from psycopg2 import IntegrityError
from core.entities.model_entities import DefaultModelEntity, DefaultModelProviderEntity
from core.entities.provider_configuration import ProviderConfigurations, ProviderConfiguration, ProviderModelBundle
from core.entities.provider_entities import CustomConfiguration, CustomProviderConfiguration, CustomModelConfiguration, \
@ -380,6 +382,7 @@ class ProviderManager:
if quota.quota_type == ProviderQuotaType.TRIAL:
# Init trial provider records if not exists
if ProviderQuotaType.TRIAL not in provider_quota_to_provider_record_dict:
try:
provider_record = Provider(
tenant_id=tenant_id,
provider_name=provider_name,
@ -391,6 +394,16 @@ class ProviderManager:
)
db.session.add(provider_record)
db.session.commit()
except IntegrityError:
db.session.rollback()
provider_record = db.session.query(Provider) \
.filter(
Provider.tenant_id == tenant_id,
Provider.provider_name == provider_name,
Provider.provider_type == ProviderType.SYSTEM.value,
Provider.quota_type == ProviderQuotaType.TRIAL.value,
Provider.is_valid == True
).first()
provider_name_to_provider_records_dict[provider_name].append(provider_record)