Skip to content

fix(workspaces): reject non-existent file IDs in add_files_to_workspace - #281

Merged
EnjoyBacon7 merged 6 commits into
devfrom
fix/workspace-ghost-file-ids
Mar 19, 2026
Merged

fix(workspaces): reject non-existent file IDs in add_files_to_workspace#281
EnjoyBacon7 merged 6 commits into
devfrom
fix/workspace-ghost-file-ids

Conversation

@EnjoyBacon7

@EnjoyBacon7 EnjoyBacon7 commented Mar 12, 2026

Copy link
Copy Markdown
Collaborator

Summary

Fixes #278.

POST /partition/{partition}/workspaces/{workspace_id}/files accepted arbitrary file_ids and silently inserted ghost workspace_files rows when the referenced files did not exist in the files table, because there was no validation before the upsert and no foreign key constraint.

Changes

Application-level validation (router + utils)

  • vectordb/utils.py — add PartitionFileManager.get_existing_file_ids(partition, file_ids) -> set[str]: queries files for which IDs actually belong to the partition
  • vectordb/vectordb.py — expose it as async MilvusDB.get_existing_file_ids(...) (Ray-callable)
  • routers/workspaces.pyadd_files_to_workspace now checks unknown IDs before insertion and returns HTTP 404 listing them: "File IDs not found in partition 'p': ['<id>']"

Database-level constraint

  • models.pyWorkspaceFile.file_id now has a ForeignKey("files.file_id", ondelete="CASCADE") constraint so the DB enforces referential integrity even if the app check is bypassed
  • Migration f1a2b3c4d5e6 — purges orphaned rows then adds the FK constraint; downgrade drops it

Tests

  • 7 new unit tests in test_workspace_file_validation.py covering: all IDs exist, some missing, none exist, empty input, cross-partition isolation, duplicate IDs in input

Test results

212 passed, 3 skipped (pre-existing skip — unrelated to this change)

Summary by CodeRabbit

  • New Features

    • Added file existence validation when adding files to workspaces; missing files return a 404 with details.
    • Endpoint now requires a partition parameter when adding files to a workspace.
  • Bug Fixes

    • Improved error handling so workspace association restoration during re-indexing won’t fail the operation.
    • Applied timeout protection to external indexing calls to improve reliability.
  • Chores

    • Database migration converting workspace-file links to a foreign-key-backed relationship for stronger integrity.

@coderabbitai

coderabbitai Bot commented Mar 12, 2026

Copy link
Copy Markdown
📝 Walkthrough

Walkthrough

This PR adds a foreign-key from workspace_files.file_idfiles.id, migrates existing data, refactors workspace-file utils to use File.id, adds partition-scoped validation in the workspace router to reject unknown file IDs, and makes workspace restoration during re-indexing best-effort (exceptions logged, not raised).

Changes

Cohort / File(s) Summary
Schema & Migration
openrag/components/indexer/vectordb/models.py, openrag/scripts/migrations/alembic/versions/f1a2b3c4d5e6_add_workspace_files_file_id_fk.py
Change WorkspaceFile.file_id from string to Integer FK → files.id with ON DELETE CASCADE. Migration converts data, removes unresolved rows, recreates index/unique constraint, and provides downgrade.
Partition-aware File Manager
openrag/components/indexer/vectordb/utils.py
Refactor methods to map external file_idFile.id within a partition. Add get_existing_file_ids(). Update add/remove/list/get_file_workspaces/remove_all/delete_workspace to use joins and FK semantics.
VectorDB Adapter
openrag/components/indexer/vectordb/vectordb.py
Add async get_existing_file_ids() and get_file_workspaces() delegating to PartitionFileManager.
API Routers
openrag/routers/workspaces.py, openrag/routers/indexer.py
add_files_to_workspace now requires partition, pre-validates file IDs via vectordb.get_existing_file_ids() and returns 404 for unknown IDs. put_file workspace snapshot uses call_ray_actor_with_timeout.
Indexer Error Handling
openrag/components/indexer/indexer.py
Workspace membership restoration after re-indexing changed from unguarded asyncio.gather(...) to a guarded try/except; failures are logged as warnings and do not abort indexing.
Tests
tests/api_tests/test_workspaces.py
Add shared helper to upload files before workspace operations; update tests to upload files prior to add/list/remove to reflect partition-scoped validation.

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
Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Possibly related PRs

Suggested reviewers

  • paultranvan

Poem

🐰
I hopped through schemas late at night,
Swapped strings for ints to make things right,
I checked each file within its partition bright,
Ghosts were chased and APIs now bite,
A tidy hop — indexed, linked, and tight.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 31.25% 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 PR title clearly and specifically summarizes the main change: adding validation to reject non-existent file IDs in add_files_to_workspace, which is the primary objective addressed by the changeset.
Linked Issues check ✅ Passed All coding objectives from issue #278 are met: application-level validation added to check file existence [workspaces.py], database-level FK constraint enforcing referential integrity [models.py, migration], and comprehensive error reporting for unknown IDs.
Out of Scope Changes check ✅ Passed All changes are directly scoped to issue #278: validation logic, FK constraint, migration, error handling, and supporting test infrastructure. No unrelated modifications detected.

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

✨ Finishing Touches
📝 Generate docstrings
  • Create stacked PR
  • Commit on current branch
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch fix/workspace-ghost-file-ids
📝 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: 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

📥 Commits

Reviewing files that changed from the base of the PR and between c1fe53a and 096bc07.

📒 Files selected for processing (8)
  • openrag/components/indexer/indexer.py
  • openrag/components/indexer/vectordb/models.py
  • openrag/components/indexer/vectordb/test_workspace_file_validation.py
  • openrag/components/indexer/vectordb/utils.py
  • openrag/components/indexer/vectordb/vectordb.py
  • openrag/routers/indexer.py
  • openrag/routers/workspaces.py
  • openrag/scripts/migrations/alembic/versions/f1a2b3c4d5e6_add_workspace_files_file_id_fk.py

Comment thread openrag/components/indexer/indexer.py Outdated
Comment thread openrag/components/indexer/vectordb/models.py Outdated
Comment thread openrag/components/indexer/vectordb/test_workspace_file_validation.py Outdated
Comment thread openrag/routers/indexer.py Outdated
Comment thread openrag/components/indexer/vectordb/test_workspace_file_validation.py Outdated
Comment thread openrag/components/indexer/vectordb/models.py Outdated
Comment thread openrag/components/indexer/vectordb/models.py Outdated
@EnjoyBacon7

Copy link
Copy Markdown
Collaborator Author

@CodeRabbit Review

@coderabbitai

coderabbitai Bot commented Mar 19, 2026

Copy link
Copy Markdown
✅ Actions performed

Review triggered.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.

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
@EnjoyBacon7
EnjoyBacon7 force-pushed the fix/workspace-ghost-file-ids branch from a8eb45f to 8a41e71 Compare March 19, 2026 12:19
@coderabbitai coderabbitai Bot added the breaking-change Change of behavior after upgrade label Mar 19, 2026
- 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

@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

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 | 🔴 Critical

Don't equate “not in another workspace” with “safe to delete the file.”

openrag/routers/workspaces.py, Lines 119-147 deletes every file_id returned 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

📥 Commits

Reviewing files that changed from the base of the PR and between 096bc07 and 8a41e71.

📒 Files selected for processing (7)
  • openrag/components/indexer/indexer.py
  • openrag/components/indexer/vectordb/models.py
  • openrag/components/indexer/vectordb/utils.py
  • openrag/components/indexer/vectordb/vectordb.py
  • openrag/routers/indexer.py
  • openrag/routers/workspaces.py
  • openrag/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

Comment thread openrag/components/indexer/vectordb/utils.py

@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.

🧹 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

📥 Commits

Reviewing files that changed from the base of the PR and between 8a41e71 and f6d6456.

📒 Files selected for processing (1)
  • tests/api_tests/test_workspaces.py

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