danswer/backend/shared_configs/contextvars.py
pablonyx a98dcbc7de
Update tenant logic (#4122)
* k

* k

* k

* quick nit

* nit
2025-02-26 03:53:46 +00:00

25 lines
679 B
Python

import contextvars
from shared_configs.configs import MULTI_TENANT
from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA
# Context variable for the current tenant id
CURRENT_TENANT_ID_CONTEXTVAR: contextvars.ContextVar[
str | None
] = contextvars.ContextVar(
"current_tenant_id", default=None if MULTI_TENANT else POSTGRES_DEFAULT_SCHEMA
)
"""Utils related to contextvars"""
def get_current_tenant_id() -> str:
tenant_id = CURRENT_TENANT_ID_CONTEXTVAR.get()
if tenant_id is None:
if not MULTI_TENANT:
return POSTGRES_DEFAULT_SCHEMA
raise RuntimeError("Tenant ID is not set. This should never happen.")
return tenant_id