Skip to content

fix(workspaces): preserve file workspace memberships on PUT and PATCH - #280

Merged
paultranvan merged 2 commits into
devfrom
fix/workspace-membership-lost-on-file-replace
Mar 17, 2026
Merged

fix(workspaces): preserve file workspace memberships on PUT and PATCH#280
paultranvan merged 2 commits into
devfrom
fix/workspace-membership-lost-on-file-replace

Conversation

@EnjoyBacon7

@EnjoyBacon7 EnjoyBacon7 commented Mar 12, 2026

Copy link
Copy Markdown
Collaborator

Summary

Fixes #279.

PUT /indexer/partition/{partition}/file/{file_id} and PATCH /indexer/partition/{partition}/file/{file_id} both use a delete-then-reinsert pattern internally. The delete step called remove_file_from_all_workspaces(), permanently wiping all workspace_files rows for the file even though the file_id never 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 — add PartitionFileManager.get_file_workspaces(file_id, partition) -> list[str] that queries workspace_files JOIN workspaces filtered by partition
  • vectordb/vectordb.py — expose it as async MilvusDB.get_file_workspaces(...) (Ray-callable)
  • indexer/indexer.pyupdate_file_metadata() snapshots workspace IDs before delete_file, then restores them via asyncio.gather after async_add_documents completes
  • routers/indexer.pyput_file() snapshots workspace IDs before delete_file.remote, then passes them as workspace_ids to add_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 skipped test_token_validation.py failure is a pre-existing Ray session lock issue unrelated to this change.

Summary by CodeRabbit

  • New Features

    • Added an endpoint to list which workspaces a file belongs to within a partition.
  • Bug Fixes

    • Files now preserve their workspace associations when updated or re-indexed, preventing loss of memberships during file replacement or metadata updates.
    • File removal now correctly clears all workspace-to-file associations within the relevant scope.

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()
@coderabbitai

coderabbitai Bot commented Mar 12, 2026

Copy link
Copy Markdown

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 06c2184e-9f9d-4b9a-97b5-f23de7c1ee34

📥 Commits

Reviewing files that changed from the base of the PR and between 0e09994 and 8db4596.

📒 Files selected for processing (2)
  • openrag/routers/workspaces.py
  • tests/api_tests/test_workspaces.py

📝 Walkthrough

Walkthrough

Snapshot 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

Cohort / File(s) Summary
VectorDB utils & API
openrag/components/indexer/vectordb/utils.py, openrag/components/indexer/vectordb/vectordb.py
Added get_file_workspaces(file_id, partition) to PartitionFileManager and MilvusDB; updated remove_file_from_all_workspaces() to delete WorkspaceFile associations scoped to a partition.
Indexer core logic
openrag/components/indexer/indexer.py
In update_file_metadata, snapshot workspaces via vectordb.get_file_workspaces() before deletion, then restore memberships after documents are re-added using vectordb.add_files_to_workspace().
HTTP routers
openrag/routers/indexer.py, openrag/routers/workspaces.py
put_file() now captures existing workspace IDs and passes them to add_file() during re-indexing; added GET /partition/{partition}/files/{file_id}/workspaces endpoint to list a file's workspace IDs. Public RPC add_file accepts optional workspace_ids.
Tests
tests/api_tests/test_workspaces.py
Added tests verifying workspace memberships are preserved after PUT (replace) and PATCH (metadata update); includes helpers for uploading files and retrieving workspace memberships.
Manifests
requirements.txt, pyproject.toml
Unchanged in behavior; listed in manifest.

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
Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Possibly related PRs

  • Add workspaces #272: Related workspace feature changes that touch PartitionFileManager and workspace-related APIs; likely overlaps in implementation and tests.

Poem

🐰 I hopped around the indexed glade,
Lost memberships the old way made,
I noted, saved, then stitched them back —
No file again will lose its track. 🥕

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 40.00% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title clearly and concisely describes the main fix: preserving file workspace memberships during PUT and PATCH operations, which directly addresses the core issue being resolved.
Linked Issues check ✅ Passed The pull request fully implements the fix described in issue #279 by snapshotting workspace memberships before deletion and restoring them after re-indexing, plus adds comprehensive integration tests verifying the fix works for both PUT and PATCH operations.
Out of Scope Changes check ✅ Passed All changes are directly related to fixing the workspace membership preservation issue; no unrelated or out-of-scope modifications are present in the changeset.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
  • 📝 Generate docstrings (stacked PR)
  • 📝 Generate docstrings (commit on current branch)
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Post copyable unit tests in a comment
  • Commit unit tests in branch fix/workspace-membership-lost-on-file-replace
📝 Coding Plan
  • Generate coding plan for human review comments

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.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

📥 Commits

Reviewing files that changed from the base of the PR and between c1fe53a and 0e09994.

📒 Files selected for processing (4)
  • openrag/components/indexer/indexer.py
  • openrag/components/indexer/vectordb/utils.py
  • openrag/components/indexer/vectordb/vectordb.py
  • openrag/routers/indexer.py

Comment on lines +198 to +209
# 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)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

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.

@paultranvan

Copy link
Copy Markdown
Collaborator

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 Ahmath-Gadji left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM. I've added a test for workspace membership preservation after put or patch

@paultranvan

Copy link
Copy Markdown
Collaborator

I suggest we do not delay the bug fix and merge this PR, but treat this issue afterwards: #286

@paultranvan
paultranvan merged commit 5f05d37 into dev Mar 17, 2026
4 checks passed
@paultranvan
paultranvan deleted the fix/workspace-membership-lost-on-file-replace branch March 17, 2026 09:29
EnjoyBacon7 added a commit that referenced this pull request Mar 19, 2026
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
EnjoyBacon7 added a commit that referenced this pull request Apr 1, 2026
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
@Ahmath-Gadji Ahmath-Gadji added the fix Fix issue label Apr 15, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

fix Fix issue

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants