Refactor Document Module and Cloud Storage to use project-based organization and add signed URL support (#346)

* Add project id in model and use project id for crud

* add support for signed url

* Add storage path in project table

* fix collection endpoint

* Fix api routes test cases for document module

* fix crud unit test for document module

* fix unit test for collections

* use project id instead of project model in document store utils

* Fix migration

* precommit

* introduce get_cloud_storage method

* fix testcase

* fix message

* implement abstract method in CloudStorage class

* use POSIX separators and Reject absolute file_path

Author: avirajsingh7

backend/app/alembic/versions/40307ab77e9f_add_storage_path_to_project_and_project_to_document_table.py

@@ -0,0 +1,108 @@
+"""add storage_path to project and project_id to document table
+
+Revision ID: 40307ab77e9f
+Revises: 8725df286943
+Create Date: 2025-08-28 10:54:30.712627
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = "40307ab77e9f"
+down_revision = "8725df286943"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+
+    op.add_column("project", sa.Column("storage_path", sa.Uuid(), nullable=True))
+
+    conn = op.get_bind()
+    conn.execute(sa.text("UPDATE project SET storage_path = gen_random_uuid()"))
+
+    op.alter_column("project", "storage_path", nullable=False)
+    op.create_unique_constraint("uq_project_storage_path", "project", ["storage_path"])
+
+    op.add_column("document", sa.Column("project_id", sa.Integer(), nullable=True))
+    op.add_column("document", sa.Column("is_deleted", sa.Boolean(), nullable=True))
+
+    conn.execute(
+        sa.text(
+            """
+        UPDATE document
+        SET is_deleted = CASE
+            WHEN deleted_at IS NULL THEN false
+            ELSE true
+        END
+    """
+        )
+    )
+    conn.execute(
+        sa.text(
+            """
+        UPDATE document
+        SET project_id = (
+            SELECT project_id FROM apikey
+            WHERE apikey.user_id = document.owner_id
+            LIMIT 1
+        )
+    """
+        )
+    )
+
+    op.alter_column("document", "is_deleted", nullable=False)
+    op.alter_column("document", "project_id", nullable=False)
+
+    op.drop_constraint("document_owner_id_fkey", "document", type_="foreignkey")
+    op.create_foreign_key(
+        None, "document", "project", ["project_id"], ["id"], ondelete="CASCADE"
+    )
+    op.drop_column("document", "owner_id")
+
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_constraint("uq_project_storage_path", "project", type_="unique")
+    op.drop_column("project", "storage_path")
+
+    op.add_column(
+        "document",
+        sa.Column("owner_id", sa.Integer(), autoincrement=False, nullable=True),
+    )
+
+    conn = op.get_bind()
+    # Backfill owner_id from project_id using api_key mapping
+    conn.execute(
+        sa.text(
+            """
+        UPDATE document d
+        SET owner_id = (
+            SELECT user_id
+            FROM apikey a
+            WHERE a.project_id = d.project_id
+            LIMIT 1
+        )
+    """
+        )
+    )
+
+    op.alter_column("document", "owner_id", nullable=False)
+
+    op.drop_constraint("document_project_id_fkey", "document", type_="foreignkey")
+    op.create_foreign_key(
+        "document_owner_id_fkey",
+        "document",
+        "user",
+        ["owner_id"],
+        ["id"],
+        ondelete="CASCADE",
+    )
+    op.drop_column("document", "is_deleted")
+    op.drop_column("document", "project_id")
+    # ### end Alembic commands ###

backend/app/api/routes/collections.py

@@ -15,7 +15,7 @@
 from sqlalchemy.exc import SQLAlchemyError
 
 from app.api.deps import CurrentUser, SessionDep, CurrentUserOrgProject
-from app.core.cloud import AmazonCloudStorage
+from app.core.cloud import get_cloud_storage
 from app.api.routes.responses import handle_openai_error
 from app.core.util import now, post_callback
 from app.crud import (
@@ -24,7 +24,7 @@
     DocumentCollectionCrud,
 )
 from app.crud.rag import OpenAIVectorStoreCrud, OpenAIAssistantCrud
-from app.models import Collection, Document
+from app.models import Collection, Document, DocumentPublic
 from app.models.collection import CollectionStatus
 from app.utils import APIResponse, load_description, get_openai_client
 
@@ -225,8 +225,8 @@ def do_create_collection(
         else WebHookCallback(request.callback_url, payload)
     )
 
-    storage = AmazonCloudStorage(current_user)
-    document_crud = DocumentCrud(session, current_user.id)
+    storage = get_cloud_storage(session=session, project_id=current_user.project_id)
+    document_crud = DocumentCrud(session, current_user.project_id)
     assistant_crud = OpenAIAssistantCrud(client)
     vector_store_crud = OpenAIVectorStoreCrud(client)
     collection_crud = CollectionCrud(session, current_user.id)
@@ -423,7 +423,7 @@ def list_collections(
 @router.post(
     "/docs/{collection_id}",
     description=load_description("collections/docs.md"),
-    response_model=APIResponse[List[Document]],
+    response_model=APIResponse[List[DocumentPublic]],
 )
 def collection_documents(
     session: SessionDep,

backend/app/api/routes/documents.py

@@ -3,14 +3,14 @@
 from typing import List
 from pathlib import Path
 
-from fastapi import APIRouter, File, UploadFile, Query
+from fastapi import APIRouter, File, UploadFile, Query, HTTPException
 from fastapi import Path as FastPath
 
 from app.crud import DocumentCrud, CollectionCrud
-from app.models import Document
+from app.models import Document, DocumentPublic, Message
 from app.utils import APIResponse, load_description, get_openai_client
 from app.api.deps import CurrentUser, SessionDep, CurrentUserOrgProject
-from app.core.cloud import AmazonCloudStorage
+from app.core.cloud import get_cloud_storage
 from app.crud.rag import OpenAIAssistantCrud
 
 logger = logging.getLogger(__name__)
@@ -20,35 +20,35 @@
 @router.get(
     "/list",
     description=load_description("documents/list.md"),
-    response_model=APIResponse[List[Document]],
+    response_model=APIResponse[List[DocumentPublic]],
 )
 def list_docs(
     session: SessionDep,
-    current_user: CurrentUser,
+    current_user: CurrentUserOrgProject,
     skip: int = Query(0, ge=0),
     limit: int = Query(100, gt=0, le=100),
 ):
-    crud = DocumentCrud(session, current_user.id)
+    crud = DocumentCrud(session, current_user.project_id)
     data = crud.read_many(skip, limit)
     return APIResponse.success_response(data)
 
 
 @router.post(
     "/upload",
     description=load_description("documents/upload.md"),
-    response_model=APIResponse[Document],
+    response_model=APIResponse[DocumentPublic],
 )
 def upload_doc(
     session: SessionDep,
-    current_user: CurrentUser,
+    current_user: CurrentUserOrgProject,
     src: UploadFile = File(...),
 ):
-    storage = AmazonCloudStorage(current_user)
+    storage = get_cloud_storage(session=session, project_id=current_user.project_id)
     document_id = uuid4()
 
     object_store_url = storage.put(src, Path(str(document_id)))
 
-    crud = DocumentCrud(session, current_user.id)
+    crud = DocumentCrud(session, current_user.project_id)
     document = Document(
         id=document_id,
         fname=src.filename,
@@ -58,10 +58,10 @@ def upload_doc(
     return APIResponse.success_response(data)
 
 
-@router.get(
+@router.delete(
     "/remove/{doc_id}",
     description=load_description("documents/delete.md"),
-    response_model=APIResponse[Document],
+    response_model=APIResponse[Message],
 )
 def remove_doc(
     session: SessionDep,
@@ -73,18 +73,21 @@ def remove_doc(
     )
 
     a_crud = OpenAIAssistantCrud(client)
-    d_crud = DocumentCrud(session, current_user.id)
+    d_crud = DocumentCrud(session, current_user.project_id)
     c_crud = CollectionCrud(session, current_user.id)
 
     document = d_crud.delete(doc_id)
     data = c_crud.delete(document, a_crud)
-    return APIResponse.success_response(data)
+
+    return APIResponse.success_response(
+        Message(message="Document Deleted Successfully")
+    )
 
 
 @router.delete(
     "/remove/{doc_id}/permanent",
     description=load_description("documents/permanent_delete.md"),
-    response_model=APIResponse[Document],
+    response_model=APIResponse[Message],
 )
 def permanent_delete_doc(
     session: SessionDep,
@@ -94,11 +97,10 @@ def permanent_delete_doc(
     client = get_openai_client(
         session, current_user.organization_id, current_user.project_id
     )
-
     a_crud = OpenAIAssistantCrud(client)
-    d_crud = DocumentCrud(session, current_user.id)
+    d_crud = DocumentCrud(session, current_user.project_id)
     c_crud = CollectionCrud(session, current_user.id)
-    storage = AmazonCloudStorage(current_user)
+    storage = get_cloud_storage(session=session, project_id=current_user.project_id)
 
     document = d_crud.read_one(doc_id)
 
@@ -107,19 +109,30 @@ def permanent_delete_doc(
     storage.delete(document.object_store_url)
     d_crud.delete(doc_id)
 
-    return APIResponse.success_response(document)
+    return APIResponse.success_response(
+        Message(message="Document permanently deleted successfully")
+    )
 
 
 @router.get(
     "/info/{doc_id}",
     description=load_description("documents/info.md"),
-    response_model=APIResponse[Document],
+    response_model=APIResponse[DocumentPublic],
 )
 def doc_info(
     session: SessionDep,
-    current_user: CurrentUser,
+    current_user: CurrentUserOrgProject,
     doc_id: UUID = FastPath(description="Document to retrieve"),
+    include_url: bool = Query(
+        False, description="Include a signed URL to access the document"
+    ),
 ):
-    crud = DocumentCrud(session, current_user.id)
-    data = crud.read_one(doc_id)
-    return APIResponse.success_response(data)
+    crud = DocumentCrud(session, current_user.project_id)
+    document = crud.read_one(doc_id)
+
+    doc_schema = DocumentPublic.model_validate(document, from_attributes=True)
+    if include_url:
+        storage = get_cloud_storage(session=session, project_id=current_user.project_id)
+        doc_schema.signed_url = storage.get_signed_url(document.object_store_url)
+
+    return APIResponse.success_response(doc_schema)

backend/app/core/cloud/__init__.py

@@ -3,4 +3,5 @@
     AmazonCloudStorageClient,
     CloudStorage,
     CloudStorageError,
+    get_cloud_storage,
 )