onyx/backend/shared_configs/contextvars.py

import contextvars

from shared_configs.configs import MULTI_TENANT, 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
    )
)

# set by every route in the API server
INDEXING_REQUEST_ID_CONTEXTVAR: contextvars.ContextVar[str | None] = (
    contextvars.ContextVar("indexing_request_id", default=None)
)

# set by every route in the API server
ONYX_REQUEST_ID_CONTEXTVAR: contextvars.ContextVar[str | None] = contextvars.ContextVar(
    "onyx_request_id", default=None
)

# Used to store cc pair id and index attempt id in multithreaded environments
INDEX_ATTEMPT_INFO_CONTEXTVAR: contextvars.ContextVar[tuple[int, int] | None] = (
    contextvars.ContextVar("index_attempt_info", default=None)
)

# Set by endpoint context middleware — used for per-endpoint DB pool attribution
CURRENT_ENDPOINT_CONTEXTVAR: contextvars.ContextVar[str | None] = (
    contextvars.ContextVar("current_endpoint", default=None)
)

# Per-request user id for usage attribution; None in workers.
CURRENT_USER_ID_CONTEXTVAR: contextvars.ContextVar[str | None] = contextvars.ContextVar(
    "current_user_id", default=None
)


def get_current_tenant_id() -> str:
    tenant_id = CURRENT_TENANT_ID_CONTEXTVAR.get()
    if tenant_id is None:
        import traceback

        if not MULTI_TENANT:
            return POSTGRES_DEFAULT_SCHEMA

        stack_trace = traceback.format_stack()
        error_message = (
            "Tenant ID is not set. This should never happen.\nStack trace:\n"
            + "".join(stack_trace)
        )
        raise RuntimeError(error_message)
    return tenant_id


def get_current_user_id() -> str | None:
    """Requesting user's id, or None outside a per-request context."""
    return CURRENT_USER_ID_CONTEXTVAR.get()