From 82a91c0cfcf40e949c8fe01c2de97fc73e00f6b0 Mon Sep 17 00:00:00 2001 From: yaojin Date: Wed, 5 Aug 2026 15:13:47 +0800 Subject: [PATCH] fix: inject tenant scope into ORM queries --- backend/app/dao/base.py | 47 +++++++++++++++++++++- backend/app/models/agent.py | 1 + backend/app/models/audit.py | 2 + backend/app/models/experience.py | 1 + backend/app/models/experience_reference.py | 1 + backend/app/models/invitation_code.py | 1 + backend/app/models/notification.py | 1 + backend/app/models/org.py | 2 + backend/app/models/plaza.py | 1 + backend/app/models/published_page.py | 1 + backend/app/models/task.py | 1 + backend/app/models/user.py | 1 + backend/tests/test_base_dao.py | 38 ++++++++++++++++- 13 files changed, 94 insertions(+), 4 deletions(-) diff --git a/backend/app/dao/base.py b/backend/app/dao/base.py index 61aef9fad..db4341a31 100644 --- a/backend/app/dao/base.py +++ b/backend/app/dao/base.py @@ -4,8 +4,9 @@ from contextvars import ContextVar from typing import Any, Generic, Type, TypeVar -from sqlalchemy import select +from sqlalchemy import event, select from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy.orm import Session, with_loader_criteria from app.database import Base, _session_ctx, async_session @@ -108,6 +109,49 @@ async def delete(self, *, id: Any) -> ModelType | None: _tenant_ctx: ContextVar[uuid.UUID | None] = ContextVar("tenant_ctx", default=None) +def _is_tenant_scoped_model(model: type[Base]) -> bool: + """Return whether model rows must be isolated whenever tenant context exists. + + Non-null ``tenant_id`` columns are tenant-owned by schema. Legacy tables + whose tenant column is nullable can opt in with ``__tenant_scoped__ = True`` + while their historic, tenant-less rows remain readable only outside a tenant + context (for example during migration or platform administration). + """ + if getattr(model, "__tenant_scoped__", False): + return True + tenant_column = model.__table__.c.get("tenant_id") + return tenant_column is not None and not tenant_column.nullable + + +@event.listens_for(Session, "do_orm_execute") +def _inject_tenant_scope(execute_state: Any) -> None: + """Apply the active tenant predicate to every tenant-owned ORM SELECT. + + This is deliberately installed on SQLAlchemy's synchronous ``Session`` + class, which is also the execution layer below ``AsyncSession``. It covers + direct API/service queries and DAO queries alike, so a missed business-level + ``tenant_id`` filter cannot disclose another tenant's rows. + """ + if not execute_state.is_select: + return + tenant_id = _tenant_ctx.get() + if tenant_id is None: + return + + statement = execute_state.statement + for mapper in execute_state.all_mappers: + model = mapper.class_ + if _is_tenant_scoped_model(model): + statement = statement.options( + with_loader_criteria( + model, + lambda cls: cls.tenant_id == tenant_id, + include_aliases=True, + ) + ) + execute_state.statement = statement + + @contextmanager def tenant_context(tenant_id: uuid.UUID): """Explicitly bind a tenant_id to the current coroutine context. @@ -193,4 +237,3 @@ async def delete_scoped(self, *, id: Any) -> ModelType | None: await db.delete(obj) await db.flush() return obj - diff --git a/backend/app/models/agent.py b/backend/app/models/agent.py index bd4e56e92..d2e88f7ae 100644 --- a/backend/app/models/agent.py +++ b/backend/app/models/agent.py @@ -23,6 +23,7 @@ class Agent(Base): """ __tablename__ = "agents" + __tenant_scoped__ = True __table_args__ = ( Index( "ix_agents_active_tenant_created_at", diff --git a/backend/app/models/audit.py b/backend/app/models/audit.py index 1f8ce517f..71181a408 100644 --- a/backend/app/models/audit.py +++ b/backend/app/models/audit.py @@ -14,6 +14,7 @@ class AuditLog(Base): """Audit trail for all operations.""" __tablename__ = "audit_logs" + __tenant_scoped__ = True id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4) tenant_id: Mapped[uuid.UUID | None] = mapped_column( @@ -50,6 +51,7 @@ class ChatMessage(Base): """Message on the unified chat substrate.""" __tablename__ = "chat_messages" + __tenant_scoped__ = True id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4) tenant_id: Mapped[uuid.UUID | None] = mapped_column( diff --git a/backend/app/models/experience.py b/backend/app/models/experience.py index f97b65ca3..8b61ff322 100644 --- a/backend/app/models/experience.py +++ b/backend/app/models/experience.py @@ -33,6 +33,7 @@ class ExperienceEntry(Base): """ __tablename__ = "experience_entries" + __tenant_scoped__ = True id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4) # A draft created while editing a published/retired entry. Publishing the diff --git a/backend/app/models/experience_reference.py b/backend/app/models/experience_reference.py index 3db65b266..88a4b76ad 100644 --- a/backend/app/models/experience_reference.py +++ b/backend/app/models/experience_reference.py @@ -22,6 +22,7 @@ class ExperienceReference(Base): """One reuse event of an experience entry by an agent.""" __tablename__ = "experience_references" + __tenant_scoped__ = True id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4) entry_id: Mapped[uuid.UUID] = mapped_column( diff --git a/backend/app/models/invitation_code.py b/backend/app/models/invitation_code.py index c7160cfbd..9288e7fb3 100644 --- a/backend/app/models/invitation_code.py +++ b/backend/app/models/invitation_code.py @@ -14,6 +14,7 @@ class InvitationCode(Base): """An invitation code that can be used to register new accounts.""" __tablename__ = "invitation_codes" + __tenant_scoped__ = True id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4) code: Mapped[str] = mapped_column(String(32), unique=True, nullable=False, index=True) diff --git a/backend/app/models/notification.py b/backend/app/models/notification.py index 722066786..7fc3a9b54 100644 --- a/backend/app/models/notification.py +++ b/backend/app/models/notification.py @@ -14,6 +14,7 @@ class Notification(Base): """A notification delivered to a user or an agent.""" __tablename__ = "notifications" + __tenant_scoped__ = True id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4) tenant_id: Mapped[uuid.UUID | None] = mapped_column( diff --git a/backend/app/models/org.py b/backend/app/models/org.py index b06930d47..df994b125 100644 --- a/backend/app/models/org.py +++ b/backend/app/models/org.py @@ -14,6 +14,7 @@ class OrgDepartment(Base): """Department from Feishu org structure.""" __tablename__ = "org_departments" + __tenant_scoped__ = True id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4) external_id: Mapped[str | None] = mapped_column(String(100), index=True) @@ -35,6 +36,7 @@ class OrgMember(Base): """Person from an identity provider's org structure.""" __tablename__ = "org_members" + __tenant_scoped__ = True id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4) diff --git a/backend/app/models/plaza.py b/backend/app/models/plaza.py index fdfaa2982..5c47b50a9 100644 --- a/backend/app/models/plaza.py +++ b/backend/app/models/plaza.py @@ -14,6 +14,7 @@ class PlazaPost(Base): """A post in the Agent Plaza social feed.""" __tablename__ = "plaza_posts" + __tenant_scoped__ = True id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4) author_id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), nullable=False, index=True) diff --git a/backend/app/models/published_page.py b/backend/app/models/published_page.py index f2cf89780..af672c029 100644 --- a/backend/app/models/published_page.py +++ b/backend/app/models/published_page.py @@ -14,6 +14,7 @@ class PublishedPage(Base): """A publicly accessible HTML page published from an agent workspace.""" __tablename__ = "published_pages" + __tenant_scoped__ = True id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4) short_id: Mapped[str] = mapped_column(String(16), unique=True, index=True, nullable=False) diff --git a/backend/app/models/task.py b/backend/app/models/task.py index aa43d86e8..811314dca 100644 --- a/backend/app/models/task.py +++ b/backend/app/models/task.py @@ -14,6 +14,7 @@ class Task(Base): """Task assigned to or managed by a digital employee.""" __tablename__ = "tasks" + __tenant_scoped__ = True id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4) tenant_id: Mapped[uuid.UUID | None] = mapped_column( diff --git a/backend/app/models/user.py b/backend/app/models/user.py index 2ef4c0624..734802dae 100644 --- a/backend/app/models/user.py +++ b/backend/app/models/user.py @@ -55,6 +55,7 @@ class User(Base): """ __tablename__ = "users" + __tenant_scoped__ = True # Note: Unique constraints for (tenant_id, username), (tenant_id, email) and (tenant_id, primary_mobile) # are handled via partial unique indexes in migration to allow NULL values diff --git a/backend/tests/test_base_dao.py b/backend/tests/test_base_dao.py index c54d922c5..d1ed2f582 100644 --- a/backend/tests/test_base_dao.py +++ b/backend/tests/test_base_dao.py @@ -1,15 +1,27 @@ from types import SimpleNamespace +import uuid import pytest +from sqlalchemy import String, create_engine, select +from sqlalchemy.orm import Mapped, Session, mapped_column -from app.dao.base import BaseDAO -from app.database import _session_ctx +from app.dao.base import BaseDAO, tenant_context +from app.database import Base, _session_ctx class DummyModel: id = "id" +class TenantScopedRecord(Base): + """Small mapped record proving the session-level isolation hook.""" + + __tablename__ = "test_tenant_scoped_records" + + id: Mapped[str] = mapped_column(String, primary_key=True) + tenant_id: Mapped[str] = mapped_column(String, nullable=False) + + class RecordingSession: def __init__(self): self.added = [] @@ -102,3 +114,25 @@ async def test_delete_uses_current_session_without_nested_lookup(monkeypatch): assert session.deleted == [session.object_to_get] assert session.flushed is True assert session.committed is True + + +def test_orm_session_injects_tenant_filter_for_direct_queries(): + """Direct ORM access cannot bypass tenant isolation by omitting WHERE.""" + engine = create_engine("sqlite://") + TenantScopedRecord.__table__.create(engine) + tenant_a = str(uuid.uuid4()) + tenant_b = str(uuid.uuid4()) + + with Session(engine) as session: + session.add_all( + [ + TenantScopedRecord(id="a", tenant_id=tenant_a), + TenantScopedRecord(id="b", tenant_id=tenant_b), + ] + ) + session.commit() + + with tenant_context(tenant_a): + records = session.scalars(select(TenantScopedRecord).order_by(TenantScopedRecord.id)).all() + + assert [record.id for record in records] == ["a"]