fix(workspaces): reject non-existent file IDs in add_files_to_workspace - #281
Conversation
📝 WalkthroughWalkthroughThis PR adds a foreign-key from Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant Router as add_files_to_workspace (Router)
participant VectorDB as MilvusDB
participant PartitionMgr as PartitionFileManager
participant DB as Database
Client->>Router: POST /partition/{partition}/workspaces/{ws}/files\n{"file_ids": ["f1","f2"]}
Router->>VectorDB: get_existing_file_ids(partition, file_ids)
VectorDB->>PartitionMgr: get_existing_file_ids(partition, file_ids)
PartitionMgr->>DB: SELECT files.file_id WHERE partition=? AND file_id IN (...)
DB-->>PartitionMgr: existing_ids
PartitionMgr-->>VectorDB: existing_ids
VectorDB-->>Router: existing_ids
alt missing files
Router->>Router: compute unknown_ids = requested - existing
Router-->>Client: 404 NOT_FOUND {"detail": "Files not found: [...]"}
else all exist
Router->>VectorDB: add_files_to_workspace(workspace_id, file_ids)
VectorDB->>PartitionMgr: add_files_to_workspace(workspace_id, file_ids)
PartitionMgr->>DB: resolve file_id -> files.id, INSERT WorkspaceFile(workspace_id, file_pk)
DB-->>PartitionMgr: success
PartitionMgr-->>VectorDB: OK
VectorDB-->>Router: OK
Router-->>Client: 200 OK {"status":"added"}
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested reviewers
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 docstrings
🧪 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: 5
🤖 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 204-209: The workspace restoration call using asyncio.gather over
vectordb.add_files_to_workspace.remote(...) should be wrapped in a try-except so
transient failures don't turn a successful metadata update into a 500; replace
the direct await asyncio.gather(...) in the update flow with a try block and on
Exception log a warning (similar to add_file()'s handling) including a message,
the exception string, and workspace_ids to indicate the file is indexed but
workspace links may be incomplete.
In `@openrag/components/indexer/vectordb/models.py`:
- Line 165: The foreign key on file_id in models.py points to a non-unique
column; fix by making the child table reference a unique key: either (A) add
partition_name to WorkspaceFile (ensure WorkspaceFile has a UNIQUE or PRIMARY
KEY on (file_id, partition_name)) and replace the single Column(String,
ForeignKey(...)) with a composite ForeignKeyConstraint referencing
("files.file_id","files.partition_name") and add a local partition_name Column,
or (B) change the foreign key to reference the globally-unique id column on
WorkspaceFile (use ForeignKey("files.id")) and update any relationship code that
expects file_id to use the id instead; update migrations accordingly.
In `@openrag/components/indexer/vectordb/test_workspace_file_validation.py`:
- Around line 50-72: The tests duplicate the SQL by implementing
PartitionFileManagerHelper.get_existing_file_ids instead of exercising the
production code; remove or replace the duplicated implementation so tests call
the real PartitionFileManager.get_existing_file_ids (or extract a shared helper
used by both) – e.g., construct a PartitionFileManager with the existing
session_factory/Session and call its get_existing_file_ids method (referencing
PartitionFileManager, get_existing_file_ids, and FileModel to locate the real
implementation), ensuring the test suite exercises the shipped code path rather
than a local copy.
In `@openrag/routers/indexer.py`:
- Around line 268-269: Replace the direct Ray actor call to
vectordb.get_file_workspaces.remote with the repository's timeout wrapper: use
call_ray_actor_with_timeout to invoke vectordb.get_file_workspaces so the call
respects actor timeouts; update the assignment to existing_workspace_ids to
receive the wrapped call's result and ensure you pass the same arguments
(file_id, partition) and any required timeout or context parameters to
call_ray_actor_with_timeout just like other usages in workspaces.py/search.py.
In
`@openrag/scripts/migrations/alembic/versions/f1a2b3c4d5e6_add_workspace_files_file_id_fk.py`:
- Around line 20-31: The upgrade currently deletes orphans and creates
fk_workspace_files_file_id only on file_id, ignoring partition scoping; modify
upgrade() to first add a partition_name column to workspace_files, populate it
by copying the partition_name from the owning workspace (join on workspace_id or
the appropriate workspace FK), ensure NOT NULL/default as required, then change
the DELETE to remove rows where (file_id, partition_name) NOT IN (SELECT
file_id, partition_name FROM files), and finally create a composite foreign key
on (file_id, partition_name) referencing files(file_id, partition_name) (replace
the single-column fk_workspace_files_file_id with a composite FK name), so both
the purge and FK use the composite key; ensure migrations handle existing nulls
and data migration ordering when adding the NOT NULL constraint.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 85a86995-da77-41b4-8a8b-3a69ccc85924
📒 Files selected for processing (8)
openrag/components/indexer/indexer.pyopenrag/components/indexer/vectordb/models.pyopenrag/components/indexer/vectordb/test_workspace_file_validation.pyopenrag/components/indexer/vectordb/utils.pyopenrag/components/indexer/vectordb/vectordb.pyopenrag/routers/indexer.pyopenrag/routers/workspaces.pyopenrag/scripts/migrations/alembic/versions/f1a2b3c4d5e6_add_workspace_files_file_id_fk.py
|
@CodeRabbit Review |
✅ Actions performedReview triggered.
|
add_files_to_workspace() would silently insert ghost workspace_files rows for file IDs that don't exist in the files table, because there was no validation before the upsert. Fix: add PartitionFileManager.get_existing_file_ids() (and expose it on MilvusDB) to query which IDs are actually present in the partition. The router now calls this before insertion and raises HTTP 404 listing the unknown IDs if any are found.
Add a proper foreign key constraint (ON DELETE CASCADE) so the database enforces referential integrity between workspace_files and files, preventing ghost rows from persisting even if the application check is somehow bypassed. The Alembic migration purges any existing orphaned rows before adding the constraint to avoid upgrade failures on instances that already have ghost data.
Cover the cases: all IDs exist, some missing, none exist, empty input, cross-partition isolation, and duplicate IDs in input.
…files.file_id (string) files.file_id is not unique — only (file_id, partition_name) is. PostgreSQL rejects a FK on a non-unique column, and the string-based approach cannot enforce partition boundaries. Changes: - WorkspaceFile.file_id: String FK → Integer FK referencing files.id - utils.py: all workspace-file methods now resolve File.id via a JOIN before inserting/querying WorkspaceFile, keeping the public string file_id interface intact - Migration f1a2b3c4d5e6 rewritten: adds a temp int column, populates via JOIN to files.id, drops old string column, renames, adds FK - indexer.py: wrap workspace restore asyncio.gather in try-except so a transient failure doesn't turn a successful metadata update into a 500 - routers/indexer.py: wrap get_file_workspaces.remote with call_ray_actor_with_timeout to respect actor timeout policy - Drop test_workspace_file_validation.py (duplicated production SQL); coverage to move to integration tests per @paultranvan
…ty, and remove_file_from_workspace cross-product - delete_workspace: remove dead 'not in files table' check — with the int FK every workspace file has a backing files row, so orphan = only in this workspace and no other - migration: scope the UPDATE JOIN through workspaces to resolve the correct partition when the same filename exists in multiple partitions; purge unresolvable rows after the UPDATE - migration: recreate uix_workspace_file unique constraint after column drop/rename (PostgreSQL drops it with the column) - remove_file_from_workspace: fix cross-product JOIN that would fail with scalar_one_or_none when multiple workspaces share a partition; resolve partition via the specific workspace row instead
a8eb45f to
8a41e71
Compare
- test_workspaces.py: upload files before adding them to workspaces, since the router now validates file existence via get_existing_file_ids - delete_workspace: remove dead 'not in files' check (int FK guarantees every workspace file has a backing files row) - migration: scope UPDATE JOIN through workspaces table to resolve the correct partition; purge unresolvable rows; recreate uix_workspace_file unique constraint after column drop/rename - remove_file_from_workspace: fix cross-product JOIN by resolving partition via the specific workspace row
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
openrag/components/indexer/vectordb/utils.py (1)
636-659:⚠️ Potential issue | 🔴 CriticalDon't equate “not in another workspace” with “safe to delete the file.”
openrag/routers/workspaces.py, Lines 119-147 deletes everyfile_idreturned here from the partition. With this definition, any file linked only to the workspace being deleted is treated as orphaned and removed, even though it can still be a valid indexed/searchable partition file. That makes workspace deletion a data-loss path unless file ownership/provenance is tracked separately.Based on learnings: "workspace association is a secondary/best-effort metadata step" and a file remains correctly indexed and searchable in the partition even if that association fails.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@openrag/components/indexer/vectordb/utils.py` around lines 636 - 659, The delete_workspace function currently identifies files "orphaned" by absence in other WorkspaceFile rows and the router then deletes those files from the partition; instead, change delete_workspace (in openrag/components/indexer/vectordb/utils.py) to only remove the Workspace and its WorkspaceFile rows and return the candidate file_ids but do NOT treat absence-from-other-workspaces as permission to delete the underlying File/partition data; update callers (e.g. the workspace delete handler in routers/workspaces.py) to require an explicit safety check before deleting partition files — either check a File ownership/provenance/managed flag (add such a column on File if needed) or require a force_delete_files parameter/confirmation — so only files with explicit managed/owned provenance are removed. Ensure references to WorkspaceFile, File, delete_workspace, and the router delete handler are updated accordingly.
🤖 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/vectordb/utils.py`:
- Around line 681-689: The loop currently skips missing File.id mappings
(file_row is None) which yields partial success; instead, resolve the full
mapping for all file_ids in this DB session and fail if any are missing: query
all File.id for the given file_ids and partition (reuse the
select(File.id).where(File.file_id == fid, File.partition_name == partition)
logic but do it as a bulk query), build a map of file_id->id, compute
missing_ids = [fid for fid in file_ids if fid not in map], and if missing_ids is
non-empty raise/return an error listing them rather than continuing; once all
IDs are resolved, perform the pg_insert into WorkspaceFile (values workspace_id
and resolved file_id) with
on_conflict_do_nothing(constraint="uix_workspace_file") as before.
---
Outside diff comments:
In `@openrag/components/indexer/vectordb/utils.py`:
- Around line 636-659: The delete_workspace function currently identifies files
"orphaned" by absence in other WorkspaceFile rows and the router then deletes
those files from the partition; instead, change delete_workspace (in
openrag/components/indexer/vectordb/utils.py) to only remove the Workspace and
its WorkspaceFile rows and return the candidate file_ids but do NOT treat
absence-from-other-workspaces as permission to delete the underlying
File/partition data; update callers (e.g. the workspace delete handler in
routers/workspaces.py) to require an explicit safety check before deleting
partition files — either check a File ownership/provenance/managed flag (add
such a column on File if needed) or require a force_delete_files
parameter/confirmation — so only files with explicit managed/owned provenance
are removed. Ensure references to WorkspaceFile, File, delete_workspace, and the
router delete handler are updated accordingly.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 2df33fe9-b14f-41a5-989e-22df32aeb27b
📒 Files selected for processing (7)
openrag/components/indexer/indexer.pyopenrag/components/indexer/vectordb/models.pyopenrag/components/indexer/vectordb/utils.pyopenrag/components/indexer/vectordb/vectordb.pyopenrag/routers/indexer.pyopenrag/routers/workspaces.pyopenrag/scripts/migrations/alembic/versions/f1a2b3c4d5e6_add_workspace_files_file_id_fk.py
🚧 Files skipped from review as they are similar to previous changes (4)
- openrag/components/indexer/indexer.py
- openrag/components/indexer/vectordb/models.py
- openrag/routers/indexer.py
- openrag/components/indexer/vectordb/vectordb.py
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/api_tests/test_workspaces.py (1)
79-89: Consider unifying upload helpers to avoid drift.This helper duplicates logic already present in
TestFileWorkspaceMembership._upload_file, and the two implementations behave differently (one waits for indexing internally, the other does not). Consolidating into one shared helper/fixture will keep test behavior consistent.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/api_tests/test_workspaces.py` around lines 79 - 89, The test helper _upload_file duplicates logic from TestFileWorkspaceMembership._upload_file and inconsistently calls wait_for_indexing; unify them by extracting a single shared helper/fixture used by both tests (e.g., a module-level function or pytest fixture named upload_file) that performs the API POST to "/indexer/partition/{partition}/file/{file_id}" with the same files/data shape and documents whether it waits for indexing; update callers to use the unified helper and ensure wait_for_indexing is invoked consistently (or controlled via a parameter) so both usages of _upload_file and TestFileWorkspaceMembership._upload_file behave identically.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@tests/api_tests/test_workspaces.py`:
- Around line 79-89: The test helper _upload_file duplicates logic from
TestFileWorkspaceMembership._upload_file and inconsistently calls
wait_for_indexing; unify them by extracting a single shared helper/fixture used
by both tests (e.g., a module-level function or pytest fixture named
upload_file) that performs the API POST to
"/indexer/partition/{partition}/file/{file_id}" with the same files/data shape
and documents whether it waits for indexing; update callers to use the unified
helper and ensure wait_for_indexing is invoked consistently (or controlled via a
parameter) so both usages of _upload_file and
TestFileWorkspaceMembership._upload_file behave identically.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 59e1670d-615d-4793-811b-ad609d23fdb4
📒 Files selected for processing (1)
tests/api_tests/test_workspaces.py
Summary
Fixes #278.
POST /partition/{partition}/workspaces/{workspace_id}/filesaccepted arbitraryfile_idsand silently inserted ghostworkspace_filesrows when the referenced files did not exist in thefilestable, because there was no validation before the upsert and no foreign key constraint.Changes
Application-level validation (router + utils)
vectordb/utils.py— addPartitionFileManager.get_existing_file_ids(partition, file_ids) -> set[str]: queriesfilesfor which IDs actually belong to the partitionvectordb/vectordb.py— expose it asasync MilvusDB.get_existing_file_ids(...)(Ray-callable)routers/workspaces.py—add_files_to_workspacenow checks unknown IDs before insertion and returnsHTTP 404listing them:"File IDs not found in partition 'p': ['<id>']"Database-level constraint
models.py—WorkspaceFile.file_idnow has aForeignKey("files.file_id", ondelete="CASCADE")constraint so the DB enforces referential integrity even if the app check is bypassedf1a2b3c4d5e6— purges orphaned rows then adds the FK constraint; downgrade drops itTests
test_workspace_file_validation.pycovering: all IDs exist, some missing, none exist, empty input, cross-partition isolation, duplicate IDs in inputTest results
212 passed, 3 skipped (pre-existing skip — unrelated to this change)
Summary by CodeRabbit
New Features
Bug Fixes
Chores