From 02251aab75bad74647ba526654950b131748eb45 Mon Sep 17 00:00:00 2001 From: pablodanswer Date: Thu, 21 Nov 2024 11:09:11 -0800 Subject: [PATCH] update --- backend/danswer/background/celery/apps/app_base.py | 10 ++++++++-- .../tests/personas/test_persona_categories.py | 1 + 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/backend/danswer/background/celery/apps/app_base.py b/backend/danswer/background/celery/apps/app_base.py index 5e5b24d2d50..6c448f64449 100644 --- a/backend/danswer/background/celery/apps/app_base.py +++ b/backend/danswer/background/celery/apps/app_base.py @@ -35,6 +35,8 @@ from danswer.utils.logger import ColoredFormatter from danswer.utils.logger import PlainFormatter from danswer.utils.logger import setup_logger +from shared_configs.configs import MULTI_TENANT +from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA from shared_configs.configs import SENTRY_DSN from shared_configs.configs import TENANT_ID_PREFIX from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR @@ -443,7 +445,11 @@ def set_tenant_id( **other_kwargs: Any, ) -> None: """Signal handler to set tenant ID in context var before task starts.""" - tenant_id = kwargs.get("tenant_id") if kwargs else None + tenant_id = ( + kwargs.get("tenant_id", POSTGRES_DEFAULT_SCHEMA) + if kwargs + else POSTGRES_DEFAULT_SCHEMA + ) CURRENT_TENANT_ID_CONTEXTVAR.set(tenant_id) @@ -457,4 +463,4 @@ def reset_tenant_id( **other_kwargs: Any, ) -> None: """Signal handler to reset tenant ID in context var after task ends.""" - CURRENT_TENANT_ID_CONTEXTVAR.reset() + CURRENT_TENANT_ID_CONTEXTVAR.set(POSTGRES_DEFAULT_SCHEMA) diff --git a/backend/tests/integration/tests/personas/test_persona_categories.py b/backend/tests/integration/tests/personas/test_persona_categories.py index fdd0e645814..1ac2d3b3000 100644 --- a/backend/tests/integration/tests/personas/test_persona_categories.py +++ b/backend/tests/integration/tests/personas/test_persona_categories.py @@ -44,6 +44,7 @@ def test_persona_category_management(reset: None) -> None: category=updated_persona_category, user_performing_action=regular_user, ) + assert exc_info.value.response is not None assert exc_info.value.response.status_code == 403 assert PersonaCategoryManager.verify(