UN-3479 [FEAT] stream Prompt Studio file uploads to storage - #1988
Merged
chandrasekharan-zipstack merged 16 commits intoMay 25, 2026
Conversation
…l instances Enables the unstract-python-client SDK migration subpackage to drive org-to-org data migration purely through admin-issued Platform API keys. Specifically: - adapter_processor_v2/models.py: AdapterInstanceModelManager.for_user returns non-frictionless adapters for service-account callers (was: all()) - permissions/permission.py: IsFrictionLessAdapter grants access to service accounts on non-frictionless adapters, keeping the friction-first check - prompt_studio/permission.py: PromptAcesssToUser short-circuits to True for service accounts so Platform API can GET/POST prompts - tool_instance_v2/views.py: get_queryset scopes via Workflow.for_user so service accounts see all tool instances under workflows they can access Plan: org-to-org data migration v1 (KB: zipstuff/org-data-migration/05).
The SDK migration subpackage relies on list-by-name GET as the cross-run idempotency check (Layer 2). Without this filter, every re-run would re-create adapters that already exist on the target org. Plan: org-to-org data migration v1 (KB: zipstuff/org-data-migration/05).
Mirror the adapter pattern from c05dc05: SDK migrator uses name-based GET against target to detect existing rows before POST. - ConnectorInstanceViewSet.get_queryset: thread CIKey.CONNECTOR_NAME through FilterHelper. - TagViewSet: declare filterset_fields=["name"] so DjangoFilterBackend honors ?name=. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Tag URLs previously bound only GET on list/detail, so the migrator (and any other API consumer) got 405 on POST tags/. Wire create/partial_update/ destroy through the same TagViewSet — permission_classes already cover the auth path; no behavior change for callers that only GET. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
ConnectorInstanceSerializer.to_representation overrides connector_mode from the catalog every response, so any client-supplied value is silently discarded. Make that explicit via extra_kwargs so DRF OPTIONS reports the field as read-only and round-trippers don't trip the choice validator (catalog mode 'FILESYSTEM' vs model choice 'FILE_SYSTEM'). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Needed by the org-to-org migration SDK's idempotent get-or-create flow: without a name filter the SDK had to list-all-then-linear-search every time. Adds WorkflowKey.WF_NAME to the existing FilterHelper.build_filter_args call — same shape as the recently-added adapter_name / connector_name / tag name filters. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Needed by the org-to-org migration SDK: after CustomTool migration the target's PromptStudioRegistry gets a freshly-minted prompt_registry_id, and downstream ToolInstance migration needs to remap source.tool_id -> target.tool_id (both are registry UUIDs). With this filter the SDK can GET /prompt-studio/registry/?custom_tool=<tool_id> on either side to resolve the registry id without needing to know it up front. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Mirrors the existing exact-match name filters added for adapter, connector, tag, and workflow lists. Migration SDK's get-or-create flow queries by name on the target before deciding fresh vs adopt; without this the SDK had to list-all-then-linear-scan every time. Pipeline keeps the icontains ?search= alongside; api_deployment keeps the icontains ?search= alongside. New filters are exact-match only. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Contributor
|
Caution Review failedPull request was closed or merged during review No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (1)
Summary by CodeRabbit
WalkthroughThis PR adds a streaming file upload helper (write_streaming), integrates it into PromptStudioFileHelper for IDE uploads, and adds tests covering bytes-path, chunked streaming, read fallback, close/error propagation, and cleanup behavior. ChangesStreaming File Upload Refactoring
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
chandrasekharan-zipstack
marked this pull request as ready for review
May 24, 2026 08:30
Contributor
|
| Filename | Overview |
|---|---|
| backend/utils/file_storage/helpers/streaming_writer.py | New helper that streams file uploads chunk-by-chunk; bytes fast-path preserved, close()-on-success-path propagates (previous review issues resolved), and cleanup logic is correct for write failures. |
| backend/utils/file_storage/helpers/prompt_studio_file_helper.py | Minimal callsite change — both upload methods now delegate to write_streaming instead of calling fs_instance.write directly; no logic changes beyond the delegation. |
| backend/utils/tests/test_prompt_studio_file_helper.py | Comprehensive unit tests covering bytes fast-path, Django UploadedFile streaming, read()-fallback, mid-stream failure + cleanup, non-callable chunks attribute, close()-propagation on success, and close()-swallow on failure paths. |
Flowchart
%%{init: {'theme': 'neutral'}}%%
flowchart TD
A[write_streaming called] --> B{file_data is bytes?}
B -- Yes --> C[fs_instance.write single-shot]
C -- success --> D[return]
C -- exception --> E[_remove_file] --> F[re-raise]
B -- No --> G[fs_instance.fs.open with block_size=8MiB]
G --> H{has callable chunks?}
H -- Yes --> I[iterate file_data.chunks chunk_size=8MiB]
H -- No --> J[iterate file_data.read 8MiB sentinel b'']
I --> K[out.write each chunk]
J --> K
K -- write exception --> L[out.close swallow] --> M[_remove_file] --> N[re-raise]
K -- loop complete --> O[out.close propagate]
O -- success --> P[return]
O -- close exception --> Q[re-raise to caller]
Prompt To Fix All With AI
Fix the following 1 code review issue. Work through them one at a time, proposing concise fixes.
---
### Issue 1 of 1
backend/utils/file_storage/helpers/streaming_writer.py:46-48
When `out.close()` raises on the success path, `_remove_file` is never called. For cloud providers (S3/GCS) this is harmless because a failed multipart commit leaves no object at the final path, so a `rm` would be a silent no-op. For local-filesystem or POSIX-backed providers, however, `close()` can fail after some data has been flushed to disk (e.g. a deferred `fsync` hitting a full disk), leaving a partial file at `file_path` that the next retry will silently overwrite rather than detect. A single best-effort `_remove_file` call after the failed `close()` — mirroring what the write-failure path already does — would keep the two paths consistent without masking the original exception.
```suggestion
# On success, propagate close() errors — provider multipart commits
# happen here, so a failure means the upload did not finalize.
try:
out.close()
except Exception:
_remove_file(fs_instance, file_path)
raise
```
Reviews (7): Last reviewed commit: "refactor(streaming_writer): tighten clea..." | Re-trigger Greptile
Greptile flagged on PR #1987 that the prior switch from `created_by=user` to `workflow__in=Workflow.for_user(user)` silently widened tool-instance visibility for regular users with shared-workflow access — they would start seeing every member's tool instances in those workflows, not just their own. Restrict the widened queryset to `is_service_account` callers so the migration SDK still gets org-wide enumeration via Platform API keys, while regular users keep their original per-creator scope. Reported-by: greptile-apps[bot] (PR #1987, P1)
…-3479] `upload_for_ide` / `upload_converted_for_ide` previously buffered the entire payload in memory before writing — for a 100 MB PDF that meant ~100 MB held per concurrent upload per worker. Replaces the single-shot write with chunked streaming via the underlying fsspec handle when the input is a file-like (Django UploadedFile / IOBase). - New `streaming_writer.write_streaming(fs_instance, path, data)` helper in its own module (kept Django-import-free so it is unit-testable without booting full settings). - Bytes inputs stay single-shot — preserves caller contract for the converted-PDF flow that passes already-materialised bytes. - Opens the destination with `block_size=8 MB` so providers use their native multipart upload primitive (GCS resumable, S3 multipart). - Iterates via the source's own `chunks()` when available, otherwise fixed-size `read()` blocks. - On mid-stream exception, attempts provider-level multipart abort (`abort_mpu` for s3fs, `discard` for gcsfs) so partial uploads don't leak orphaned parts. Unrecognised providers log and rethrow. Benefits the SDK org-migration files phase (companion PR in unstract-python-client) and the FE concurrent-upload path simultaneously — no API surface change, no behavior change for callers. Branched off `feat/org-migration-platform-api-gaps` so the full stack (platform-API filters + service-account perms + streaming write) can be tested together. PR targets the same parent branch. 6 unit tests cover: bytes fast-path, chunks() streaming with multipart-hint propagation, read() fallback for sources without chunks(), abort_mpu on stream error, discard fallback when abort_mpu absent, and a no-coalesce assertion that catches the "collapse to one buffer" regression.
…t.mock Rewrites the streaming-writer test suite as bare pytest: - Removes ``unittest.mock.MagicMock`` / ``patch`` imports. - Hand-rolled ``FakeHandle`` / ``FakeFs`` / ``FakeStorage`` classes expose only the surface the helper actually touches (write, fs.open, abort_mpu/discard, close). - Module-level test functions with pytest fixtures, no test-class grouping needed. - Adds one extra test (``test_unknown_abort_method_does_not_mask_original_error``) that asserts the original exception propagates and the leak warning is emitted when no abort hook is exposed. Self-sufficient under the existing ``uv run pytest`` runner; no new dev dependency.
chandrasekharan-zipstack
force-pushed
the
feat/prompt-studio-streaming-upload
branch
from
May 24, 2026 08:39
7882fc7 to
fb50704
Compare
Provider-native multipart commits (s3fs, gcsfs) happen inside close() — a close failure on the success path means the upload did not finalize. The previous code logged-and-swallowed any close error in finally, so a silent partial upload looked like a successful write. Track success/failure via a flag: on success, let close() raise; on the exception path, keep the swallow-and-log behaviour so a close failure cannot mask the original error reaching the caller. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
hasattr() returns True even for a non-callable data attribute named ``chunks``, which would then raise TypeError when invoked. Switch to callable(getattr(...)) so plain attributes correctly fall through to the read-based fallback. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Replace _safe_abort's fsspec-internal introspection (abort_mpu / discard) with a uniform fs.rm cleanup on the public API. Local partial files are removed; orphaned S3/GCS multipart parts are left to the bucket's AbortIncompleteMultipartUpload lifecycle rule. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
muhammad-ali-e
approved these changes
May 25, 2026
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@backend/utils/file_storage/helpers/streaming_writer.py`:
- Around line 30-36: The code currently creates chunks_iter (using
file_data.chunks or iter(...)) before entering the try block, so if
file_data.chunks() raises the opened file handle (out from fs_instance.fs.open)
remains open and cleanup (rm) is skipped; move the creation of chunks_iter
inside the protected try/finally (or try/except) block immediately after opening
out (where STREAMING_CHUNK_SIZE, file_data.chunks, and the iter(lambda:
file_data.read(...)) logic live) so any exception from file_data.chunks()
triggers the existing cleanup path that closes out and calls rm.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: a9580bb6-ccc5-4121-bdb4-e4ec9aa3b5e9
📒 Files selected for processing (3)
backend/utils/file_storage/helpers/prompt_studio_file_helper.pybackend/utils/file_storage/helpers/streaming_writer.pybackend/utils/tests/test_prompt_studio_file_helper.py
kirtimanmishrazipstack
approved these changes
May 25, 2026
…s_iter Move chunks_iter construction inside the try so a hypothetical non-generator chunks() that raises eagerly still triggers cleanup. No behavioural change for current callers (UploadedFile.chunks is a generator; BytesIO fallback uses iter() with a sentinel). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
|
Contributor
Test ResultsSummary
Runner Tests - Full Report
SDK1 Tests - Full Report
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Streams Prompt Studio file uploads to storage instead of buffering the whole payload in memory. Touches only the upload-side helpers; no API surface change, no request/response shape change.
streaming_writer.write_streaming(fs_instance, path, data)helper inbackend/utils/file_storage/helpers/streaming_writer.py. Lives outsideprompt_studio_file_helper.pyso it is unit-testable without booting Django settings (the helper transitively importsfile_management→backend.__init__→ celery → settings).PromptStudioFileHelper.upload_for_ideandupload_converted_for_idedelegate to the new helper.UploadedFile,IOBase) are streamed: destination is opened viafs_instance.fs.open(..., block_size=8 * 1024 * 1024)so providers use their native multipart upload primitive; source is iterated viachunks()when available, else fixed-sizeread().abort_mpufor s3fs,discardfor gcsfs) so partial uploads don't leak orphaned parts in the bucket.Why
fs_instance.write(..., data=file_data.read())materialises the entire upload payload in memory before the storage layer touches it. A 100 MB PDF held ~100 MB per concurrent upload per worker; under FE parallel uploads or the org-migration SDK's per-file roundtrip pattern (which adds another base64 inflation on top), this drives worker memory pressure unnecessarily.Tracked in UN-3479. Surfaced while building the SDK's
filesphase against the existing endpoints — the same upload path is hot for the FE concurrent-upload flow.How
write_streamingbranches onisinstance(file_data, bytes)to preserve the bytes fast path. Everything else goes through the streaming branch.block_sizeis set explicitly because fsspec defaults vary per provider (GCS, S3, MinIO, Azure) — without it some providers buffer to memory until much larger thresholds._safe_abortwalks a small list of known per-provider abort hooks (abort_mpu,discard); unrecognised handles are logged and rethrown without abort (better to leak a part than mask the original exception).close()always runs infinally; any close-side failure is logged but not allowed to mask the primary exception.Can this PR break any existing features. If yes, please list possible items. If no, please explain why. (PS: Admins do not merge the PR without this section filled)
No — change is behavior-preserving for all caller types:
upload_for_ide) hit the same single-shotfs_instance.writepath as before.upload_for_ide, SDK migration upload via the same endpoint) receive identical written bytes; only the per-chunk granularity at the storage layer changes.block_size; this is the supported fsspec path, not a custom protocol.The only edge-case to watch is the abort-on-failure path: if a future provider exposes a different abort hook name, we'd leak the part. That's documented in
_safe_abort's logger.warning and is recoverable via bucket lifecycle rules.Database Migrations
None.
Env Config
None.
Relevant Docs
unstract-python-client:docs/internal/files-migration-plan.md§6.~/Documents/Obsidian Vault/zipstuff/org-data-migration/Related Issues or PRs
feat/org-migration-platform-api-gaps— this PR targets that branch (notmain) so the platform-API filter additions + service-account perms + this streaming change can be reviewed and tested as a single coherent stack against the SDK.Zipstack/unstract-python-client#15(feat/org-migration). The SDK'sfilesphase exercisesupload_for_ideon the target side.Dependencies Versions
None.
Notes on Testing
backend/utils/tests/test_prompt_studio_file_helper.py:fs.opencall).UploadedFile-like source streams viachunks(),block_sizepropagated tofs.open.chunks()falls back toread()-block iteration.abort_mpuand still closes the handle.discardwhenabort_mpuis absent.writecalls of<= chunk_size, proving the helper doesn't accumulate in a single buffer before flushing.backend/utils/testscontinue to pass.http://localhost:8000).Screenshots
N/A — backend-only refactor.
Checklist
I have read and understood the Contribution Guidelines.