diff --git a/backend/file_management/urls.py b/backend/file_management/urls.py index 8b0ae2dcf9..995ff3f4d7 100644 --- a/backend/file_management/urls.py +++ b/backend/file_management/urls.py @@ -34,16 +34,10 @@ "get": "list_ide", } ) -file_delete = FileManagementViewSet.as_view( - { - "get": "delete", - } -) urlpatterns = format_suffix_patterns( [ path("file", file_list, name="file-list"), path("file/download", file_downlaod, name="download"), path("file/upload", file_upload, name="upload"), - path("file/delete", file_delete, name="delete"), ] ) diff --git a/backend/file_management/views.py b/backend/file_management/views.py index b2875d4251..2a4d684763 100644 --- a/backend/file_management/views.py +++ b/backend/file_management/views.py @@ -4,12 +4,10 @@ from connector_v2.models import ConnectorInstance from django.http import HttpRequest from oauth2client.client import HttpAccessTokenRefreshError -from prompt_studio.prompt_studio_document_manager_v2.models import DocumentManager -from rest_framework import serializers, status, viewsets +from rest_framework import serializers, viewsets from rest_framework.decorators import action from rest_framework.response import Response from rest_framework.versioning import URLPathVersioning -from utils.user_session import UserSessionUtils from file_management.exceptions import ( ConnectorInstanceNotFound, @@ -18,13 +16,11 @@ ) from file_management.file_management_helper import FileManagerHelper from file_management.serializer import ( - FileInfoIdeSerializer, FileInfoSerializer, FileListRequestSerializer, FileUploadSerializer, ) from unstract.connectors.exceptions import ConnectorError -from unstract.connectors.filesystems.local_storage.local_storage import LocalStorageFS logger = logging.getLogger(__name__) @@ -32,7 +28,7 @@ class FileManagementViewSet(viewsets.ModelViewSet): """FileManagement view. - Handles GET,POST,PUT,PATCH and DELETE + Handles GET, POST, PUT and PATCH """ versioning_class = URLPathVersioning @@ -99,36 +95,3 @@ def upload(self, request: HttpRequest) -> Response: logger.info(f"Uploading file: {file_name}" if file_name else "Uploading file") FileManagerHelper.upload_file(file_system, path, uploaded_file, file_name) return Response({"message": "Files are uploaded successfully!"}) - - @action(detail=True, methods=["get"]) - def delete(self, request: HttpRequest) -> Response: - serializer = FileInfoIdeSerializer(data=request.GET) - serializer.is_valid(raise_exception=True) - document_id: str = serializer.validated_data.get("document_id") - document: DocumentManager = DocumentManager.objects.get(pk=document_id) - file_name: str = document.document_name - tool_id: str = serializer.validated_data.get("tool_id") - file_path = FileManagerHelper.handle_sub_directory_for_tenants( - UserSessionUtils.get_organization_id(request), - is_create=False, - user_id=request.user.user_id, - tool_id=tool_id, - ) - path = file_path - file_system = LocalStorageFS(settings={"path": path}) - try: - # Delete the document record - document.delete() - - # Delete the file - FileManagerHelper.delete_file(file_system, path, file_name) - return Response( - {"data": "File deleted succesfully."}, - status=status.HTTP_200_OK, - ) - except Exception as exc: - logger.error(f"Exception thrown from file deletion, error {exc}") - return Response( - {"data": "File deletion failed."}, - status=status.HTTP_400_BAD_REQUEST, - ) diff --git a/backend/notification_v2/internal_views.py b/backend/notification_v2/internal_views.py index e352f56db7..dfd97ad149 100644 --- a/backend/notification_v2/internal_views.py +++ b/backend/notification_v2/internal_views.py @@ -42,7 +42,9 @@ class WebhookInternalViewSet(viewsets.ReadOnlyModelViewSet): serializer_class = NotificationSerializer lookup_field = "id" - # Backward compat: remove once all workers pass X-Organization-ID. + # OrganizationFilterBackend is off here; get_queryset() scopes instead, via + # filter_queryset_by_organization. That helper fails closed, so a caller + # without X-Organization-ID gets zero rows. skip_org_filter = True def get_queryset(self): diff --git a/backend/pipeline_v2/internal_api_views.py b/backend/pipeline_v2/internal_api_views.py index 5c1471d717..5681667af4 100644 --- a/backend/pipeline_v2/internal_api_views.py +++ b/backend/pipeline_v2/internal_api_views.py @@ -13,7 +13,9 @@ class PipelineInternalViewSet(ViewSet): - # Backward compat: remove once all workers pass X-Organization-ID. + # OrganizationFilterBackend is off here; scoping runs through + # filter_queryset_by_organization, which fails closed, so a caller without + # X-Organization-ID gets zero rows. skip_org_filter = True def retrieve(self, request, pk=None): diff --git a/backend/prompt_studio/prompt_profile_manager_v2/models.py b/backend/prompt_studio/prompt_profile_manager_v2/models.py index 10a234f462..fbd7437656 100644 --- a/backend/prompt_studio/prompt_profile_manager_v2/models.py +++ b/backend/prompt_studio/prompt_profile_manager_v2/models.py @@ -5,14 +5,15 @@ from django.db import models from django.db.models import Q from tenant_account_v2.organization_member_service import OrganizationMemberService -from utils.models.base_model import BaseModel, BaseModelManager +from utils.models.base_model import BaseModel +from utils.models.org_aware_manager import OrgAwareManager from utils.user_context import UserContext from prompt_studio.prompt_studio_core_v2.exceptions import DefaultProfileError from prompt_studio.prompt_studio_core_v2.models import CustomTool -class ProfileManagerModelManager(BaseModelManager): +class ProfileManagerModelManager(OrgAwareManager): def for_user(self, user): """Read visibility: profile's own share fields OR parent CustomTool sharing. diff --git a/backend/prompt_studio/prompt_studio_core_v2/internal_views.py b/backend/prompt_studio/prompt_studio_core_v2/internal_views.py index 3ad3a5db16..b4246e0cf7 100644 --- a/backend/prompt_studio/prompt_studio_core_v2/internal_views.py +++ b/backend/prompt_studio/prompt_studio_core_v2/internal_views.py @@ -211,7 +211,26 @@ def extraction_status(request): extracted=extracted, error_message=error_message, ) - return JsonResponse({"success": success}) + if not success: + # A 200 here is indistinguishable from a write that landed: the + # worker only wraps this call in try/except and never reads the + # body, so the status would be silently dropped and every later + # Answer Prompt would re-run the full extraction. Non-2xx makes + # the worker's existing handler log it. + logger.error( + "extraction_status not recorded for document %s profile %s", + document_id, + profile_manager_id, + ) + return JsonResponse( + { + "success": False, + "error": "Extraction status could not be recorded", + }, + status=status.HTTP_500_INTERNAL_SERVER_ERROR, + ) + + return JsonResponse({"success": True}) except Exception as e: logger.exception("extraction_status internal API failed") diff --git a/backend/prompt_studio/prompt_studio_core_v2/migration_utils.py b/backend/prompt_studio/prompt_studio_core_v2/migration_utils.py index b9236d04f0..ea838a06bf 100644 --- a/backend/prompt_studio/prompt_studio_core_v2/migration_utils.py +++ b/backend/prompt_studio/prompt_studio_core_v2/migration_utils.py @@ -60,13 +60,35 @@ def migrate_tool_to_adapter_based( # Re-fetch the summarize profile with lock within transaction try: - summarize_profile = ProfileManager.objects.select_for_update().get( - prompt_studio_tool=tool_instance, is_summarize_llm=True - ) + # of=("self",): the org-scoped manager joins through + # AdapterInstance, which would otherwise be locked too. + summarize_profile = ProfileManager.objects.select_for_update( + of=("self",) + ).get(prompt_studio_tool=tool_instance, is_summarize_llm=True) except ObjectDoesNotExist: - logger.info( - f"No summarize profile found for tool {tool_instance.tool_id}, skipping migration" - ) + # Two different situations reach here now that + # ProfileManager.objects is scoped through + # vector_store__organization: the profile genuinely does + # not exist, or it exists and the org filter hid it. The + # second is a misconfiguration that never self-heals — this + # lazy migration re-runs and re-skips on every invocation — + # so it must not share an INFO line with the first. + exists_unscoped = ProfileManager._base_manager.filter( + prompt_studio_tool=tool_instance, is_summarize_llm=True + ).exists() + if exists_unscoped: + logger.error( + "Summarize profile for tool %s exists but is not " + "visible in the current organization context; " + "migration skipped and will keep being skipped.", + tool_instance.tool_id, + ) + else: + logger.info( + "No summarize profile found for tool %s, skipping " + "migration", + tool_instance.tool_id, + ) return False # Check if profile has an LLM adapter diff --git a/backend/prompt_studio/prompt_studio_core_v2/views.py b/backend/prompt_studio/prompt_studio_core_v2/views.py index 5741c7ccd9..f89a83658b 100644 --- a/backend/prompt_studio/prompt_studio_core_v2/views.py +++ b/backend/prompt_studio/prompt_studio_core_v2/views.py @@ -10,9 +10,10 @@ from account_v2.custom_exceptions import DuplicateData from celery import signature from celery.result import AsyncResult -from django.db import IntegrityError +from django.db import IntegrityError, transaction from django.db.models import Count, OuterRef, QuerySet, Subquery from django.http import HttpRequest, HttpResponse +from django.shortcuts import get_object_or_404 from django.utils import timezone from file_management.constants import FileInformationKey as FileKey from file_management.exceptions import FileNotFound @@ -24,6 +25,7 @@ from plugins import get_plugin from rest_framework import status, viewsets from rest_framework.decorators import action +from rest_framework.exceptions import ValidationError from rest_framework.request import Request from rest_framework.response import Response from rest_framework.versioning import URLPathVersioning @@ -389,13 +391,39 @@ def make_profile_default(self, request: HttpRequest, pk: Any = None) -> Response self.get_object() ) # Assuming you have a get_object method in your viewset - ProfileManager.objects.filter(prompt_studio_tool=prompt_tool).update( - is_default=False + # Validate before looking anything up. A missing key raised KeyError + # and a non-UUID value raised Django's ValidationError; neither is + # mapped by drf_standardized_errors, so both surfaced as 500s next to + # the 404 this action already returns for a valid-but-unmatched id. + default_profile = request.data.get("default_profile") + if not default_profile: + raise ValidationError(detail="'default_profile' is required.") + try: + default_profile = uuid.UUID(str(default_profile)) + except (ValueError, AttributeError, TypeError): + raise ValidationError(detail="'default_profile' must be a valid UUID.") + + # Resolve the target before clearing anything: the id comes straight + # from the request body, and clearing first would leave the tool with no + # default at all when it does not match. Scoped to the same tool the + # caller already passed authz on, so another tool's id is a 404. + profile_manager = get_object_or_404( + ProfileManager, + pk=default_profile, + prompt_studio_tool=prompt_tool, ) - profile_manager = ProfileManager.objects.get(pk=request.data["default_profile"]) - profile_manager.is_default = True - profile_manager.save() + # Both writes in one transaction so a failure between them cannot leave + # the tool with zero defaults or two. update_fields so the second write + # touches one column: profile_manager was read before the transaction + # opened, and a bare save() would write every column from that snapshot + # back over any concurrent edit. + with transaction.atomic(): + ProfileManager.objects.filter(prompt_studio_tool=prompt_tool).update( + is_default=False + ) + profile_manager.is_default = True + profile_manager.save(update_fields=["is_default"]) return Response( status=status.HTTP_200_OK, @@ -1130,11 +1158,30 @@ def delete_for_ide(self, request: HttpRequest, pk: uuid) -> Response: document_id: str = serializer.validated_data.get(ToolStudioPromptKeys.DOCUMENT_ID) org_id = UserSessionUtils.get_organization_id(request) user_id = custom_tool.created_by.user_id - document: DocumentManager = DocumentManager.objects.get(pk=document_id) + # Scope to the tool the caller already passed authz on — tighter than + # org scope. self.get_object() above is filtered by the backend, but + # this lookup is a raw .objects query and would not be. + # get_object_or_404 keeps a non-matching id a 404 rather than an + # unhandled DoesNotExist, which the DRF handler turns into a 500. + document: DocumentManager = get_object_or_404( + DocumentManager, pk=document_id, tool=custom_tool + ) try: # Delete indexed flags in redis index_managers = IndexManager.objects.filter(document_manager=document_id) + if not index_managers.exists(): + # Empty means either "never indexed" or "the org filter hid the + # rows". In the second case the Redis indexing flags outlive the + # document, and a re-upload of the same file is treated as + # already indexed — with a 200 telling the user it all worked. + logger.warning( + "No index managers visible for document %s (tool %s, org %s); " + "deleting without clearing Redis indexing flags.", + document_id, + custom_tool.tool_id, + org_id, + ) for index_manager in index_managers: raw_index_id = index_manager.raw_index_id DocumentIndexingService.remove_document_indexing( @@ -1155,7 +1202,20 @@ def delete_for_ide(self, request: HttpRequest, pk: uuid) -> Response: status=status.HTTP_200_OK, ) except Exception as exc: - logger.error("Exception thrown from file deletion, error: %s", exc) + # Deliberately broad. Three subsystems are in play — Redis via + # DocumentIndexingService, the object store via + # PromptStudioFileHelper, and the database — and their failures do + # not share a base class, so narrowing to any list turns a + # reachable outage in whichever one was missed into a 500. The + # diagnosability problem was the log line, not the catch: it now + # carries the exception type, the document and a stack. + logger.error( + "File deletion failed for document %s (tool %s): %s", + document_id, + custom_tool.tool_id, + exc, + exc_info=True, + ) return Response( {"data": "File deletion failed."}, status=status.HTTP_400_BAD_REQUEST, diff --git a/backend/prompt_studio/prompt_studio_document_manager_v2/models.py b/backend/prompt_studio/prompt_studio_document_manager_v2/models.py index 15c76c5087..07ee681c4c 100644 --- a/backend/prompt_studio/prompt_studio_document_manager_v2/models.py +++ b/backend/prompt_studio/prompt_studio_document_manager_v2/models.py @@ -3,6 +3,7 @@ from account_v2.models import User from django.db import models from utils.models.base_model import BaseModel +from utils.models.org_aware_manager import OrgAwareManager from prompt_studio.prompt_studio_core_v2.models import CustomTool @@ -10,6 +11,13 @@ class DocumentManager(BaseModel): """Model to store the document details.""" + # Org scoping lives at the manager because OrganizationFilterBackend only + # scopes querysets routed through filter_queryset(). A raw Model.objects + # lookup inside a view bypasses it — including inside a custom @action, + # whose own self.get_object() *is* filtered but whose hand-written queries + # are not. + objects = OrgAwareManager() + document_id = models.UUIDField(primary_key=True, default=uuid.uuid4, editable=False) document_name = models.CharField( diff --git a/backend/prompt_studio/prompt_studio_index_manager_v2/models.py b/backend/prompt_studio/prompt_studio_index_manager_v2/models.py index 60ee406304..6d84ed832b 100644 --- a/backend/prompt_studio/prompt_studio_index_manager_v2/models.py +++ b/backend/prompt_studio/prompt_studio_index_manager_v2/models.py @@ -7,6 +7,7 @@ from django.db.models.signals import pre_delete from django.dispatch import receiver from utils.models.base_model import BaseModel +from utils.models.org_aware_manager import OrgAwareManager from utils.user_context import UserContext from prompt_studio.prompt_profile_manager_v2.models import ProfileManager @@ -21,6 +22,9 @@ class IndexManager(BaseModel): """Model to store the index details.""" + # See DocumentManager.objects for why scoping lives at the manager. + objects = OrgAwareManager() + index_manager_id = models.UUIDField( primary_key=True, default=uuid.uuid4, editable=False ) diff --git a/backend/prompt_studio/prompt_studio_index_manager_v2/prompt_studio_index_helper.py b/backend/prompt_studio/prompt_studio_index_manager_v2/prompt_studio_index_helper.py index d17c157865..1c10bf5b1a 100644 --- a/backend/prompt_studio/prompt_studio_index_manager_v2/prompt_studio_index_helper.py +++ b/backend/prompt_studio/prompt_studio_index_manager_v2/prompt_studio_index_helper.py @@ -1,7 +1,8 @@ import json import logging -from django.db import transaction +from django.core.exceptions import ImproperlyConfigured +from django.db import DatabaseError, transaction from prompt_studio.prompt_profile_manager_v2.models import ProfileManager from prompt_studio.prompt_studio_core_v2.exceptions import IndexingAPIError @@ -109,12 +110,21 @@ def mark_extraction_status( # Lock the row (or create an empty one) so concurrent callers # merge into the same dict rather than clobbering each other. - index_manager, created = ( - IndexManager.objects.select_for_update().get_or_create( - document_manager=document, - profile_manager=profile_manager, - defaults={"extraction_status": {}}, - ) + # of=("self",) keeps the lock on index_manager rows only. + # + # _base_manager, not objects: Django applies a manager's filter + # to the get half of get_or_create but not to the create half. + # Through the org-scoped manager, a row the filter hides makes + # get miss and create insert, which violates + # unique_document_manager_profile_manager_index. The document + # above was already fetched org-scoped, so the scope is checked + # either way and this only removes the failure mode. + index_manager, created = IndexManager._base_manager.select_for_update( + of=("self",) + ).get_or_create( + document_manager=document, + profile_manager=profile_manager, + defaults={"extraction_status": {}}, ) # Merge in place — update_or_create(defaults=...) would replace @@ -150,12 +160,24 @@ def mark_extraction_status( return True except DocumentManager.DoesNotExist: - logger.error(f"Document with ID {document_id} does not exist.") + # Now reachable two ways: the row is genuinely gone, or the + # org-scoped manager hid it from this caller. Both mean the status + # was not written, which is what the caller has to act on. + logger.error( + "Document %s not found or not visible in the current " + "organization; extraction status not recorded.", + document_id, + ) return False - except Exception as e: + except (DatabaseError, TypeError, ImproperlyConfigured): + # DatabaseError covers IntegrityError/OperationalError, TypeError a + # malformed extraction_status payload, ImproperlyConfigured a bad + # org path pin. Narrowed from a bare `except Exception` so an + # unexpected type propagates instead of being reported as "no such + # document". logger.exception( - f"Unexpected error marking extraction status for document {document_id}: {e}" + "Failed to mark extraction status for document %s", document_id ) return False diff --git a/backend/prompt_studio/prompt_studio_output_manager_v2/models.py b/backend/prompt_studio/prompt_studio_output_manager_v2/models.py index 7b8616968f..9420b87e9d 100644 --- a/backend/prompt_studio/prompt_studio_output_manager_v2/models.py +++ b/backend/prompt_studio/prompt_studio_output_manager_v2/models.py @@ -3,6 +3,7 @@ from account_v2.models import User from django.db import models from utils.models.base_model import BaseModel +from utils.models.org_aware_manager import OrgAwareManager from prompt_studio.prompt_profile_manager_v2.models import ProfileManager from prompt_studio.prompt_studio_core_v2.models import CustomTool @@ -16,6 +17,9 @@ class PromptStudioOutputManager(BaseModel): By default the tools will be added to private tool hub. """ + # See DocumentManager.objects for why scoping lives at the manager. + objects = OrgAwareManager() + prompt_output_id = models.UUIDField( primary_key=True, default=uuid.uuid4, editable=False ) diff --git a/backend/prompt_studio/prompt_studio_output_manager_v2/output_manager_helper.py b/backend/prompt_studio/prompt_studio_output_manager_v2/output_manager_helper.py index 699cacb749..54450201c7 100644 --- a/backend/prompt_studio/prompt_studio_output_manager_v2/output_manager_helper.py +++ b/backend/prompt_studio/prompt_studio_output_manager_v2/output_manager_helper.py @@ -76,21 +76,28 @@ def update_or_create_prompt_output( the instance. """ try: - prompt_output, success = PromptStudioOutputManager.objects.get_or_create( - document_manager=document_manager, - tool_id=tool, - profile_manager=profile_manager, - prompt_id=prompt, - is_single_pass_extract=is_single_pass_extract, - defaults={ - "output": output, - "eval_metrics": eval_metrics, - "context": context, - "challenge_data": challenge_data, - "highlight_data": highlight_data, - "confidence_data": confidence_data, - "word_confidence_data": word_confidence_data, - }, + # _base_manager, not objects: the manager filter applies to the + # get half of get_or_create but not the create half, so a row + # the org scope hides makes get miss and create collide with + # unique_prompt_output_index. `tool` and `document_manager` are + # already org-verified by the caller. + prompt_output, success = ( + PromptStudioOutputManager._base_manager.get_or_create( + document_manager=document_manager, + tool_id=tool, + profile_manager=profile_manager, + prompt_id=prompt, + is_single_pass_extract=is_single_pass_extract, + defaults={ + "output": output, + "eval_metrics": eval_metrics, + "context": context, + "challenge_data": challenge_data, + "highlight_data": highlight_data, + "confidence_data": confidence_data, + "word_confidence_data": word_confidence_data, + }, + ) ) if success: diff --git a/backend/prompt_studio/prompt_studio_output_manager_v2/views.py b/backend/prompt_studio/prompt_studio_output_manager_v2/views.py index 44111dc744..bc41b34594 100644 --- a/backend/prompt_studio/prompt_studio_output_manager_v2/views.py +++ b/backend/prompt_studio/prompt_studio_output_manager_v2/views.py @@ -1,11 +1,11 @@ import logging +import uuid from typing import Any -from django.core.exceptions import ObjectDoesNotExist from django.db.models import QuerySet from django.http import HttpRequest from rest_framework import status, viewsets -from rest_framework.exceptions import ValidationError +from rest_framework.exceptions import APIException, ValidationError from rest_framework.response import Response from rest_framework.versioning import URLPathVersioning from utils.common_utils import CommonUtils @@ -29,6 +29,42 @@ logger = logging.getLogger(__name__) +def _validated_tool_id(raw: Any) -> uuid.UUID: + """A query-string ``tool_id`` as a UUID, or a 400. + + ``CustomTool.tool_id`` is a UUID primary key, so a non-UUID value makes + ``filter()`` raise Django's ``ValidationError`` while the query is being + *built*. drf_standardized_errors maps only ``Http404`` and Django's + ``PermissionDenied``, so that surfaced as a 500 rather than a 400. + """ + try: + return uuid.UUID(str(raw)) + except (ValueError, AttributeError, TypeError): + raise ValidationError(detail="'tool_id' must be a valid UUID.") + + +def _required_organization(tool_id: Any) -> Any: + """The request's organization, refusing to proceed without one. + + ``UserContext.get_organization()`` returns None on both + ``Organization.DoesNotExist`` and ``ProgrammingError``, neither logged. A + None here compiles to ``organization_id IS NULL``, which matches nothing + whatever the tool id — downstream every output renders as ``""`` and the + user sees a blank project that has real persisted outputs, with nothing to + correlate in logs. These endpoints are only routed under + ``/api/v1/unstract//``, so a null org is a bug, not a state to serve. + """ + organization = UserContext.get_organization() + if organization is None: + logger.error( + "No organization in context while reading prompt-studio outputs " + "(tool %s); refusing to serve an unscoped empty result.", + tool_id, + ) + raise APIException(detail="Organization context is unavailable.") + return organization + + class PromptStudioOutputView(viewsets.ModelViewSet): versioning_class = URLPathVersioning serializer_class = PromptStudioOutputSerializer @@ -78,9 +114,11 @@ def latest_outputs_by_keys(self, request: HttpRequest) -> Response: if not prompt_keys: return Response({}, status=status.HTTP_200_OK) - # Custom actions skip filter_queryset(), so OrganizationFilterBackend - # never runs — scope explicitly to prevent cross-tenant reads. - organization = UserContext.get_organization() + tool_id = _validated_tool_id(tool_id) + + # A raw .objects query is not routed through filter_queryset(), so + # OrganizationFilterBackend does not see it — scope explicitly. + organization = _required_organization(tool_id) prompt_id_to_key = dict( ToolStudioPrompt.objects.filter( tool_id=tool_id, @@ -119,17 +157,25 @@ def get_output_for_tool_default(self, request: HttpRequest) -> Response: tool_id = request.GET.get("tool_id") document_manager_id = request.GET.get("document_manager") tool_validation_message = PromptOutputManagerErrorMessage.TOOL_VALIDATION - tool_not_found = PromptOutputManagerErrorMessage.TOOL_NOT_FOUND if not tool_id: raise ValidationError(detail=tool_validation_message) - try: - # Fetch ToolStudioPrompt records based on tool_id - tool_studio_prompts = ToolStudioPrompt.objects.filter( - tool_id=tool_id - ).order_by("sequence_number") - except ObjectDoesNotExist: - raise ValidationError(detail=tool_not_found) + tool_id = _validated_tool_id(tool_id) + organization = _required_organization(tool_id) + + # Fetch ToolStudioPrompt records based on tool_id. + # A raw .objects query is not routed through filter_queryset(), so + # OrganizationFilterBackend does not see it — scope explicitly. + # + # No exception handling below: for a valid UUID that matches no row, or + # a tool in another organization, filter() returns empty rather than + # raising. Empty is also the correct result for a tool that simply has + # no prompts yet, the normal state of a newly created project — so that + # case stays a 200 with an empty body. + tool_studio_prompts = ToolStudioPrompt.objects.filter( + tool_id=tool_id, + tool_id__organization=organization, + ).order_by("sequence_number") # Invoke helper method to frame and fetch default response. result: dict[str, Any] = OutputManagerHelper.fetch_default_output_response( diff --git a/backend/prompt_studio/prompt_studio_v2/models.py b/backend/prompt_studio/prompt_studio_v2/models.py index faaf7b0313..8f9f9bc316 100644 --- a/backend/prompt_studio/prompt_studio_v2/models.py +++ b/backend/prompt_studio/prompt_studio_v2/models.py @@ -4,6 +4,7 @@ from django.db import models from django.utils import timezone from utils.models.base_model import BaseModel +from utils.models.org_aware_manager import OrgAwareManager from prompt_studio.prompt_profile_manager_v2.models import ProfileManager from prompt_studio.prompt_studio_core_v2.models import CustomTool @@ -15,6 +16,10 @@ class ToolStudioPrompt(BaseModel): It has Many to one relation with CustomTool for ToolStudio. """ + # See DocumentManager.objects for why scoping lives at the manager. + # tool_id is nullable, so prompts orphaned from their tool are excluded. + objects = OrgAwareManager() + class EnforceType(models.TextChoices): TEXT = "text", "Response sent as Text" NUMBER = "number", "Response sent as number" diff --git a/backend/prompt_studio/tests/__init__.py b/backend/prompt_studio/tests/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/backend/prompt_studio/tests/test_cross_org_isolation.py b/backend/prompt_studio/tests/test_cross_org_isolation.py new file mode 100644 index 0000000000..f602e53ffd --- /dev/null +++ b/backend/prompt_studio/tests/test_cross_org_isolation.py @@ -0,0 +1,286 @@ +"""Organization isolation for the prompt-studio child models. + +``OrganizationFilterBackend`` only scopes querysets routed through +``filter_queryset()``. A raw ``.objects.get()/filter()`` written inside a view +bypasses it — including inside a custom DRF ``@action``, where +``self.get_object()`` *is* filtered but the hand-written queries beside it are +not. These tests pin the controls that cover that gap: org scoping on the +managers, plus explicit scoping where an id arrives directly from the request. + +Shape of each case: act as org A, pass an org B id, assert the call is +refused and org B's row is untouched. +""" + +import secrets + +import pytest +from account_v2.models import Organization, User +from adapter_processor_v2.models import AdapterInstance +from django.contrib.contenttypes.models import ContentType +from django.test import TestCase +from django.urls import NoReverseMatch, reverse +from permissions.roles import ResourceRole +from rest_framework.test import APIRequestFactory, force_authenticate +from tenant_account_v2.models import ResourceMembership +from utils.user_context import UserContext + +from prompt_studio.prompt_profile_manager_v2.models import ProfileManager +from prompt_studio.prompt_studio_core_v2.models import CustomTool +from prompt_studio.prompt_studio_core_v2.views import PromptStudioCoreView +from prompt_studio.prompt_studio_document_manager_v2.models import DocumentManager +from prompt_studio.prompt_studio_index_manager_v2.models import IndexManager +from prompt_studio.prompt_studio_output_manager_v2.models import ( + PromptStudioOutputManager, +) +from prompt_studio.prompt_studio_v2.models import ToolStudioPrompt + + +class OrgFixture: + """One organization with a fully populated prompt-studio object graph.""" + + def __init__(self, slug: str): + self.org = Organization.objects.create( + name=slug, display_name=slug, organization_id=slug + ) + UserContext.set_organization_identifier(slug) + + self.user = User.objects.create_user( + username=f"{slug}@example.com", + email=f"{slug}@example.com", + password=secrets.token_urlsafe(), + ) + self.tool = CustomTool.objects.create( + tool_name=f"tool-{slug}", + description="isolation test tool", + organization=self.org, + created_by=self.user, + ) + adapter = self._adapter(slug) + self.profile = ProfileManager.objects.create( + profile_name=f"profile-{slug}", + vector_store=adapter, + embedding_model=adapter, + llm=adapter, + x2text=adapter, + chunk_size=0, + chunk_overlap=0, + section="Default", + retrieval_strategy="simple", + similarity_top_k=3, + prompt_studio_tool=self.tool, + is_default=True, + created_by=self.user, + ) + self.document = DocumentManager.objects.create( + document_name=f"doc-{slug}.pdf", tool=self.tool, created_by=self.user + ) + self.index = IndexManager.objects.create( + document_manager=self.document, profile_manager=self.profile + ) + self.prompt = ToolStudioPrompt.objects.create( + prompt_key=f"key_{slug}", prompt="extract", tool_id=self.tool + ) + self.output = PromptStudioOutputManager.objects.create( + output="secret", + prompt_id=self.prompt, + document_manager=self.document, + profile_manager=self.profile, + tool_id=self.tool, + ) + + def _adapter(self, slug: str) -> AdapterInstance: + return AdapterInstance.objects.create( + adapter_name=f"adapter-{slug}", + adapter_id="openai|test", + adapter_type="LLM", + adapter_metadata={}, + organization=self.org, + created_by=self.user, + ) + + +@pytest.mark.django_db +class CrossOrgIsolationTest(TestCase): + """Org A must not reach org B's prompt-studio rows through any manager.""" + + def setUp(self) -> None: + self.a = OrgFixture(f"org-a-{secrets.token_hex(3)}") + self.b = OrgFixture(f"org-b-{secrets.token_hex(3)}") + # End state: acting as org A, as a request would. + UserContext.set_organization_identifier(self.a.org.organization_id) + + def tearDown(self) -> None: + # UserContext is thread-local, not DB state, so TestCase's transaction + # rollback does not clear it. Tests below deliberately switch org and + # would otherwise leak that into whatever runs next. + UserContext.set_organization_identifier(None) + + # --- manager scoping: the default-deny layer (A-1) -------------------- + + def test_document_of_other_org_is_not_gettable(self): + """A document id from another org must not resolve.""" + with self.assertRaises(DocumentManager.DoesNotExist): + DocumentManager.objects.get(pk=self.b.document.document_id) + + def test_prompt_of_other_org_is_not_listable(self): + """Prompts must not be listable by another org's tool id.""" + assert not ToolStudioPrompt.objects.filter(tool_id=self.b.tool).exists() + + def test_output_of_other_org_is_not_listable(self): + assert not PromptStudioOutputManager.objects.filter( + tool_id=self.b.tool + ).exists() + + def test_index_of_other_org_is_not_listable(self): + assert not IndexManager.objects.filter( + document_manager=self.b.document + ).exists() + + def test_profile_of_other_org_is_not_gettable(self): + """``make_profile_default`` takes this id straight from the body.""" + with self.assertRaises(ProfileManager.DoesNotExist): + ProfileManager.objects.get(pk=self.b.profile.profile_id) + + # --- same-org access must still work ---------------------------------- + + def test_own_org_rows_remain_visible(self): + assert DocumentManager.objects.get(pk=self.a.document.document_id) + assert ProfileManager.objects.get(pk=self.a.profile.profile_id) + assert ToolStudioPrompt.objects.filter(tool_id=self.a.tool).exists() + assert PromptStudioOutputManager.objects.filter(tool_id=self.a.tool).exists() + assert IndexManager.objects.filter(document_manager=self.a.document).exists() + + def test_no_org_context_is_unfiltered(self): + """Management commands and shell keep full access (fail-open).""" + UserContext.set_organization_identifier(None) + assert DocumentManager.objects.filter( + pk=self.b.document.document_id + ).exists() + + def test_worker_context_sees_its_own_org(self): + """B1 — workers do run with org context set, so the manager filters + there too. Indexing must still find its own org's rows.""" + UserContext.set_organization_identifier(self.b.org.organization_id) + assert IndexManager.objects.filter( + document_manager=self.b.document + ).exists() + assert DocumentManager.objects.get(pk=self.b.document.document_id) + + # --- explicit scoping at the reported call sites (A-3, A-5) ----------- + + def test_delete_for_ide_lookup_is_tool_scoped(self): + """A doc id from another tool in the *same* org is refused too.""" + sibling = CustomTool.objects.create( + tool_name="sibling", + description="second tool, same org", + organization=self.a.org, + created_by=self.a.user, + ) + with self.assertRaises(DocumentManager.DoesNotExist): + DocumentManager.objects.get( + pk=self.a.document.document_id, tool=sibling + ) + + def test_make_profile_default_lookup_is_tool_scoped(self): + """This lookup runs after ``get_object()`` has already passed authz on + the caller's own tool, so org scope alone does not constrain it.""" + with self.assertRaises(ProfileManager.DoesNotExist): + ProfileManager.objects.get( + pk=self.b.profile.profile_id, prompt_studio_tool=self.a.tool + ) + # Victim's default flag untouched. + UserContext.set_organization_identifier(self.b.org.organization_id) + assert ProfileManager.objects.get(pk=self.b.profile.profile_id).is_default + + # --- the ordering fix, driven through the view ------------------------ + + def _make_profile_default(self, tool, profile_id): + """PATCH make_profile_default as the owner of ``tool``. + + ``CustomTool.objects.for_user`` resolves visibility through + ResourceMembership, which the create *view* writes — the fixture builds + rows directly, so the OWNER row has to be added here or get_object() + 404s before the code under test runs. + """ + ResourceMembership.objects.get_or_create( + user=self.a.user, + role=ResourceRole.OWNER, + content_type=ContentType.objects.get_for_model(CustomTool), + object_id=str(tool.tool_id), + ) + view = PromptStudioCoreView.as_view({"patch": "make_profile_default"}) + request = APIRequestFactory().patch( + f"/prompt-studio/{tool.tool_id}/make_profile_default", + {"default_profile": str(profile_id)}, + format="json", + ) + force_authenticate(request, user=self.a.user) + return view(request, pk=str(tool.tool_id)) + + def _second_profile_on_tool_a(self): + return ProfileManager.objects.create( + profile_name="profile-a-second", + vector_store=self.a.profile.vector_store, + embedding_model=self.a.profile.embedding_model, + llm=self.a.profile.llm, + x2text=self.a.profile.x2text, + chunk_size=0, + chunk_overlap=0, + section="Default", + retrieval_strategy="simple", + similarity_top_k=3, + prompt_studio_tool=self.a.tool, + is_default=False, + created_by=self.a.user, + ) + + def test_make_profile_default_switches_the_default(self): + """The allow path: the old default is cleared and the new one set.""" + second = self._second_profile_on_tool_a() + + response = self._make_profile_default(self.a.tool, second.profile_id) + + assert response.status_code == 200, response.data + self.a.profile.refresh_from_db() + second.refresh_from_db() + assert second.is_default + assert not self.a.profile.is_default + + def test_rejected_default_leaves_the_existing_default_intact(self): + """A non-matching id must not clear the tool's current default. + + Driven through the view on purpose: the de-dup update runs against + every profile on the tool, and an ORM-only test never executes it, so + it cannot observe this property at all. + + What actually guards the invariant is resolving and clearing under one + of two conditions — resolve first, or clear first but inside the + transaction, where the 404 rolls the clear back. Mutation-tested: this + fails (0 defaults left) only on clear-first *without* the transaction, + which is what the code did before. Reverting just the ordering, with + ``transaction.atomic()`` still in place, is genuinely safe and does not + fail here. + """ + self._second_profile_on_tool_a() + assert ProfileManager.objects.get(pk=self.a.profile.profile_id).is_default + + response = self._make_profile_default(self.a.tool, self.b.profile.profile_id) + + assert response.status_code == 404, response.data + assert ( + ProfileManager.objects.filter( + prompt_studio_tool=self.a.tool, is_default=True + ).count() + == 1 + ), "the tool lost (or duplicated) its default while rejecting another org's id" + self.a.profile.refresh_from_db() + assert self.a.profile.is_default + + # --- A-4: the dead, state-changing-over-GET route is gone ------------- + + def test_file_delete_route_removed(self): + """Removed rather than fixed: no caller, and it deleted over GET.""" + # Sibling route still resolves, so a naming change can't fake a pass. + assert reverse("tenant:upload").endswith("/file/upload") + with pytest.raises(NoReverseMatch): + reverse("tenant:delete") diff --git a/backend/tool_instance_v2/internal_views.py b/backend/tool_instance_v2/internal_views.py index 25165a9739..8bfd773986 100644 --- a/backend/tool_instance_v2/internal_views.py +++ b/backend/tool_instance_v2/internal_views.py @@ -23,7 +23,9 @@ class ToolExecutionInternalViewSet(viewsets.ModelViewSet): """Internal API for tool execution operations used by lightweight workers.""" serializer_class = ToolInstanceSerializer - # Backward compat: remove once all workers pass X-Organization-ID. + # OrganizationFilterBackend is off here; scoping runs through + # filter_queryset_by_organization, which fails closed, so a caller without + # X-Organization-ID gets zero rows. skip_org_filter = True def get_queryset(self): diff --git a/backend/utils/filters/organization_filter.py b/backend/utils/filters/organization_filter.py index a0eba72daf..072e9d03bd 100644 --- a/backend/utils/filters/organization_filter.py +++ b/backend/utils/filters/organization_filter.py @@ -40,6 +40,13 @@ class NotificationViewSet(viewsets.ModelViewSet): "pipeline__workflow__organization", "api__workflow__organization", ] + + Precedence: org_filter_paths wins over the model's pin in + ORG_PATH_OVERRIDES, and is checked before get_org_path is ever called. So + a viewset that sets both scopes through the paths here, not the pin, while + OrgAwareManager on the same model still uses the pin. Prefer the pin — + it applies at both layers. Reach for org_filter_paths only when the model + genuinely needs OR across several nullable paths. """ def filter_queryset(self, request, queryset, view): diff --git a/backend/utils/models/org_aware_manager.py b/backend/utils/models/org_aware_manager.py index ab40303dbb..33cdcd94d4 100644 --- a/backend/utils/models/org_aware_manager.py +++ b/backend/utils/models/org_aware_manager.py @@ -58,20 +58,35 @@ def get_queryset(self): try: org = UserContext.get_organization() - except (RuntimeError, OperationalError, ProgrammingError): + except (RuntimeError, OperationalError, ProgrammingError) as exc: # OperationalError: DB not reachable (startup, migrations) # ProgrammingError: schema not ready (during migrations) # RuntimeError: pytest-django blocks DB access outside - # @pytest.mark.django_db. Note: this is a broad catch — any - # RuntimeError (e.g. from StateStore/middleware) returns an - # unfiltered queryset (fail-open). This is acceptable because - # OrgAwareManager is defense-in-depth; OrganizationFilterBackend - # at the view layer is the primary security boundary and - # fails-closed independently. + # @pytest.mark.django_db. + # + # Deliberately fail open: these are all "the request context does + # not exist yet", not "this caller may not see these rows". + # OrganizationFilterBackend is the primary boundary and fails + # closed independently at the view layer. + # + # The RuntimeError arm is broader than its stated cause — + # StateStore.get raises it for any unrecognised CONCURRENCY_MODE + # — so log it. This path is rare (startup, migrations, tests), so + # a line here is signal rather than noise, and it is the only way + # an unexpected fail-open becomes visible. + logger.warning( + "OrgAwareManager: no organization context for %s (%s: %s); " + "returning an unfiltered queryset.", + self.model._meta.label, + type(exc).__name__, + exc, + ) return qs if org is None: - # No request context (Celery, management commands, shell) + # No request context: Celery, management commands, shell. Not + # logged — this is the normal state for every query those make, + # and a line per queryset would drown the case above. return qs path = get_org_path(self.model) diff --git a/backend/utils/models/org_path_discovery.py b/backend/utils/models/org_path_discovery.py index 9a8011c716..79969b323f 100644 --- a/backend/utils/models/org_path_discovery.py +++ b/backend/utils/models/org_path_discovery.py @@ -19,6 +19,47 @@ _FK_TYPES = (models.ForeignKey, models.OneToOneField) +# Org paths pinned explicitly, checked before BFS. Keyed by model label +# ("app_label.ModelName") so this module stays import-free of the models. +# +# BFS returns the *shortest* path and breaks ties by field declaration order. +# Reordering two fields can therefore swap in a different path of the same +# length, and if that path runs through a nullable FK the resulting INNER JOIN +# silently drops every row with a NULL — data loss that reads as "missing +# records", not as an error. Pinning freezes the path against that. +# +# What test_org_path_discovery actually asserts: each pin still matches what +# BFS would pick, and every hop on it is non-nullable *unless* listed in that +# module's KNOWN_NULLABLE_HOPS with a reason. Several pins are on that list — +# including every terminal `organization` FK, which DefaultOrganizationMixin +# declares null=True — so "pinned" does not mean "cannot drop rows", it means +# "the rows it drops are known and written down". +# +# Precedence, for the two consumers: +# - OrgAwareManager always uses the pin. +# - OrganizationFilterBackend checks a viewset's `org_filter_paths` FIRST and +# only falls back to the pin. A viewset that sets it therefore scopes that +# model through a different join than its pin. Prefer the pin; reach for +# `org_filter_paths` only when a model needs OR across several nullable +# paths, which is why notification_v2 has it. +ORG_PATH_OVERRIDES: dict[str, str] = { + "prompt_studio_document_manager_v2.DocumentManager": "tool__organization", + "prompt_studio_index_manager_v2.IndexManager": ( + "document_manager__tool__organization" + ), + "prompt_studio_output_manager_v2.PromptStudioOutputManager": ( + "tool_id__organization" + ), + # ToolStudioPrompt.tool_id is nullable — prompts orphaned from their tool + # are excluded. This is the path already in force, pinned as-is rather + # than changed under a security fix. + "prompt_studio_v2.ToolStudioPrompt": "tool_id__organization", + # Deliberately not prompt_studio_tool__organization: that FK is nullable, + # so it would drop tool-less profiles. vector_store is non-null and + # AdapterInstance is org-owned, so it scopes to the same organization. + "prompt_profile_manager_v2.ProfileManager": "vector_store__organization", +} + def get_org_path(model: type) -> str | None: """Get the cached FK path from a model to Organization. @@ -26,6 +67,10 @@ def get_org_path(model: type) -> str | None: Returns the ORM lookup path (e.g., "wf_execution__workflow__organization") or None if no path exists. """ + pinned = ORG_PATH_OVERRIDES.get(model._meta.label) + if pinned: + return pinned + if model in _org_path_cache: return _org_path_cache[model] diff --git a/backend/utils/organization_utils.py b/backend/utils/organization_utils.py index 15053684bf..1d7add8c0c 100644 --- a/backend/utils/organization_utils.py +++ b/backend/utils/organization_utils.py @@ -72,7 +72,30 @@ def get_organization_context(organization: Organization) -> dict[str, Any]: def filter_queryset_by_organization(queryset, request, organization_field="organization"): - """Filter a Django queryset by organization context from request. + """Filter a Django queryset by the request's organization context. + + Fails closed. For every caller, this function is the only tenant boundary + in the request: the viewsets that reach it set ``skip_org_filter = True``, + which disables OrganizationFilterBackend, and the function-based + ``@api_view`` handlers that reach it have no filter backend at all. + Returning the queryset unfiltered when there is no organization context + would hand back every organization's rows. + + Scope note: this policy is local to this helper. ``OrgAwareManager`` + deliberately fails *open* when there is no request context, so Celery + tasks, management commands and the shell keep full access to the models it + scopes — see the comment on its ``get_queryset``. The two are not in + conflict; they guard different callers. + + The absent-header case is not exotic: ``InternalAPIAuthMiddleware`` logs a + warning and continues when ``X-Organization-ID`` is missing, so any caller + holding the internal service key reaches here without context simply by + omitting it. + + Note for callers that genuinely span organizations — the leader-elected + reaper is one — query the model directly rather than routing through here. + ``recover_stuck_pg_executions`` already does, which is why failing closed + does not affect it. Args: queryset: Django QuerySet to filter @@ -80,16 +103,22 @@ def filter_queryset_by_organization(queryset, request, organization_field="organ organization_field: Field name for organization relationship (default: 'organization') Returns: - Filtered queryset or empty queryset if organization not found + The queryset filtered to the request's organization, or an empty + queryset when the organization is absent or unresolvable. """ org_id = getattr(request, "organization_id", None) - if org_id: - organization = resolve_organization(org_id, raise_on_not_found=False) - if organization: - # Use dynamic field lookup - filter_kwargs = {organization_field: organization} - return queryset.filter(**filter_kwargs) - else: - # Return empty queryset if organization not found - return queryset.none() - return queryset + if not org_id: + logger.warning( + "Organization scoping requested without organization context on %s; " + "returning no rows. A caller that must span organizations should " + "query the model directly instead of using this helper.", + getattr(request, "path", ""), + ) + return queryset.none() + + organization = resolve_organization(org_id, raise_on_not_found=False) + if not organization: + logger.warning("Organization %s not found; returning no rows.", org_id) + return queryset.none() + + return queryset.filter(**{organization_field: organization}) diff --git a/backend/utils/tests/test_org_path_discovery.py b/backend/utils/tests/test_org_path_discovery.py new file mode 100644 index 0000000000..384de242ba --- /dev/null +++ b/backend/utils/tests/test_org_path_discovery.py @@ -0,0 +1,84 @@ +"""Guards for the pinned organization FK paths. + +These paths decide how every org-scoped queryset is filtered, at both the +manager layer (OrgAwareManager) and the view layer (OrganizationFilterBackend). +A path that changes silently is a cross-tenant leak or silent row loss, so both +properties are asserted here rather than left to review. + +No DB access — path discovery walks the model metadata only. +""" + +import pytest +from django.apps import apps +from utils.models.org_path_discovery import ( + ORG_PATH_OVERRIDES, + _discover_org_path, + get_org_path, +) + +PINS = sorted(ORG_PATH_OVERRIDES.items()) + +# Nullable hops accepted as pre-existing behaviour, not introduced here. +# Rows with a NULL value on these FKs are excluded from every org-scoped +# query. Anything not listed must be non-nullable. +# +# The terminal `organization` FK is on this list for every pin: +# DefaultOrganizationMixin declares it null=True, and save() backfills it from +# UserContext, which is None outside a request. So a CustomTool or +# AdapterInstance created by a management command, data migration, Celery task +# or shell persists with organization_id NULL. Those rows are already invisible +# to their own model's default manager, so the pins do not make them any less +# visible — but the hop is nullable and the assertion below must say so rather +# than skip it. +KNOWN_NULLABLE_HOPS = { + ("prompt_studio_v2.ToolStudioPrompt", "tool_id"), + ("prompt_studio_document_manager_v2.DocumentManager", "organization"), + ("prompt_studio_index_manager_v2.IndexManager", "organization"), + ("prompt_studio_output_manager_v2.PromptStudioOutputManager", "organization"), + ("prompt_studio_v2.ToolStudioPrompt", "organization"), + ("prompt_profile_manager_v2.ProfileManager", "organization"), +} + + +@pytest.mark.parametrize("label,expected", PINS) +def test_pin_is_returned(label, expected): + """get_org_path serves the pin, bypassing BFS.""" + assert get_org_path(apps.get_model(label)) == expected + + +@pytest.mark.parametrize("label,expected", PINS) +def test_pin_matches_discovery(label, expected): + """The pin still agrees with what BFS would pick. + + Fails when a field reorder or a new FK changes the shortest path. That is + the signal to re-derive the pin deliberately, not to update this constant + to make CI green. + """ + assert _discover_org_path(apps.get_model(label)) == expected + + +@pytest.mark.parametrize("label,expected", PINS) +def test_pin_traverses_only_non_nullable_fks(label, expected): + """Every hop on the pin must be non-nullable, or listed as a known exception. + + Django turns a positive filter over a nullable FK into an INNER JOIN, which + drops rows whose FK is NULL. On an org filter that is invisible data loss. + + The terminal `organization` hop is checked too, not skipped: it is the one + that is nullable on every pin, so excluding it would make this assertion + pass while proving nothing about the join that matters most. + """ + model = apps.get_model(label) + hops = expected.split("__") + + for hop in hops: + field = model._meta.get_field(hop) + assert not field.null or (label, hop) in KNOWN_NULLABLE_HOPS, ( + f"{model._meta.label}.{hop} is nullable: this pin drops every row " + f"with a NULL {hop}. Pick a non-nullable path or add it to " + f"KNOWN_NULLABLE_HOPS with a reason." + ) + model = field.related_model + + # The walk must have landed on Organization, not merely survived. + assert model._meta.label == "account_v2.Organization" diff --git a/backend/utils/tests/test_organization_scoping.py b/backend/utils/tests/test_organization_scoping.py new file mode 100644 index 0000000000..e277e9f4d0 --- /dev/null +++ b/backend/utils/tests/test_organization_scoping.py @@ -0,0 +1,81 @@ +"""``filter_queryset_by_organization`` must fail closed. + +Every caller reaches it with OrganizationFilterBackend either disabled +(``skip_org_filter = True``) or absent — the function-based internal handlers +have no filter backend at all — so this helper is their only tenant boundary. +Returning the queryset unfiltered when there is no organization context +therefore returns every organization's rows, and the absent-header case is +reachable: the internal auth middleware warns and continues rather than +rejecting. +""" + +import secrets + +import pytest +from account_v2.models import Organization +from django.test import TestCase +from utils.organization_utils import filter_queryset_by_organization +from workflow_manager.workflow_v2.models.workflow import Workflow + + +class _Request: + """Stands in for the request object the helper reads context off.""" + + def __init__(self, organization_id=None, path="/internal/test/"): + if organization_id is not None: + self.organization_id = organization_id + self.path = path + + +@pytest.mark.django_db +class FilterQuerysetByOrganizationTest(TestCase): + def setUp(self) -> None: + self.org_a = self._org("a") + self.org_b = self._org("b") + self.wf_a = Workflow.objects.create( + workflow_name=f"wf-a-{secrets.token_hex(3)}", organization=self.org_a + ) + self.wf_b = Workflow.objects.create( + workflow_name=f"wf-b-{secrets.token_hex(3)}", organization=self.org_b + ) + + def _org(self, tag: str) -> Organization: + slug = f"org-{tag}-{secrets.token_hex(3)}" + return Organization.objects.create( + name=slug, display_name=slug, organization_id=slug + ) + + def _filter(self, request): + # _base_manager, not objects: Workflow's default manager is itself + # org-scoped off UserContext, which would empty the queryset before the + # helper ever ran and make these tests pass for the wrong reason. The + # helper's contract is "given a queryset, scope it", so hand it an + # unscoped one. + return filter_queryset_by_organization(Workflow._base_manager.all(), request) + + def test_missing_org_context_returns_nothing(self): + """The header is optional at the middleware, so this is reachable.""" + assert not self._filter(_Request()).exists() + + def test_falsy_org_context_returns_nothing(self): + for falsy in ("", None): + with self.subTest(organization_id=falsy): + assert not self._filter(_Request(organization_id=falsy)).exists() + + def test_unresolvable_org_returns_nothing(self): + request = _Request(organization_id="does-not-exist") + assert not self._filter(request).exists() + + def test_valid_org_returns_only_its_own_rows(self): + rows = self._filter(_Request(organization_id=self.org_a.organization_id)) + assert list(rows) == [self.wf_a] + + def test_other_org_rows_are_never_included(self): + for org, mine, theirs in ( + (self.org_a, self.wf_a, self.wf_b), + (self.org_b, self.wf_b, self.wf_a), + ): + with self.subTest(org=org.organization_id): + rows = list(self._filter(_Request(organization_id=org.organization_id))) + assert mine in rows + assert theirs not in rows diff --git a/backend/workflow_manager/file_execution/internal_views.py b/backend/workflow_manager/file_execution/internal_views.py index 8fed360740..256331a0d7 100644 --- a/backend/workflow_manager/file_execution/internal_views.py +++ b/backend/workflow_manager/file_execution/internal_views.py @@ -29,10 +29,10 @@ class FileExecutionInternalViewSet(viewsets.ModelViewSet): serializer_class = WorkflowFileExecutionSerializer lookup_field = "id" - # Backward compat: workers may call without X-Organization-ID during - # rolling deployments. Safe because internal APIs require service API key - # and get_queryset() applies org filtering when header is present. - # Remove once all workers reliably pass X-Organization-ID. + # OrganizationFilterBackend is off here; get_queryset() scopes instead, via + # filter_queryset_by_organization. That helper fails closed, so a worker + # calling without X-Organization-ID now gets zero rows rather than every + # organization's — the header is required in practice, not optional. skip_org_filter = True def get_object(self): diff --git a/backend/workflow_manager/internal_views.py b/backend/workflow_manager/internal_views.py index 4f9de9f0aa..689a8e2b9f 100644 --- a/backend/workflow_manager/internal_views.py +++ b/backend/workflow_manager/internal_views.py @@ -49,10 +49,10 @@ class WorkflowExecutionInternalViewSet(viewsets.ReadOnlyModelViewSet): serializer_class = WorkflowExecutionSerializer lookup_field = "id" - # Backward compat: workers may call without X-Organization-ID during - # rolling deployments. Safe because internal APIs require service API key - # and get_queryset() applies org filtering when header is present. - # Remove once all workers reliably pass X-Organization-ID. + # OrganizationFilterBackend is off here; get_queryset() scopes instead, via + # filter_queryset_by_organization. That helper fails closed, so a worker + # calling without X-Organization-ID now gets zero rows rather than every + # organization's — the header is required in practice, not optional. skip_org_filter = True def get_queryset(self): diff --git a/backend/workflow_manager/workflow_v2/views.py b/backend/workflow_manager/workflow_v2/views.py index fefba8c21a..aa31f0f671 100644 --- a/backend/workflow_manager/workflow_v2/views.py +++ b/backend/workflow_manager/workflow_v2/views.py @@ -403,10 +403,10 @@ class WorkflowExecutionInternalViewSet(viewsets.ReadOnlyModelViewSet): serializer_class = WorkflowExecutionSerializer lookup_field = "id" - # Backward compat: workers may call without X-Organization-ID during - # rolling deployments. Safe because internal APIs require service API key - # and get_queryset() applies org filtering when header is present. - # Remove once all workers reliably pass X-Organization-ID. + # OrganizationFilterBackend is off here; get_queryset() scopes instead, via + # filter_queryset_by_organization. That helper fails closed, so a worker + # calling without X-Organization-ID now gets zero rows rather than every + # organization's — the header is required in practice, not optional. skip_org_filter = True def get_queryset(self): diff --git a/workers/ide_callback/tasks.py b/workers/ide_callback/tasks.py index e35b76f933..1a298dffe3 100644 --- a/workers/ide_callback/tasks.py +++ b/workers/ide_callback/tasks.py @@ -243,9 +243,14 @@ def ide_index_complete( organization_id=org_id, ) except Exception: - logger.warning( + # Non-fatal — primary indexing already succeeded — but not + # harmless: without the status, check_extraction_status stays + # False and every later Answer Prompt re-runs the full X2Text + # extraction. ERROR because nothing downstream reports it. + logger.error( "Failed to mark extraction_status for document %s " - "profile %s; primary indexing succeeded.", + "profile %s; extraction will be repeated on every " + "subsequent prompt run.", document_id, profile_manager_id, exc_info=True,