mirror of
https://gh-proxy.org/https://github.com/klinsmaya/SQLBot.git
synced 2026-11-08 01:06:59 +08:00
perf: Assistant cors config
This commit is contained in:
@@ -1,12 +1,15 @@
|
||||
|
||||
|
||||
from fastapi import FastAPI
|
||||
from sqlmodel import Session, select
|
||||
from apps.datasource.models.datasource import CoreDatasource
|
||||
from apps.system.models.system_model import AssistantModel
|
||||
from apps.system.schemas.auth import CacheName, CacheNamespace
|
||||
from apps.system.schemas.system_schema import UserInfoDTO
|
||||
from common.core.sqlbot_cache import cache
|
||||
|
||||
from common.core.db import engine
|
||||
from starlette.middleware.cors import CORSMiddleware
|
||||
from common.core.config import settings
|
||||
|
||||
@cache(namespace=CacheNamespace.EMBEDDED_INFO, cacheName=CacheName.ASSISTANT_INFO, keyExpression="assistant_id")
|
||||
async def get_assistant_info(*, session: Session, assistant_id: int) -> AssistantModel | None:
|
||||
@@ -23,3 +26,28 @@ def get_assistant_ds(*, session: Session, assistant: AssistantModel):
|
||||
# filter private ds if offline
|
||||
return db_ds_list
|
||||
pass
|
||||
|
||||
|
||||
def init_dynamic_cors(app: FastAPI):
|
||||
try:
|
||||
with Session(engine) as session:
|
||||
list_result = session.exec(select(AssistantModel).order_by(AssistantModel.create_time)).all()
|
||||
seen = set()
|
||||
unique_domains = []
|
||||
for item in list_result:
|
||||
if item.domain:
|
||||
for domain in item.domain.split(','):
|
||||
domain = domain.strip()
|
||||
if domain and domain not in seen:
|
||||
seen.add(domain)
|
||||
unique_domains.append(domain)
|
||||
cors_middleware = None
|
||||
for middleware in app.user_middleware:
|
||||
if middleware.cls == CORSMiddleware:
|
||||
cors_middleware = middleware
|
||||
break
|
||||
if cors_middleware:
|
||||
updated_origins = list(set(settings.all_cors_origins + unique_domains))
|
||||
cors_middleware.kwargs['allow_origins'] = updated_origins
|
||||
except Exception as e:
|
||||
return False, e
|
||||
@@ -6,6 +6,7 @@ from fastapi.routing import APIRoute
|
||||
from starlette.middleware.cors import CORSMiddleware
|
||||
from starlette.exceptions import HTTPException as StarletteHTTPException
|
||||
from apps.api import api_router
|
||||
from apps.system.crud.assistant import init_dynamic_cors
|
||||
from apps.system.middleware.auth import TokenMiddleware
|
||||
from common.core.config import settings
|
||||
from common.core.response_middleware import ResponseMiddleware, exception_handler
|
||||
@@ -27,6 +28,7 @@ async def lifespan(app: FastAPI):
|
||||
run_migrations()
|
||||
FastAPICache.init(InMemoryBackend())
|
||||
print("✅ FastAPICache 初始化完成")
|
||||
init_dynamic_cors(app)
|
||||
yield
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user