Strict Tenant ID Enforcement #3871

pablonyx · 2025-02-01T01:05:45Z

Description

fixes https://linear.app/danswer/issue/DAN-1390/clean-up-of-tenant-id-assumption

New Default behavior / assumptions
- When a request comes in, the HTTP middleware sets CURRENT_TENANT_ID_CONTEXTVAR to the correct tenant ID.
- In Celery tasks, we override the Task class to be TenantAwareTask and you pass tenant_id=... in kwargs, the task automatically sets the contextvar before running.
- This means that, most of the time, you do not need to pass tenant_id around yourself. Just call:
```
with get_session_with_current_tenant() as db_session:
    ...
```
  or
```
r = get_redis_client()
```
  and it will pick up the ID from the contextvar.
Shared utils
- For truly global logic not tied to any tenant, use these dedicated helpers:
  - get_session_with_shared_schema(): forces usage of the “public” schema in Postgres.
  - get_shared_redis_client() / get_shared_redis_replica_client(): skip per-tenant prefixes in Redis.
  - get_shared_kv_store(): global K/V store that ignores the per-tenant context.
Overriding Tenant ID
- If you’re doing multi-tenant “manager” work, you can explicitly pass a tenant ID:
```
with get_session_with_current_tenant("some_tenant") as db_session:
    ...
```
- This directly sets the contextvar to some_tenant for the duration of that block.
BG processes

Celery uses TenantAwareTask:

@shared_task(base=TenantAwareTask)
def my_task(*, tenant_id: str):
    # tenant_id is picked up and set into the contextvar automatically
    ...

When you call my_task.delay(tenant_id="..."), Celery sets that ID in the contextvar before running the task body.

How Has This Been Tested?

[Describe the tests you ran to verify your changes]

Backporting (check the box to trigger backport action)

Note: You have to check that the action passes, otherwise resolve the conflicts manually and tag the patches.

This PR should be backported (make sure to check that the backport attempt succeeds)
[Optional] Override Linear Check

vercel · 2025-02-01T01:05:50Z

The latest updates on your projects. Learn more about Vercel for Git ↗︎

Name	Status	Preview	Comments	Updated (UTC)
internal-search	✅ Ready (Inspect)	Visit Preview	💬 Add feedback	Feb 6, 2025 6:35am

Weves

As discussed, let's have a separate function for accessing a Postgres / Redis session with the public tenant to not overload the get_session_with_tenant etc. Let's also remove the tenant_id override option to enforce consistency.

Weves · 2025-02-01T18:51:31Z

backend/shared_configs/contextvars.py

 )


 """Utils related to contextvars"""


-def get_current_tenant_id() -> str:
+def current_tenant_id(strict: bool = True) -> str:


nit: prefer get_current_tenant_id for consistency. We usually have a get / fetch prefix (e.g. make them verbs) for functions. For example get_llm / get_session / get_redis_client.

also is there ever a reason to pass in strict=False?

Definitely agree in the abstract - didn't set this originally since we have a dependency called get_current_tenant_id, but with this refactory I think it makes sense to remove that function anyways.

Weves · 2025-02-01T18:54:09Z

backend/shared_configs/contextvars.py

+CURRENT_TENANT_ID_CONTEXTVAR: contextvars.ContextVar[
+    str | None
+] = contextvars.ContextVar(
+    "current_tenant_id", default=None if MULTI_TENANT else POSTGRES_DEFAULT_SCHEMA


thoughts on this always being None and enforcing that it be set on both multi and single tenant to bring behavior closer together? E.g. middleware always sets this to POSTGRES_DEFAULT_SCHEMA.

Not sure if this is better, but just a thought.

rkuo-danswer · 2025-02-08T23:03:02Z

backend/ee/onyx/background/celery/apps/primary.py

@@ -21,7 +21,7 @@
 def perform_ttl_management_task(
    retention_limit_days: int, *, tenant_id: str | None
 ) -> None:
-    with get_session_with_tenant(tenant_id) as db_session:
+    with get_session_with_current_tenant(tenant_id) as db_session:


i don't believe saying "current tenant" has any semantic meaning to the function name when we're passing in a tenant id. if the function takes a tenant id, it doesn't care where the tenant id came from.

backend/ee/onyx/server/middleware/tenant_tracking.py

rkuo-danswer · 2025-02-08T23:11:27Z

backend/onyx/background/celery/apps/beat.py

@@ -19,6 +19,7 @@

 celery_app = Celery(__name__)
 celery_app.config_from_object("onyx.background.celery.configs.beat")
+celery_app.Task = app_base.TenantAwareTask  # type: ignore [misc]


Monkey patching Task seems too invasive as there are already built in ways to subclass Task behavior. Wouldn't something like https://docs.celeryq.dev/en/stable/userguide/tasks.html#task-inheritance accomplish this just fine?

rkuo-danswer · 2025-02-08T23:18:48Z

backend/onyx/background/celery/tasks/indexing/tasks.py

@@ -68,12 +69,15 @@
 def check_for_indexing(self: Task, *, tenant_id: str | None) -> int | None:
    """a lightweight task used to kick off indexing tasks.
    Occcasionally does some validation of existing state to clear up error conditions"""
+    print(f"TENANT ID, {tenant_id}")


debug logging?

rkuo-danswer · 2025-02-08T23:24:30Z

backend/onyx/db/engine.py

+
+        with Session(bind=connection, expire_on_commit=False) as session:
+            try:
+                yield session


i've seen some infailedsqltransaction errors involving these search_path operations and i'm wondering if we might need a rollback somewhere. Don't have an exact direction to go on this but might be worth investigating.

backend/onyx/setup.py

vercel bot deployed to Preview February 1, 2025 01:08 View deployment

vercel bot deployed to Preview February 1, 2025 01:55 View deployment

Weves reviewed Feb 1, 2025

View reviewed changes

vercel bot deployed to Preview February 1, 2025 20:59 View deployment

pablonyx force-pushed the tenant_enforcement branch from dab29ee to 3b3898a Compare February 1, 2025 22:16

pablonyx enabled auto-merge February 1, 2025 22:16

vercel bot deployed to Preview February 1, 2025 22:25 View deployment

vercel bot deployed to Preview February 5, 2025 02:31 View deployment

pablonyx force-pushed the tenant_enforcement branch from eaaadfd to 9e8e87e Compare February 5, 2025 17:51

vercel bot deployed to Preview February 5, 2025 17:55 View deployment

vercel bot deployed to Preview February 5, 2025 18:45 View deployment

pablonyx added 8 commits February 5, 2025 12:23

update

12c8993

update

a31a454

various updates

40f83fd

update

15b8d29

shared task override

d31e80f

save for now

a70304f

update + functional background processing

5c886bd

temporary mods

bb06274

pablonyx force-pushed the tenant_enforcement branch from deb1a8b to bb06274 Compare February 5, 2025 20:23

vercel bot deployed to Preview February 5, 2025 20:28 View deployment

quick fix

df5bd4c

vercel bot deployed to Preview February 5, 2025 20:53 View deployment

update get redis client

3ab9081

vercel bot deployed to Preview February 5, 2025 21:56 View deployment

mostly functioning!

0555bba

vercel bot deployed to Preview February 5, 2025 22:30 View deployment

pablonyx added 2 commits February 5, 2025 15:08

remove logs

210779b

misc

d7123fd

vercel bot deployed to Preview February 5, 2025 23:19 View deployment

k

38726ab

vercel bot deployed to Preview February 6, 2025 06:35 View deployment

rkuo-danswer reviewed Feb 8, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Strict Tenant ID Enforcement #3871

Strict Tenant ID Enforcement #3871

pablonyx commented Feb 1, 2025 •

edited

Loading

vercel bot commented Feb 1, 2025 •

edited

Loading

Weves left a comment

Weves Feb 1, 2025

Weves Feb 1, 2025

pablonyx Feb 5, 2025

Weves Feb 1, 2025

rkuo-danswer Feb 8, 2025

rkuo-danswer Feb 8, 2025

rkuo-danswer Feb 8, 2025

rkuo-danswer Feb 8, 2025

Strict Tenant ID Enforcement #3871

Are you sure you want to change the base?

Strict Tenant ID Enforcement #3871

Conversation

pablonyx commented Feb 1, 2025 • edited Loading

Description

How Has This Been Tested?

Backporting (check the box to trigger backport action)

vercel bot commented Feb 1, 2025 • edited Loading

Weves left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

pablonyx commented Feb 1, 2025 •

edited

Loading

vercel bot commented Feb 1, 2025 •

edited

Loading