From 90588ebae8151a6450e65b4521cf09e8961e44a9 Mon Sep 17 00:00:00 2001
From: Wu Clan <jianhengwu0407@gmail.com>
Date: Thu, 30 Oct 2025 19:24:46 +0800
Subject: [PATCH] Add the celery rabbitmq vhost config (#894)

---
 backend/app/task/celery.py | 15 ++++++++++-----
 backend/core/conf.py       |  1 +
 2 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/backend/app/task/celery.py b/backend/app/task/celery.py
index f21f0ac9..e0f403b3 100644
--- a/backend/app/task/celery.py
+++ b/backend/app/task/celery.py
@@ -27,15 +27,20 @@ def init_celery() -> celery.Celery:
     celery.app.trace.build_tracer = celery_aio_pool.build_async_tracer
     celery.app.trace.reset_worker_optimizations()
 
+    broker_url = f'amqp://{settings.CELERY_RABBITMQ_USERNAME}:{settings.CELERY_RABBITMQ_PASSWORD}@{settings.CELERY_RABBITMQ_HOST}:{settings.CELERY_RABBITMQ_PORT}/{settings.CELERY_RABBITMQ_VHOST}'
+    if settings.CELERY_BROKER == 'redis':
+        broker_url = f'redis://:{settings.REDIS_PASSWORD}@{settings.REDIS_HOST}:{settings.REDIS_PORT}/{settings.CELERY_BROKER_REDIS_DATABASE}'
+
+    result_backend = f'db+postgresql+psycopg://{settings.DATABASE_USER}:{settings.DATABASE_PASSWORD}@{settings.DATABASE_HOST}:{settings.DATABASE_PORT}/{settings.DATABASE_SCHEMA}'
+    if settings.DATABASE_TYPE == 'mysql':
+        result_backend = result_backend.replace('postgresql+psycopg', 'mysql+pymysql')
+
     # https://docs.celeryq.dev/en/stable/userguide/configuration.html
     app = celery.Celery(
         'fba_celery',
-        broker_url=f'redis://:{settings.REDIS_PASSWORD}@{settings.REDIS_HOST}:{settings.REDIS_PORT}/{settings.CELERY_BROKER_REDIS_DATABASE}'
-        if settings.CELERY_BROKER == 'redis'
-        else f'amqp://{settings.CELERY_RABBITMQ_USERNAME}:{settings.CELERY_RABBITMQ_PASSWORD}@{settings.CELERY_RABBITMQ_HOST}:{settings.CELERY_RABBITMQ_PORT}',
+        broker_url=broker_url,
         broker_connection_retry_on_startup=True,
-        result_backend=f'db+{settings.DATABASE_TYPE}+{"pymysql" if settings.DATABASE_TYPE == "mysql" else "psycopg"}'
-        f'://{settings.DATABASE_USER}:{settings.DATABASE_PASSWORD}@{settings.DATABASE_HOST}:{settings.DATABASE_PORT}/{settings.DATABASE_SCHEMA}',
+        result_backend=result_backend,
         result_extended=True,
         database_engine_options={'echo': settings.DATABASE_ECHO},
         # result_expires=0,
diff --git a/backend/core/conf.py b/backend/core/conf.py
index bd290d7d..0ca91655 100644
--- a/backend/core/conf.py
+++ b/backend/core/conf.py
@@ -210,6 +210,7 @@ class Settings(BaseSettings):
 
     # 基础配置
     CELERY_BROKER: Literal['rabbitmq', 'redis'] = 'redis'
+    CELERY_RABBITMQ_VHOST: str = ''
     CELERY_REDIS_PREFIX: str = 'fba:celery'
     CELERY_TASK_MAX_RETRIES: int = 5