-
Notifications
You must be signed in to change notification settings - Fork 709
UN-3815 [FIX] Apply organization scoping to prompt-studio child models #2213
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
447d0c9
09d320b
14f94cd
18a53f9
14b7e68
0dce94e
cc41956
8b182c9
6ed54b3
34ff8c4
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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. | ||
|
|
||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [Medium] [Lens 16, 3] —
But There is also a null-policy split: when Fix: drop |
||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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( | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Instead of returning a JsonResponse, can you raise it as an error to keep it consistent with the rest of the codebase and let DRF validation handle it?
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. NOT RESOLVED — no push since this was raised, so just re-flagging it rather than treating it as ignored. Still worth doing alongside the other changes on this handler; note the retryable-500 point at |
||
| { | ||
| "success": False, | ||
| "error": "Extraction status could not be recorded", | ||
| }, | ||
| status=status.HTTP_500_INTERNAL_SERVER_ERROR, | ||
| ) | ||
|
|
||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [Medium] [Lens 8] — surfacing the failure is right; mapping a permanent condition onto a retryable status is not
Fix: have This is a new consequence of the correct fix to prior finding #5, not a regression on it. |
||
| return JsonResponse({"success": True}) | ||
|
|
||
| except Exception as e: | ||
| logger.exception("extraction_status internal API failed") | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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) | ||
|
athul-rs marked this conversation as resolved.
|
||
| 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. | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [Medium] [Lens 3, 10] — this diagnostic is dead code: the unfixed lookup above short-circuits first (prior finding #13, NOT RESOLVED) An identical, unguarded So the Fix: hoist the |
||
| 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 | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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 @@ | |
| 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. | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [Low] [Lens 16] — this comment and its own test's docstring disagree on whether resolve-before-clear is load-bearing Here: "clearing first would leave the tool with no default at all when it does not match".
Your correction in the previous thread was right, and I re-confirmed it by mutation — ordering-only revert with the transaction retained -> 13 passed; full revert to Suggest qualifying the comment — the ordering is belt-and-braces given the surrounding transaction — or dropping the causal claim and keeping "resolve before clearing". |
||
| profile_manager = get_object_or_404( | ||
| ProfileManager, | ||
| pk=default_profile, | ||
| prompt_studio_tool=prompt_tool, | ||
|
athul-rs marked this conversation as resolved.
|
||
| ) | ||
|
|
||
| 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 @@ | |
| 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 | ||
| ) | ||
|
athul-rs marked this conversation as resolved.
|
||
|
|
||
| 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, | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [Medium] [Lens 10] — this warning fires on the ordinary never-indexed path, and the case it warns about still returns 200 (prior finding #12, PARTIALLY RESOLVED)
And in the case it does warn about, the handler still returns Fix: probe |
||
| ) | ||
| for index_manager in index_managers: | ||
| raw_index_id = index_manager.raw_index_id | ||
| DocumentIndexingService.remove_document_indexing( | ||
|
|
@@ -1155,7 +1202,20 @@ | |
| 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( | ||
|
Check failure on line 1212 in backend/prompt_studio/prompt_studio_core_v2/views.py
|
||
| "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, | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
[Low] [Lens 16] — prior finding NOT fixed: this still claims PUT and PATCH
The PR edited this exact line (
Handles GET,POST,PUT,PATCH and DELETE->Handles GET, POST, PUT and PATCH), removing only the DELETE half.urls.py:37-42binds three paths, allget/post, andFileManagementViewSetdefines onlylist,download(get) andupload(post) — there is noupdate/partial_updateand no PUT/PATCH route anywhere in the repo.Suggest
Handles GET and POST., or dropping the line — the routes are two files away and it will rot again on the next route change.