Feat/file quota - #245
Conversation
Ahmath-Gadji
added
breaking-change
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds a per-user file quota system: DB schema (users.file_quota, users.file_count, files.created_by), Alembic migrations and merge, DEFAULT_FILE_QUOTA config/env, quota enforcement dependency applied to upload endpoints, API for viewing/updating quotas, vectordb/ORM changes to track counts, CLI checker, and tests. Changes
Sequence DiagramsequenceDiagram
participant Client as Client API
participant Router as Router (add_file)
participant QuotaCheck as check_user_file_quota
participant VectorDB as VectorDB / PartitionFileManager
participant TaskMgr as TaskStateManager
participant DB as Database
Client->>Router: POST /add_file
Router->>QuotaCheck: Depends(check_user_file_quota)
QuotaCheck->>VectorDB: get_user_by_id(user_id)
VectorDB->>DB: query user.file_quota, user.file_count
DB-->>VectorDB: user data
VectorDB-->>QuotaCheck: user data
QuotaCheck->>TaskMgr: get_user_pending_task_count(user_id)
TaskMgr-->>QuotaCheck: pending_count
QuotaCheck->>QuotaCheck: total = file_count + pending_count
alt total < file_quota (or unlimited)
QuotaCheck-->>Router: allow
Router->>VectorDB: add_file_to_partition(..., user_id)
VectorDB->>DB: insert file (created_by), increment file_count
DB-->>VectorDB: success
VectorDB-->>Router: file added
Router-->>Client: 200 OK
else total >= file_quota
QuotaCheck-->>Router: deny (403)
Router-->>Client: 403 Quota exceeded
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Tip Issue Planner is now in beta. Read the docs and try it out! Share your feedback on Discord. 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
openrag/components/indexer/vectordb/vectordb.py (1)
886-893:⚠️ Potential issue | 🔴 Critical
delete_usercallsdelete_partitionwithout the now-requireduser_id, causing aTypeErrorat runtime.
delete_partitioninPartitionFileManager(utils.py line 302) now requiresuser_id: int, but line 892 calls it without that argument. Any attempt to delete a user who owns partitions will crash.Proposed fix
async def delete_user(self, user_id: int): self._check_user_exists(user_id) user_partitions = [ p["partition"] for p in self.partition_file_manager.list_user_partitions(user_id) if p["role"] == "owner" ] for partition in user_partitions: - self.partition_file_manager.delete_partition(partition) + self.partition_file_manager.delete_partition(partition, user_id=user_id) self.partition_file_manager.delete_user(user_id)openrag/components/indexer/indexer.py (1)
164-189:⚠️ Potential issue | 🟠 Major
user=Nonewill crashdelete_filewhich requiresuser: dict.
update_file_metadatadeclaresuser: dict | None = None(line 170), but callsself.delete_fileat line 183 expectinguser: dict(non-optional, required parameter at line 149). Ifupdate_file_metadatais ever called without auserargument,delete_filewill invokeuser.get("id")onNone, raising anAttributeError.The current router call site always passes
user, but the type signature creates a contract violation that could surface in future usage.Either make
userrequired inupdate_file_metadataor add a guard before callingdelete_file.Proposed fix — make user required to match delete_file
`@ray.method`(concurrency_group="update") async def update_file_metadata( self, file_id: str, metadata: dict, partition: str, - user: dict | None = None, + user: dict, ):
🤖 Fix all issues with AI agents
In `@docs/content/docs/documentation/data_model.md`:
- Around line 67-71: The markdown table rows for `file_quota` and `file_count`
are malformed (inconsistent pipe spacing) and there's no blank line before the
`**Relationships**` header which causes MD table parsing errors; fix by
normalizing the table row formatting for the `file_quota` and `file_count` lines
so pipes and cell spacing are consistent, then insert a single blank line after
the end of the table and before the `**Relationships**` line so
`**Relationships**` is parsed as a normal heading/paragraph rather than a table
row.
In `@openrag/routers/partition.py`:
- Around line 71-72: The base abstract method delete_partition currently has the
signature async def delete_partition(self, partition: str) but the concrete
implementation expects async def delete_partition(self, partition: str, user_id:
int); update the abstract declaration in the vectordb base class to async def
delete_partition(self, partition: str, user_id: int) (include the user_id type
hint), adjust any abstractmethod decorator if present, and ensure all concrete
subclasses (e.g., the implementation that defines delete_partition(self,
partition: str, user_id: int)) continue to match this signature so the interface
is consistent.
In `@openrag/routers/users.py`:
- Around line 59-95: get_current_user_info currently treats DEFAULT_FILE_QUOTA <
0 as the only signal for "unlimited", which causes inconsistent UX when
DEFAULT_FILE_QUOTA == 0 because check_user_file_quota uses DEFAULT_FILE_QUOTA <=
0 to disable quota enforcement; update get_current_user_info (and its logic
around user_quota assignment) to use DEFAULT_FILE_QUOTA <= 0 when deciding to
set user_quota = float("inf") so the /info response reports unlimited
consistently with check_user_file_quota (refer to function get_current_user_info
and the global DEFAULT_FILE_QUOTA).
In `@openrag/routers/utils.py`:
- Around line 184-187: The check_user_file_quota function currently declares
vectordb=Depends(get_vectordb) but never uses it, causing unnecessary overhead;
remove the unused dependency by deleting the vectordb parameter from the
check_user_file_quota signature (and any corresponding parameter in
callers/route declarations if applicable) so only user=Depends(current_user)
remains, or alternatively if vectordb is intended to be used implement its usage
inside check_user_file_quota (referencing get_vectordb) — ensure to update any
imports/usages accordingly and run tests to confirm no broken dependency
injections.
🧹 Nitpick comments (3)
tests/api_tests/test_indexer.py (1)
486-518: Good enforcement test — verify the 403 error detail message too.The test correctly validates that the 6th upload is rejected with 403. Consider also asserting on the response body to confirm it's a quota-related rejection (not some other 403 like a permissions issue).
💡 Optional enhancement
# 6th file should be rejected due to quota response = self._upload_file(api_client, partition_name, "file-5", user_token, "Content 5") assert response.status_code == 403, ( f"6th file should be rejected with 403 Forbidden, got {response.status_code}: {response.text}" ) + assert "quota" in response.json().get("detail", "").lower(), ( + f"403 should be quota-related, got: {response.text}" + )openrag/routers/users.py (1)
262-277: Consider returning the updated user quota details instead of just a message.
vectordb.update_user_quotareturns a dict with the updated user details (includingfile_quotaandfile_count), but the response only includes a plain text message. Returning the actual updated state would let clients confirm the change without a follow-up GET request.Proposed change
await vectordb.update_user_quota.remote(user_id, file_quota) - + updated_user = await vectordb.update_user_quota.remote(user_id, file_quota) logger.debug("Updated user quota", user_id=user_id, file_quota=file_quota) return JSONResponse( status_code=status.HTTP_200_OK, - content={"message": f"Quota for user {user_id} updated to {file_quota}"}, + content=updated_user, )openrag/components/indexer/vectordb/vectordb.py (1)
52-62: Abstract base class signatures are out of sync withMilvusDBimplementations.
BaseVectorDB.delete_file(line 61) andBaseVectorDB.delete_partition(line 53) lack theuser_idparameter that the concreteMilvusDBimplementations now require (lines 581, 760). Any code relying on the abstract interface would call these methods with the wrong arity.Proposed fix
`@abstractmethod` - async def delete_partition(self, partition: str): + async def delete_partition(self, partition: str, user_id: int): pass ... `@abstractmethod` - async def delete_file(self, file_id: str, partition: str): + async def delete_file(self, file_id: str, partition: str, user_id: int): pass
| | `file_quota` | Integer (nullable) | Max files allowed for that user | | ||
| | `file_count`| Integer (default=0) | Number of uploaded files.| | ||
| **Relationships** | ||
| - `memberships`: one-to-many → `PartitionMembership` | ||
|
|
There was a problem hiding this comment.
Missing blank line before **Relationships** breaks markdown table rendering.
Static analysis (MD055, MD056, MD058) correctly flags that line 69 (**Relationships**) is being parsed as a malformed table row because there's no blank line separating it from the preceding table. Also, line 68 has inconsistent pipe spacing.
📝 Proposed fix
-| `file_count`| Integer (default=0) | Number of uploaded files.|
-**Relationships**
- - `memberships`: one-to-many → `PartitionMembership`
+| `file_count` | Integer (default=0) | Number of uploaded files. |
+
+**Relationships**
+- `memberships`: one-to-many → `PartitionMembership`🧰 Tools
🪛 markdownlint-cli2 (0.20.0)
[warning] 69-69: Table pipe style
Expected: leading_and_trailing; Actual: no_leading_or_trailing; Missing leading pipe
(MD055, table-pipe-style)
[warning] 69-69: Table pipe style
Expected: leading_and_trailing; Actual: no_leading_or_trailing; Missing trailing pipe
(MD055, table-pipe-style)
[warning] 69-69: Table column count
Expected: 3; Actual: 1; Too few cells, row will be missing data
(MD056, table-column-count)
[warning] 69-69: Tables should be surrounded by blank lines
(MD058, blanks-around-tables)
🤖 Prompt for AI Agents
In `@docs/content/docs/documentation/data_model.md` around lines 67 - 71, The
markdown table rows for `file_quota` and `file_count` are malformed
(inconsistent pipe spacing) and there's no blank line before the
`**Relationships**` header which causes MD table parsing errors; fix by
normalizing the table row formatting for the `file_quota` and `file_count` lines
so pipes and cell spacing are consistent, then insert a single blank line after
the end of the table and before the `**Relationships**` line so
`**Relationships**` is parsed as a normal heading/paragraph rather than a table
row.
| async def check_user_file_quota( | ||
| user=Depends(current_user), | ||
| vectordb=Depends(get_vectordb), | ||
| ): |
There was a problem hiding this comment.
🛠️ Refactor suggestion | 🟠 Major
vectordb dependency is injected but never used in check_user_file_quota.
The vectordb=Depends(get_vectordb) parameter at line 186 is never referenced in the function body. This adds unnecessary overhead on every quota-checked request (Ray actor lookup) with no benefit.
♻️ Proposed fix
async def check_user_file_quota(
user=Depends(current_user),
- vectordb=Depends(get_vectordb),
):🤖 Prompt for AI Agents
In `@openrag/routers/utils.py` around lines 184 - 187, The check_user_file_quota
function currently declares vectordb=Depends(get_vectordb) but never uses it,
causing unnecessary overhead; remove the unused dependency by deleting the
vectordb parameter from the check_user_file_quota signature (and any
corresponding parameter in callers/route declarations if applicable) so only
user=Depends(current_user) remains, or alternatively if vectordb is intended to
be used implement its usage inside check_user_file_quota (referencing
get_vectordb) — ensure to update any imports/usages accordingly and run tests to
confirm no broken dependency injections.
|
|
||
| # Increment uploader's file_count | ||
| if user_id: | ||
| session.query(User).filter(User.id == user_id).update( |
There was a problem hiding this comment.
Nice enhancement with the argument synchronize_session=False that allows fast updates.
But we are still doing the updates manually. If i'm not mistaken this is not the same as using SQL triggers.
| user.file_count = func.greatest(0, User.file_count - file_count) | ||
|
|
||
| session.delete(partition_obj) # Will delete all files due to cascade | ||
| # Count files per uploader before cascade deletes them |
There was a problem hiding this comment.
I did miss this case in my implementation: discounting per user when a partition is deleted.
It would be nice to have a test that goes along with it
|
I've tested the migration script and got this error: INFO [alembic.runtime.migration] Will assume transactional DDL.
ERROR [alembic.util.messaging] Multiple head revisions are present for given argument 'head'; please specify a specific target revision, '<branchname>@head' to narrow to a specific head, or 'heads' for all heads
FAILED: Multiple head revisions are present for given argument 'head'; please specify a specific target revision, '<branchname>@head'
to narrow to a specific head, or 'heads' for all heads
In fact, theses features relationships and file_quota have been developed independently. I've created a new migration script that merges them. See my commit |
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Fix all issues with AI agents
In `@docs/content/docs/documentation/sql_migration.mdx`:
- Around line 90-94: The migration example is missing the database startup step
so applying migrations can fail; update the "Apply migrations after merge"
snippet to run the rdb service first by adding a `docker compose up -d rdb` step
before the `docker compose run --no-deps --build --rm --entrypoint "uv run
alembic -c /app/openrag/scripts/migrations/alembic/alembic.ini upgrade head"
openrag; docker compose down` command so the rdb is available for Alembic to
apply migrations.
In `@openrag/components/indexer/vectordb/utils.py`:
- Around line 582-603: The update_user_quota method lacks a null-guard for the
queried User which can cause an AttributeError if the user was deleted after a
prior existence check; inside update_user_quota (using self.Session and querying
User), verify that user is not None after s.query(...).first(), and if it is
None log a clear message via self.logger and raise or return a controlled "user
not found" error (e.g., raise a ValueError or a NotFoundError) instead of
proceeding to access attributes; ensure no commit/refresh is attempted when user
is None so callers receive a deterministic error rather than an AttributeError.
🧹 Nitpick comments (3)
openrag/components/indexer/vectordb/vectordb.py (1)
871-878: Consider using keyword arguments when forwarding topartition_file_manager.create_user.Line 878 forwards all parameters positionally. If the signature of
PartitionFileManager.create_userever changes parameter order, this will silently break.♻️ Suggested improvement
- return self.partition_file_manager.create_user(display_name, external_user_id, is_admin, file_quota) + return self.partition_file_manager.create_user( + display_name=display_name, + external_user_id=external_user_id, + is_admin=is_admin, + file_quota=file_quota, + )tests/api_tests/test_indexer.py (2)
409-456: Good test scaffolding with proper cleanup.The helper methods and try/finally cleanup pattern are well-structured. Swallowing exceptions in cleanup helpers is appropriate for test teardown.
One minor nit:
import ioon Line 432 could be moved to the file-level imports since it's also used intest_file_count_stable_on_replace(Line 672).
670-714: Good test for file replacement stability.Verifies that PUT (replace) doesn't double-count by checking
file_countremains 1 after replacement.Consider adding a test for
file_quota=0(blocks all uploads) as an edge case — currently onlyfile_quota=-1(unlimited) andfile_quota=5(bounded) are covered.
1866dc3 to
27d594f
Compare
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Fix all issues with AI agents
In `@openrag/scripts/check_file_counts.py`:
- Line 8: The import list includes an unused symbol 'text' from sqlalchemy which
is causing lint failures; remove 'text' from the import statement so it reads
only "from sqlalchemy import create_engine, func" (ensuring any usages of 'text'
elsewhere are not present or are handled), then run linters to confirm the
pipeline passes.
- Around line 16-23: The build_database_url function constructs a plain-text
PostgreSQL URL using raw user and password values which breaks when credentials
contain special characters; update build_database_url to percent-encode
credentials by importing urllib.parse.quote_plus and applying it to
args.user/user and args.password/password (encode both user and password before
interpolating into the returned f-string) so the generated
"postgresql://{user}:{password}@{host}:{port}/{db_name}" URL is safe for special
characters.
In `@tests/api_tests/test_indexer.py`:
- Around line 703-705: In test_indexer.py adjust the partition delete to use the
same user auth as the other test: replace the admin/default api_client.delete
call that deletes partition_a with a call that includes
headers={"Authorization": f"Bearer {user_token}"}, mirroring the behavior in
test_file_count_decrements_on_delete so the code path exercised is the
uploader-authenticated delete (look for the delete call around partition_a in
the test and the user_token variable used earlier).
🧹 Nitpick comments (2)
openrag/scripts/check_file_counts.py (2)
11-13: Consider using absolute imports per coding guidelines.The
sys.path.inserthack is fragile. If this script is meant to be run from the project root (or installed as part of the package), an absolute import would be more robust and consistent with the project's import conventions.-# Add parent dirs so we can import the models -sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) -from components.indexer.vectordb.utils import File, User +from openrag.components.indexer.vectordb.utils import File, UserAs per coding guidelines: "Use absolute imports from the
openrag/directory as the Python path root, avoiding relative imports across packages."
26-28: Engine is never disposed.For a CLI script this is unlikely to cause issues, but it's good practice to dispose the engine to cleanly release connections, especially if the script is ever imported and called programmatically.
Proposed fix
def check_file_counts(database_url, fix=False): engine = create_engine(database_url) - Session = sessionmaker(bind=engine) - - with Session() as session: - ... - - return 0 if not has_mismatch else 1 + try: + Session = sessionmaker(bind=engine) + with Session() as session: + ... + return 0 if not has_mismatch else 1 + finally: + engine.dispose()
| def build_database_url(args): | ||
| host = args.host or os.environ.get("POSTGRES_HOST", "localhost") | ||
| port = args.port or os.environ.get("POSTGRES_PORT", "5432") | ||
| user = args.user or os.environ.get("POSTGRES_USER", "root") | ||
| password = args.password or os.environ.get("POSTGRES_PASSWORD", "root_password") | ||
| collection = args.collection or os.environ.get("VDB_COLLECTION_NAME", "vdb_test") | ||
| db_name = f"partitions_for_collection_{collection}" | ||
| return f"postgresql://{user}:{password}@{host}:{port}/{db_name}" |
There was a problem hiding this comment.
Password with special characters will break the connection URL.
If the PostgreSQL password contains characters like @, /, %, or :, the f-string interpolation on line 23 will produce a malformed URL. Use urllib.parse.quote_plus to percent-encode the credentials.
Proposed fix
+from urllib.parse import quote_plus
+
def build_database_url(args):
host = args.host or os.environ.get("POSTGRES_HOST", "localhost")
port = args.port or os.environ.get("POSTGRES_PORT", "5432")
user = args.user or os.environ.get("POSTGRES_USER", "root")
password = args.password or os.environ.get("POSTGRES_PASSWORD", "root_password")
collection = args.collection or os.environ.get("VDB_COLLECTION_NAME", "vdb_test")
db_name = f"partitions_for_collection_{collection}"
- return f"postgresql://{user}:{password}@{host}:{port}/{db_name}"
+ return f"postgresql://{quote_plus(user)}:{quote_plus(password)}@{host}:{port}/{db_name}"📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def build_database_url(args): | |
| host = args.host or os.environ.get("POSTGRES_HOST", "localhost") | |
| port = args.port or os.environ.get("POSTGRES_PORT", "5432") | |
| user = args.user or os.environ.get("POSTGRES_USER", "root") | |
| password = args.password or os.environ.get("POSTGRES_PASSWORD", "root_password") | |
| collection = args.collection or os.environ.get("VDB_COLLECTION_NAME", "vdb_test") | |
| db_name = f"partitions_for_collection_{collection}" | |
| return f"postgresql://{user}:{password}@{host}:{port}/{db_name}" | |
| from urllib.parse import quote_plus | |
| def build_database_url(args): | |
| host = args.host or os.environ.get("POSTGRES_HOST", "localhost") | |
| port = args.port or os.environ.get("POSTGRES_PORT", "5432") | |
| user = args.user or os.environ.get("POSTGRES_USER", "root") | |
| password = args.password or os.environ.get("POSTGRES_PASSWORD", "root_password") | |
| collection = args.collection or os.environ.get("VDB_COLLECTION_NAME", "vdb_test") | |
| db_name = f"partitions_for_collection_{collection}" | |
| return f"postgresql://{quote_plus(user)}:{quote_plus(password)}@{host}:{port}/{db_name}" |
🤖 Prompt for AI Agents
In `@openrag/scripts/check_file_counts.py` around lines 16 - 23, The
build_database_url function constructs a plain-text PostgreSQL URL using raw
user and password values which breaks when credentials contain special
characters; update build_database_url to percent-encode credentials by importing
urllib.parse.quote_plus and applying it to args.user/user and
args.password/password (encode both user and password before interpolating into
the returned f-string) so the generated
"postgresql://{user}:{password}@{host}:{port}/{db_name}" URL is safe for special
characters.
| self._cleanup_user(api_client, editor["id"]) | ||
| self._cleanup_user(api_client, owner["id"]) | ||
|
|
||
| def test_partition_delete_only_decrements_its_files(self, api_client): |
There was a problem hiding this comment.
It's a nice test but in my comment the idea was more about testing the discounting per user: #245 (comment)
- Test
- user A uploads a 2 files. Likewise for user B.
- Then we delete the partition and check the discounting per user.
There was a problem hiding this comment.
Well that was discounting a user after a partition removal :)
But I see your point, I updated the test
| return f"postgresql://{user}:{password}@{host}:{port}/{db_name}" | ||
|
|
||
|
|
||
| def check_file_counts(database_url, fix=False): |
There was a problem hiding this comment.
I guess this is supposed to be scheduled to run periodically. I think that's the missing piece.
There was a problem hiding this comment.
yes, but not by the script itself, for now let's keep it manual
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@tests/api_tests/test_indexer.py`:
- Around line 473-476: The test blocks currently only wait when
"task_status_url" is present, which misses responses that include only
"task_id"; update each upload handling block to always attempt to extract a task
id via get_task_id(response.json()) and then call wait_for_task(api_client,
task_id, headers={"Authorization": f"Bearer {user_token}"}), or wrap get_task_id
in a try/except to skip waiting only if no task info is present—apply this
change to every occurrence in the test class where the current code checks for
"task_status_url" so uploads that return {"task_id": "..."} are properly
awaited.
🧹 Nitpick comments (2)
tests/api_tests/test_indexer.py (2)
430-442: Moveimport ioto the top of the file.
iois used in_upload_file(line 432) and again intest_file_count_stable_on_replace(line 728). Hoist it to the module-level imports alongsideos,time, andPathfor consistency.Proposed fix
At the top of the file:
import os import time +import io from pathlib import PathThen remove the local imports at lines 432 and 728.
458-458: Remove unusedtmp_pathfixture parameter.Both
test_unlimited_quota_user_can_uploadandtest_quota_limit_blocks_excess_uploadsaccepttmp_pathbut never use it. This creates unnecessary temp directories on each run.Proposed fix
- def test_unlimited_quota_user_can_upload(self, api_client, tmp_path): + def test_unlimited_quota_user_can_upload(self, api_client):- def test_quota_limit_blocks_excess_uploads(self, api_client, tmp_path): + def test_quota_limit_blocks_excess_uploads(self, api_client):Also applies to: 482-482
| data = response.json() | ||
| if "task_status_url" in data: | ||
| task_id = get_task_id(data) | ||
| wait_for_task(api_client, task_id, headers={"Authorization": f"Bearer {user_token}"}) |
There was a problem hiding this comment.
Task wait is skipped when response contains task_id instead of task_status_url.
All quota tests only check if "task_status_url" in data, but get_task_id() (line 52‑58) also handles a bare "task_id" key. If the API ever returns {"task_id": "..."} without task_status_url, the test won't wait for indexing to finish, causing racy assertions on file_count. This pattern repeats across every upload in this class (~10 occurrences).
Consider aligning with the existing tests (e.g., line 157) that unconditionally call get_task_id + wait_for_task:
Proposed fix (apply to all similar blocks)
data = response.json()
- if "task_status_url" in data:
- task_id = get_task_id(data)
+ task_id = get_task_id(data)
+ if task_id:
wait_for_task(api_client, task_id, headers={"Authorization": f"Bearer {user_token}"})Or, if some responses genuinely have no task info, guard with a try/except around get_task_id.
🤖 Prompt for AI Agents
In `@tests/api_tests/test_indexer.py` around lines 473 - 476, The test blocks
currently only wait when "task_status_url" is present, which misses responses
that include only "task_id"; update each upload handling block to always attempt
to extract a task id via get_task_id(response.json()) and then call
wait_for_task(api_client, task_id, headers={"Authorization": f"Bearer
{user_token}"}), or wrap get_task_id in a try/except to skip waiting only if no
task info is present—apply this change to every occurrence in the test class
where the current code checks for "task_status_url" so uploads that return
{"task_id": "..."} are properly awaited.
requires database migration to add file_quota column to users table.
Features:
- Add file_quota field to users table (nullable integer)
- Add DEFAULT_FILE_QUOTA env var to set global default quota
- Add PATCH /users/{user_id}/quota endpoint to update user quotas
- Add quota enforcement on file upload (indexed files + pending tasks)
- Admins bypass quota checks; quota<=0 for a user means unlimited; quota > 0 to limit user's quota
API Changes:
- POST /users/ now accepts optional file_quota parameter
- GET /users/ and GET /users/{id} return file_quota in response
- New PATCH /users/{user_id}/quota endpoint for quota management
- GET /users/info now renders additional fields: file_count, pending_files, total_files and file_quota
Tests:
- test_update_user_quota: verify quota update from 10 to 12
- test_user_default_quota: verify None quota defaults to 10
- TestUserQuotaEnforcement class in test_indexer.py:
- test_unlimited_quota_user_can_exceed_default: user with quota=0 uploads 11 files
- test_quota_limit_blocks_excess_uploads: user with quota=5 blocked on 6th file
Perform migration: Refer to the doc
https://github.com/linagora/openrag/blob/0b5f84cd880e7c75db6f78d12a7681227721c2ad/docs/content/docs/documentation/sql_migration.mdx?plain=1#L47-L53
Replace SQL triggers with application-level file_count management. Add created_by column to files table to track who uploaded each file. Increment/decrement the uploader's file_count in app code, using func.greatest to clamp decrements to zero for race-condition safety. Remove user_id from delete_file/delete_partition signatures since the uploader is now looked up from the file record. Fix delete_user to also clean up Milvus data via self.delete_partition(). Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Verifies that deleting a partition only decrements the uploader's file_count by the number of files in that partition, leaving files in other partitions unaffected. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Compares each user's materialized file_count against the actual count from the files table and displays results in a color-coded table. Supports --fix to correct mismatched values. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Two users each upload files to the same partition. On deletion, verify each user's file_count is decremented independently. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
5dea0be to
95c9d2f
Compare
2 participants
PR reopen for additional commits before merge. See original PR: #234
Summary by CodeRabbit
New Features
Documentation
Tests
Tools