From 0824eaa014b7e22bb1968c2ff3b770c926cb1af4 Mon Sep 17 00:00:00 2001 From: Ayush8923 <80516839+Ayush8923@users.noreply.github.com> Date: Thu, 3 Sep 2026 15:38:32 +0530 Subject: [PATCH 1/4] fix(*): knowledge base pre signed url --- backend/app/api/routes/documents.py | 3 + backend/app/core/cloud/storage.py | 28 +++++++-- backend/app/models/document.py | 7 ++- backend/app/services/doctransform/job.py | 9 ++- backend/app/services/documents/helpers.py | 18 +++++- backend/app/tests/core/cloud/test_storage.py | 54 +++++++++++++++- .../tests/services/documents/test_helpers.py | 62 +++++++++++++++++++ docs/wiki/modules/knowledge-base.md | 1 + 8 files changed, 169 insertions(+), 13 deletions(-) diff --git a/backend/app/api/routes/documents.py b/backend/app/api/routes/documents.py index f2e3f2d6b..648a7c537 100644 --- a/backend/app/api/routes/documents.py +++ b/backend/app/api/routes/documents.py @@ -179,6 +179,9 @@ async def upload_doc( source_document, from_attributes=True ) document_schema.signed_url = storage.get_signed_url( + source_document.object_store_url, filename=source_document.fname + ) + document_schema.preview_url = storage.get_signed_url( source_document.object_store_url ) diff --git a/backend/app/core/cloud/storage.py b/backend/app/core/cloud/storage.py index a0b451edd..a65d43fc5 100644 --- a/backend/app/core/cloud/storage.py +++ b/backend/app/core/cloud/storage.py @@ -7,7 +7,7 @@ import functools as ft from pathlib import Path from dataclasses import dataclass, asdict -from urllib.parse import ParseResult, urlparse, urlunparse +from urllib.parse import ParseResult, quote, urlparse, urlunparse from abc import ABC, abstractmethod from typing import Any @@ -32,6 +32,18 @@ def _mask(value: str | None) -> str: logger = logging.getLogger(__name__) +def _attachment_disposition(filename: str) -> str: + """ + Build a Content-Disposition header that forces a download. + """ + ascii_fallback = ( + filename.encode("ascii", "replace").decode("ascii").replace('"', "") + ) + return ( + f"attachment; filename=\"{ascii_fallback}\"; filename*=UTF-8''{quote(filename)}" + ) + + class CloudStorageError(Exception): pass @@ -151,7 +163,9 @@ def get_file_size_kb(self, url: str) -> float: pass @abstractmethod - def get_signed_url(self, url: str, expires_in: int = 3600) -> str: + def get_signed_url( + self, url: str, expires_in: int = 3600, filename: str | None = None + ) -> str: """Generate a signed URL with an optional expiry""" pass @@ -255,11 +269,14 @@ def get_file_size_kb(self, url: str) -> float: # Maximum allowed expiry for signed URLs (24 hours) MAX_SIGNED_URL_EXPIRY = 86400 - def get_signed_url(self, url: str, expires_in: int = 3600) -> str: + def get_signed_url( + self, url: str, expires_in: int = 3600, filename: str | None = None + ) -> str: """ Generate a signed S3 URL for the given file. :param url: S3 url (e.g., s3://bucket/key) :param expires_in: Expiry time in seconds (default: 1 hour, max: 24 hours) + :param filename: When set, the URL forces a download under this name :return: Signed URL as string """ # Cap expiry at maximum allowed value to prevent excessively long-lived URLs @@ -267,9 +284,12 @@ def get_signed_url(self, url: str, expires_in: int = 3600) -> str: name = SimpleStorageName.from_url(url) try: + params: dict[str, str] = {"Bucket": name.Bucket, "Key": name.Key} + if filename: + params["ResponseContentDisposition"] = _attachment_disposition(filename) signed_url = self.aws.client.generate_presigned_url( "get_object", - Params={"Bucket": name.Bucket, "Key": name.Key}, + Params=params, ExpiresIn=expires_in, ) logger.info( diff --git a/backend/app/models/document.py b/backend/app/models/document.py index 3f3c80996..79c73706e 100644 --- a/backend/app/models/document.py +++ b/backend/app/models/document.py @@ -82,7 +82,12 @@ class Document(DocumentBase, table=True): class DocumentPublic(DocumentBase): id: UUID = Field(description="The unique identifier of the document") signed_url: str | None = Field( - default=None, description="A signed URL for accessing the document" + default=None, + description="A signed URL that downloads the document under its original name", + ) + preview_url: str | None = Field( + default=None, + description="A signed URL that renders the document inline in the browser", ) inserted_at: datetime = Field( description="The timestamp when the document was inserted" diff --git a/backend/app/services/doctransform/job.py b/backend/app/services/doctransform/job.py index d87a38a32..854142e33 100644 --- a/backend/app/services/doctransform/job.py +++ b/backend/app/services/doctransform/job.py @@ -245,11 +245,14 @@ def execute_job( ), ) - signed_url = None + url_update: dict[str, str] = {} try: get_signed_url = getattr(storage, "get_signed_url", None) if callable(get_signed_url): - signed_url = get_signed_url(created.object_store_url) + url_update["signed_url"] = get_signed_url( + created.object_store_url, filename=created.fname + ) + url_update["preview_url"] = get_signed_url(created.object_store_url) except Exception as e: logger.warning( "[doc_transform] failed to generate signed URL for doc %s: %s", @@ -259,7 +262,7 @@ def execute_job( transformed_public = TransformedDocumentPublic.model_validate( created, - update={"signed_url": signed_url} if signed_url else None, + update=url_update or None, ) success_payload = build_success_payload(job_for_payload, transformed_public) diff --git a/backend/app/services/documents/helpers.py b/backend/app/services/documents/helpers.py index ffbb1096b..ea566d282 100644 --- a/backend/app/services/documents/helpers.py +++ b/backend/app/services/documents/helpers.py @@ -179,6 +179,18 @@ def _to_public_schema(doc: Document) -> PublicDoc: return TransformedDocumentPublic.model_validate(doc, from_attributes=True) +def _attach_urls(schema: PublicDoc, document: Document, storage: CloudStorage) -> None: + """ + Two links for the same object: the download one carries a Content-Disposition + that forces a save under the original filename, the preview one is left bare + so browsers render it inline (iframe previews). + """ + schema.signed_url = storage.get_signed_url( + document.object_store_url, filename=document.fname + ) + schema.preview_url = storage.get_signed_url(document.object_store_url) + + def build_document_schema( *, document: Document, @@ -187,7 +199,7 @@ def build_document_schema( ) -> PublicDoc: schema = _to_public_schema(document) if include_url and storage: - schema.signed_url = storage.get_signed_url(document.object_store_url) + _attach_urls(schema, document, storage) return schema @@ -201,7 +213,7 @@ def build_document_schemas( for doc in documents: schema = _to_public_schema(doc) if include_url and storage: - schema.signed_url = storage.get_signed_url(doc.object_store_url) + _attach_urls(schema, doc, storage) out.append(schema) return out @@ -225,7 +237,7 @@ def build_job_schema( object_url = doc.object_store_url if include_url and storage and object_url: - transformed_doc_schema.signed_url = storage.get_signed_url(object_url) + _attach_urls(transformed_doc_schema, doc, storage) job_schema = DocTransformationJobPublic.model_validate(job, from_attributes=True) return job_schema.model_copy( diff --git a/backend/app/tests/core/cloud/test_storage.py b/backend/app/tests/core/cloud/test_storage.py index 95e9ff7ff..e7ffb8cab 100644 --- a/backend/app/tests/core/cloud/test_storage.py +++ b/backend/app/tests/core/cloud/test_storage.py @@ -1,8 +1,14 @@ """Tests for app.core.cloud.storage helpers.""" -from unittest.mock import patch +from types import SimpleNamespace +from unittest.mock import MagicMock, patch +from uuid import uuid4 -from app.core.cloud.storage import GCS_SCOPES, build_gcp_sa_credentials +from app.core.cloud.storage import ( + GCS_SCOPES, + AmazonCloudStorage, + build_gcp_sa_credentials, +) def test_build_gcp_sa_credentials_passes_key_and_scopes(): @@ -14,3 +20,47 @@ def test_build_gcp_sa_credentials_passes_key_and_scopes(): mock_from_info.assert_called_once_with(sa_key, scopes=list(GCS_SCOPES)) assert creds is mock_from_info.return_value + + +def _amazon_storage_with_mock_client(mock_client): + storage = AmazonCloudStorage(project_id=1, storage_path=uuid4()) + storage.aws = SimpleNamespace(client=mock_client) + return storage + + +def test_get_signed_url_without_filename_omits_content_disposition(): + mock_client = MagicMock() + storage = _amazon_storage_with_mock_client(mock_client) + + storage.get_signed_url("s3://bucket/key.pdf") + + params = mock_client.generate_presigned_url.call_args.kwargs["Params"] + assert "ResponseContentDisposition" not in params + + +def test_get_signed_url_with_filename_forces_attachment(): + mock_client = MagicMock() + storage = _amazon_storage_with_mock_client(mock_client) + + storage.get_signed_url("s3://bucket/key.pdf", filename="report.pdf") + + disposition = mock_client.generate_presigned_url.call_args.kwargs["Params"][ + "ResponseContentDisposition" + ] + assert ( + disposition + == "attachment; filename=\"report.pdf\"; filename*=UTF-8''report.pdf" + ) + + +def test_get_signed_url_encodes_non_ascii_filename(): + mock_client = MagicMock() + storage = _amazon_storage_with_mock_client(mock_client) + + storage.get_signed_url("s3://bucket/key.pdf", filename="रिपोर्ट.pdf") + + disposition = mock_client.generate_presigned_url.call_args.kwargs["Params"][ + "ResponseContentDisposition" + ] + assert disposition.startswith('attachment; filename="') + assert "filename*=UTF-8''%E0%A4%B0" in disposition diff --git a/backend/app/tests/services/documents/test_helpers.py b/backend/app/tests/services/documents/test_helpers.py index 3da9ba30c..05ea93089 100644 --- a/backend/app/tests/services/documents/test_helpers.py +++ b/backend/app/tests/services/documents/test_helpers.py @@ -1,10 +1,14 @@ import csv from io import BytesIO +from uuid import uuid4 import pytest from fastapi import HTTPException, UploadFile +from app.models import Document from app.services.documents.helpers import ( + build_document_schema, + build_document_schemas, calculate_file_size, validate_upload, ) @@ -338,3 +342,61 @@ def test_missing_filename_raises_400(self) -> None: validate_upload(src=src, target_format=None, transformer=None) assert exc_info.value.status_code == 400 assert exc_info.value.detail == "Uploaded file has no filename" + + +def make_document(fname: str = "policy.pdf") -> Document: + return Document( + id=uuid4(), + fname=fname, + object_store_url="s3://bucket/key", + project_id=1, + ) + + +class FakeStorage: + """Records how each signed URL was requested.""" + + def __init__(self) -> None: + self.calls: list[dict] = [] + + def get_signed_url( + self, url: str, expires_in: int = 3600, filename: str | None = None + ) -> str: + self.calls.append({"url": url, "filename": filename}) + return f"https://signed/{'download' if filename else 'preview'}" + + +class TestDocumentUrls: + def test_download_url_carries_filename_preview_url_does_not(self) -> None: + storage = FakeStorage() + document = make_document() + + schema = build_document_schema( + document=document, include_url=True, storage=storage + ) + + assert schema.signed_url == "https://signed/download" + assert schema.preview_url == "https://signed/preview" + assert [c["filename"] for c in storage.calls] == ["policy.pdf", None] + + def test_no_urls_when_include_url_is_false(self) -> None: + storage = FakeStorage() + + schema = build_document_schema( + document=make_document(), include_url=False, storage=storage + ) + + assert schema.signed_url is None + assert schema.preview_url is None + assert storage.calls == [] + + def test_list_attaches_both_urls_per_document(self) -> None: + storage = FakeStorage() + documents = [make_document("a.pdf"), make_document("b.pdf")] + + schemas = build_document_schemas( + documents=documents, include_url=True, storage=storage + ) + + assert all(s.signed_url and s.preview_url for s in schemas) + assert [c["filename"] for c in storage.calls] == ["a.pdf", None, "b.pdf", None] diff --git a/docs/wiki/modules/knowledge-base.md b/docs/wiki/modules/knowledge-base.md index c6b8462df..735bc0215 100644 --- a/docs/wiki/modules/knowledge-base.md +++ b/docs/wiki/modules/knowledge-base.md @@ -33,6 +33,7 @@ All paths relative to `backend/app/`. - OpenAI vector stores / file uploads, object storage (`core/cloud/storage.py`), Zerox/OCR for transforms. ## Gotchas +- Document responses carry two links for the same object: `signed_url` downloads, `preview_url` renders inline. `put()` stores the upload's `Content-Type`, so a bare presigned URL for a PDF opens in a tab (CSV/XLSX datasets happen to download, which is why only KB docs looked broken). `get_signed_url(..., filename=...)` adds `ResponseContentDisposition: attachment`; `_attach_urls()` in `services/documents/helpers.py` sets both fields, and `api/routes/documents.py` plus `services/doctransform/job.py` do the same for their responses. An attachment URL will not render in an `