✨ Add S3 file upload functionality and update document creation process

Author: monicasmith463

backend/app/api/routes/documents.py

@@ -1,109 +1,45 @@
-import uuid
 from typing import Any
-
-from fastapi import APIRouter, HTTPException
-from sqlmodel import func, select
-
 from app.api.deps import CurrentUser, SessionDep
-from app.models import Document, DocumentCreate, DocumentPublic, DocumentsPublic, DocumentUpdate, Message
+from app.models import Document, DocumentCreate, DocumentPublic
+from fastapi import APIRouter, BackgroundTasks, File, UploadFile, HTTPException
+from app.s3 import upload_file_to_s3, generate_s3_url
 
 router = APIRouter(prefix="/documents", tags=["documents"])
-
-
-@router.get("/", response_model=DocumentsPublic)
-def read_documents(
-    session: SessionDep, current_user: CurrentUser, skip: int = 0, limit: int = 100
-) -> Any:
-    """
-    Retrieve documents.
-    """
-
-    if current_user.is_superuser:
-        count_statement = select(func.count()).select_from(Document)
-        count = session.exec(count_statement).one()
-        statement = select(Document).offset(skip).limit(limit)
-        documents = session.exec(statement).all()
-    else:
-        count_statement = (
-            select(func.count())
-            .select_from(Document)
-            .where(Document.owner_id == current_user.id)
-        )
-        count = session.exec(count_statement).one()
-        statement = (
-            select(Document)
-            .where(Document.owner_id == current_user.id)
-            .offset(skip)
-            .limit(limit)
-        )
-        documents = session.exec(statement).all()
-
-    return DocumentsPublic(data=documents, count=count)
-
-
-@router.get("/{id}", response_model=DocumentPublic)
-def read_document(session: SessionDep, current_user: CurrentUser, id: uuid.UUID) -> Any:
-    """
-    Get document by ID.
-    """
-    document = session.get(Document, id)
-    if not document:
-        raise HTTPException(status_code=404, detail="Document not found")
-    if not current_user.is_superuser and (document.owner_id != current_user.id):
-        raise HTTPException(status_code=400, detail="Not enough permissions")
-    return document
-
-
+    
 @router.post("/", response_model=DocumentPublic)
 def create_document(
-    *, session: SessionDep, current_user: CurrentUser, document_in: DocumentCreate
-) -> Any:
-    """
-    Create new document.
-    """
+    *, session: SessionDep, current_user: CurrentUser,
+      background_tasks: BackgroundTasks,
+      file: UploadFile = File(...), 
+):
+    key = None
+    try:
+        user_id = current_user.id
+        key = upload_file_to_s3(file, current_user.id)
+    except Exception as e:
+        raise HTTPException(500, f"Failed to upload file: {key}. Error: {str(e)}")
+
+    try:
+        url = generate_s3_url(key)
+    except Exception as e:
+        raise HTTPException(500, f"Could not generate URL for file key: {key}")
+
+    document_in = DocumentCreate(
+        filename=file.filename,
+        content_type=file.content_type,
+        size=file.size,  
+        s3_url=url,
+    )
+
     document = Document.model_validate(document_in, update={"owner_id": current_user.id})
-    session.add(document)
-    session.commit()
-    session.refresh(document)
-    return document
 
 
-@router.put("/{id}", response_model=DocumentPublic)
-def update_document(
-    *,
-    session: SessionDep,
-    current_user: CurrentUser,
-    id: uuid.UUID,
-    document_in: DocumentUpdate,
-) -> Any:
-    """
-    Update an document.
-    """
-    document = session.get(Document, id)
-    if not document:
-        raise HTTPException(status_code=404, detail="Document not found")
-    if not current_user.is_superuser and (document.owner_id != current_user.id):
-        raise HTTPException(status_code=400, detail="Not enough permissions")
-    update_dict = document_in.model_dump(exclude_unset=True)
-    document.sqlmodel_update(update_dict)
     session.add(document)
     session.commit()
     session.refresh(document)
-    return document
 
+    # 3. Kick off background job
+    print("Document created, starting background task...")
+    # background_tasks.add_task(generate_questions, document.id)
 
-@router.delete("/{id}")
-def delete_document(
-    session: SessionDep, current_user: CurrentUser, id: uuid.UUID
-) -> Message:
-    """
-    Delete an document.
-    """
-    document = session.get(Document, id)
-    if not document:
-        raise HTTPException(status_code=404, detail="Document not found")
-    if not current_user.is_superuser and (document.owner_id != current_user.id):
-        raise HTTPException(status_code=400, detail="Not enough permissions")
-    session.delete(document)
-    session.commit()
-    return Message(message="Document deleted successfully")
+    return document

backend/app/core/config.py

@@ -57,6 +57,11 @@ def all_cors_origins(self) -> list[str]:
     POSTGRES_PASSWORD: str = ""
     POSTGRES_DB: str = ""
 
+    AWS_ACCESS_KEY_ID: str = ""
+    AWS_SECRET_ACCESS_KEY: str = ""
+    AWS_REGION: str = ""
+    S3_BUCKET_NAME: str = ""
+
     @computed_field  # type: ignore[prop-decorator]
     @property
     def SQLALCHEMY_DATABASE_URI(self) -> PostgresDsn:

backend/app/s3.py

@@ -7,9 +7,13 @@
 
 def upload_file_to_s3(file: UploadFile, user_id: str) -> str:
     extension = file.filename.split(".")[-1]
+    bucket = settings.S3_BUCKET_NAME
     key = f"documents/{user_id}/{uuid.uuid4()}.{extension}"
 
-    s3.upload_fileobj(file.file, settings.S3_BUCKET_NAME, key)
+    try:
+        s3.upload_fileobj(file.file, bucket, key)
+    except Exception as e:
+        raise Exception(f"Failed to upload file to S3: {str(e)}")
 
     return key
 

backend/app/tests/api/routes/test_documents.py

@@ -0,0 +1,52 @@
+import uuid
+
+from fastapi.testclient import TestClient
+from sqlmodel import Session
+from unittest.mock import patch
+from app.core.config import settings
+from app.tests.utils.document import create_random_document
+import io
+
+
+def skip_test_create_document(
+    client: TestClient, superuser_token_headers: dict[str, str]
+) -> None:
+    '''Test creating a document with a file upload with the real S3 service.'''
+    file_content = b"%PDF-1.4 test file content"
+
+    response = client.post(
+        f"{settings.API_V1_STR}/documents/",
+        headers=superuser_token_headers,
+        files={
+            "file": ("example.pdf", io.BytesIO(file_content), "application/pdf")
+        },
+    )
+
+    assert response.status_code == 200
+    content = response.json()
+    assert "id" in content, "actual response: " + str(content)
+    # assert content["title"] == metadata["title"]
+    # assert content["description"] == metadata["description"]
+    # assert "id" in content
+    # assert "owner_id" in content
+
+def test_create_document(
+    client: TestClient, superuser_token_headers: dict[str, str]
+) -> None:
+    '''Test creating a document with a file upload using mocked S3.'''
+    file_content = b"%PDF-1.4 test file content"
+
+    with patch("app.api.routes.documents.upload_file_to_s3", return_value="document-slug"):
+        response = client.post(
+            f"{settings.API_V1_STR}/documents/",
+            headers=superuser_token_headers,
+            files={
+                "file": ("example.pdf", io.BytesIO(file_content), "application/pdf")
+            },
+        )
+
+    assert response.status_code == 200, f"Unexpected response: {response.content}"
+    content = response.json()
+    assert "id" in content, "actual response: " + str(content)
+    assert "document-slug" in content["s3_url"], "S3 URL should match mocked value"
+    assert content["filename"] == "example.pdf", "Filename should match uploaded file"

backend/pyproject.toml

@@ -18,6 +18,7 @@ dependencies = [
     "sqlmodel<1.0.0,>=0.0.21",
     # Pin bcrypt until passlib supports the latest
     "bcrypt==4.3.0",
+    "boto3",
     "pydantic-settings<3.0.0,>=2.2.1",
     "sentry-sdk[fastapi]<2.0.0,>=1.40.6",
     "pyjwt<3.0.0,>=2.8.0",

backend/requirements.txt

@@ -0,0 +1,81 @@
+alembic==1.15.2
+annotated-types==0.7.0
+anyio==4.6.0
+-e file:///Users/monicasmith/Documents/personal-projects/study-assistant/backend
+bcrypt==4.3.0
+boto3==1.39.9
+botocore==1.39.9
+cachetools==5.5.0
+certifi==2024.8.30
+cfgv==3.4.0
+chardet==5.2.0
+charset-normalizer==3.3.2
+click==8.1.7
+coverage==7.6.1
+cssselect==1.2.0
+cssutils==2.11.1
+distlib==0.3.8
+dnspython==2.6.1
+email-validator==2.2.0
+emails==0.6
+fastapi==0.115.0
+fastapi-cli==0.0.5
+filelock==3.16.1
+h11==0.14.0
+httpcore==1.0.5
+httptools==0.6.1
+httpx==0.28.1
+identify==2.6.1
+idna==3.10
+iniconfig==2.0.0
+jinja2==3.1.6
+jmespath==1.0.1
+lxml==5.3.0
+mako==1.3.5
+markdown-it-py==3.0.0
+markupsafe==2.1.5
+mdurl==0.1.2
+more-itertools==10.5.0
+mypy==1.11.2
+mypy-extensions==1.0.0
+nodeenv==1.9.1
+packaging==24.1
+passlib==1.7.4
+platformdirs==4.3.6
+pluggy==1.5.0
+pre-commit==3.8.0
+premailer==3.10.0
+psycopg==3.2.2
+psycopg-binary==3.2.2
+pydantic==2.9.2
+pydantic-core==2.23.4
+pydantic-settings==2.9.1
+pygments==2.18.0
+pyjwt==2.10.1
+pytest==7.4.4
+python-dateutil==2.9.0.post0
+python-dotenv==1.0.1
+python-multipart==0.0.20
+pyyaml==6.0.2
+requests==2.32.3
+rich==13.8.1
+ruff==0.6.7
+s3transfer==0.13.1
+sentry-sdk==1.45.1
+shellingham==1.5.4
+six==1.16.0
+sniffio==1.3.1
+sqlalchemy==2.0.35
+sqlmodel==0.0.24
+starlette==0.38.6
+tenacity==8.5.0
+typer==0.12.5
+types-passlib==1.7.7.20240819
+typing-extensions==4.12.2
+typing-inspection==0.4.0
+urllib3==2.2.3
+uvicorn==0.30.6
+uvloop==0.20.0
+virtualenv==20.26.5
+watchfiles==0.24.0
+websockets==13.1