fix(workspaces): preserve file workspace memberships on PUT and PATCH - #280
Conversation
PUT /file/{file_id} and PATCH /file/{file_id} both delete-then-reinsert
the file internally, which wiped all workspace_files rows for that file_id.
Fix: snapshot the file's workspace memberships before the delete step,
then restore them after re-indexing completes.
- Add PartitionFileManager.get_file_workspaces() in vectordb/utils.py
- Expose it as MilvusDB.get_file_workspaces() in vectordb/vectordb.py
- update_file_metadata() in indexer.py: snapshot + asyncio.gather restore
- put_file() in routers/indexer.py: snapshot + pass to add_file.remote()
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (2)
📝 WalkthroughWalkthroughSnapshot workspace memberships before deleting a file during replace/metadata updates, then restore those memberships after re-indexing by adding files back to the captured workspace IDs. Changes
Sequence Diagram(s)sequenceDiagram
participant Router as PUT Handler<br/>(Router)
participant Indexer as Indexer Task
participant VectorDB as VectorDB API
participant PFM as PartitionFileManager
Router->>VectorDB: get_file_workspaces(file_id, partition)
VectorDB->>PFM: query workspaces for file
PFM-->>VectorDB: [workspace_ids]
VectorDB-->>Router: [workspace_ids]
Router->>Indexer: delete_file(file_id, partition)
Indexer->>VectorDB: remove_file_from_all_workspaces(file_id, partition)
VectorDB->>PFM: delete WorkspaceFile rows
PFM-->>VectorDB: deleted
Router->>Indexer: add_file(path, metadata, partition, workspace_ids)
Indexer->>VectorDB: add_documents(...)
Indexer->>VectorDB: add_files_to_workspace(ws_id, [file_id]) for each ws_id
VectorDB->>PFM: create WorkspaceFile rows
PFM-->>VectorDB: created
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related PRs
Poem
🚥 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 unit tests (beta)
📝 Coding Plan
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 |
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@openrag/components/indexer/indexer.py`:
- Around line 198-209: Capture workspace_ids from the original partition as
before but, before re-adding, fetch the file's current partition (e.g., via
vectordb.get_file or a get_file_partition helper) and only attempt restoration
if the current partition still equals the original partition; wrap the restore
loop that calls vectordb.add_files_to_workspace.remote(...) in a try/except so
any exception is logged as a warning/debug and swallowed (best-effort) so a
successful reindex does not fail; keep using the existing symbols workspace_ids,
file_id, partition, vectordb.get_file_workspaces,
vectordb.add_files_to_workspace and the surrounding
delete_file/async_add_documents flow.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: edce78f9-cd20-4eb7-b155-faa284a36355
📒 Files selected for processing (4)
openrag/components/indexer/indexer.pyopenrag/components/indexer/vectordb/utils.pyopenrag/components/indexer/vectordb/vectordb.pyopenrag/routers/indexer.py
| # Snapshot workspace memberships before deletion so they can be restored. | ||
| workspace_ids = await vectordb.get_file_workspaces.remote(file_id, partition) | ||
|
|
||
| await self.delete_file(file_id, partition) | ||
| await vectordb.async_add_documents.remote(docs, user=user) | ||
|
|
||
| # Restore workspace memberships that existed before the delete. | ||
| if workspace_ids: | ||
| await asyncio.gather( | ||
| *[vectordb.add_files_to_workspace.remote(ws_id, [file_id]) for ws_id in workspace_ids] | ||
| ) | ||
| log.debug("Restored workspace memberships after metadata update.", workspace_ids=workspace_ids) |
There was a problem hiding this comment.
Guard workspace restoration on partition moves, and keep it best-effort.
workspace_ids here are captured from the original partition, but openrag/routers/indexer.py Lines 339-346 allow PATCH to move the file to a different partition. Re-adding those IDs on Lines 205-208 would recreate workspace_files rows for workspaces that still belong to the old partition, and any exception here currently makes the PATCH fail after the file has already been reindexed successfully.
Proposed fix
# Snapshot workspace memberships before deletion so they can be restored.
workspace_ids = await vectordb.get_file_workspaces.remote(file_id, partition)
await self.delete_file(file_id, partition)
await vectordb.async_add_documents.remote(docs, user=user)
# Restore workspace memberships that existed before the delete.
- if workspace_ids:
- await asyncio.gather(
- *[vectordb.add_files_to_workspace.remote(ws_id, [file_id]) for ws_id in workspace_ids]
- )
- log.debug("Restored workspace memberships after metadata update.", workspace_ids=workspace_ids)
+ target_partition = metadata.get("partition", partition)
+ if workspace_ids and target_partition == partition:
+ try:
+ await asyncio.gather(
+ *[vectordb.add_files_to_workspace.remote(ws_id, [file_id]) for ws_id in workspace_ids]
+ )
+ log.debug(
+ "Restored workspace memberships after metadata update.",
+ workspace_ids=workspace_ids,
+ )
+ except Exception as ws_err:
+ log.warning(
+ "Failed to restore workspace memberships after metadata update.",
+ error=str(ws_err),
+ workspace_ids=workspace_ids,
+ )
+ elif workspace_ids:
+ log.info(
+ "Skipping workspace restoration because the file moved to a different partition.",
+ old_partition=partition,
+ new_partition=target_partition,
+ workspace_ids=workspace_ids,
+ )Based on learnings: workspace association is a best-effort metadata step and failures should be logged without turning a successfully indexed file into a failed operation.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@openrag/components/indexer/indexer.py` around lines 198 - 209, Capture
workspace_ids from the original partition as before but, before re-adding, fetch
the file's current partition (e.g., via vectordb.get_file or a
get_file_partition helper) and only attempt restoration if the current partition
still equals the original partition; wrap the restore loop that calls
vectordb.add_files_to_workspace.remote(...) in a try/except so any exception is
logged as a warning/debug and swallowed (best-effort) so a successful reindex
does not fail; keep using the existing symbols workspace_ids, file_id,
partition, vectordb.get_file_workspaces, vectordb.add_files_to_workspace and the
surrounding delete_file/async_add_documents flow.
|
The fix sounds good for the current behaviour, but why do we delete/create the file entry in postgres on PUT/PATCH? Why not a SQL update? |
- Add GET /partition/{partition}/files/{file_id}/workspaces endpoint to
list all workspaces a file belongs to
- Add integration tests verifying workspace memberships are preserved
after PUT (file replace) and PATCH (metadata update) operations
Ahmath-Gadji
left a comment
There was a problem hiding this comment.
LGTM. I've added a test for workspace membership preservation after put or patch
|
I suggest we do not delay the bug fix and merge this PR, but treat this issue afterwards: #286 |
Closes #286. PATCH (metadata update) no longer deletes the file. Instead, it fetches existing chunks with their Milvus _id and vectors, merges new metadata, and upserts back into Milvus — no re-embedding. The PostgreSQL file record is updated in-place, so workspace FK references and file_count are never disturbed. PUT (file replace) now deletes only Milvus chunks while preserving the PostgreSQL File row. New chunks are embedded and inserted, then the PG row is updated in-place. Workspace associations survive without the snapshot/restore workaround from PR #280. New methods: - PartitionFileManager.update_file_metadata_in_db() — in-place PG update - PartitionFileManager.update_file_in_partition() — in-place PG update for PUT (metadata + relationship_id + parent_id) - MilvusDB.delete_file_chunks() — Milvus-only delete, PG untouched - MilvusDB.upsert_file_metadata() — fetch chunks with vectors, merge metadata, upsert to Milvus + update PG - MilvusDB.add_documents_for_existing_file() — embed + insert new chunks, update existing PG row (no duplicate check) - Indexer.replace_file_documents() — routes to the above - get_file_chunks() gains include_vectors parameter
Closes #286. PATCH (metadata update) no longer deletes the file. Instead, it fetches existing chunks with their Milvus _id and vectors, merges new metadata, and upserts back into Milvus — no re-embedding. The PostgreSQL file record is updated in-place, so workspace FK references and file_count are never disturbed. PUT (file replace) now deletes only Milvus chunks while preserving the PostgreSQL File row. New chunks are embedded and inserted, then the PG row is updated in-place. Workspace associations survive without the snapshot/restore workaround from PR #280. New methods: - PartitionFileManager.update_file_metadata_in_db() — in-place PG update - PartitionFileManager.update_file_in_partition() — in-place PG update for PUT (metadata + relationship_id + parent_id) - MilvusDB.delete_file_chunks() — Milvus-only delete, PG untouched - MilvusDB.upsert_file_metadata() — fetch chunks with vectors, merge metadata, upsert to Milvus + update PG - MilvusDB.add_documents_for_existing_file() — embed + insert new chunks, update existing PG row (no duplicate check) - Indexer.replace_file_documents() — routes to the above - get_file_chunks() gains include_vectors parameter
Summary
Fixes #279.
PUT /indexer/partition/{partition}/file/{file_id}andPATCH /indexer/partition/{partition}/file/{file_id}both use a delete-then-reinsert pattern internally. The delete step calledremove_file_from_all_workspaces(), permanently wiping allworkspace_filesrows for the file even though thefile_idnever changed. After the operation the file existed in the vector DB but had lost all its workspace associations.Fix
Snapshot the file's current workspace memberships before the delete step, then restore them after re-indexing completes.
Changes
vectordb/utils.py— addPartitionFileManager.get_file_workspaces(file_id, partition) -> list[str]that queriesworkspace_files JOIN workspacesfiltered by partitionvectordb/vectordb.py— expose it asasync MilvusDB.get_file_workspaces(...)(Ray-callable)indexer/indexer.py—update_file_metadata()snapshots workspace IDs beforedelete_file, then restores them viaasyncio.gatherafterasync_add_documentscompletesrouters/indexer.py—put_file()snapshots workspace IDs beforedelete_file.remote, then passes them asworkspace_idstoadd_file.remote(...)(which already handles re-association after indexing)Testing
All 205 unit tests pass (
uv run pytest --ignore=openrag/test_token_validation.py). The skippedtest_token_validation.pyfailure is a pre-existing Ray session lock issue unrelated to this change.Summary by CodeRabbit
New Features
Bug Fixes