diff --git a/app/common/filetype_ext.py b/app/common/filetype_ext.py index 0cc68297..42c58d79 100644 --- a/app/common/filetype_ext.py +++ b/app/common/filetype_ext.py @@ -1,7 +1,7 @@ from typing import Final import filetype # type: ignore[import-untyped] -from filetype.types import image # type: ignore[import-untyped] +from filetype.types import archive, image # type: ignore[import-untyped] FILE_HEADER_SIZE: Final[int] = 8192 @@ -17,6 +17,10 @@ image.Webp(), ] +SUPPORTED_DOCUMENT_FORMATS: list[filetype.Type] = [ + archive.Pdf(), +] + def match_filetype(obj: bytes, matchers: list[filetype.Type]) -> filetype.Type | None: return filetype.match(obj, matchers) @@ -24,3 +28,7 @@ def match_filetype(obj: bytes, matchers: list[filetype.Type]) -> filetype.Type | def match_image_filetype(obj: bytes) -> filetype.Type | None: return match_filetype(obj, SUPPORTED_IMAGE_FORMATS) + + +def match_document_filetype(obj: bytes) -> filetype.Type | None: + return match_filetype(obj, SUPPORTED_DOCUMENT_FORMATS) diff --git a/app/storage_v2/dependencies/uploads_dep.py b/app/storage_v2/dependencies/uploads_dep.py index a5c2c521..f3acfd02 100644 --- a/app/storage_v2/dependencies/uploads_dep.py +++ b/app/storage_v2/dependencies/uploads_dep.py @@ -4,7 +4,11 @@ from starlette import status from app.common.fastapi_ext import Responses, with_responses -from app.common.filetype_ext import FILE_HEADER_SIZE, match_image_filetype +from app.common.filetype_ext import ( + FILE_HEADER_SIZE, + match_document_filetype, + match_image_filetype, +) class FileFormatResponses(Responses): @@ -31,3 +35,20 @@ async def validate_image_upload(upload: UploadFile) -> UploadFile: ValidatedImageUpload = Annotated[UploadFile, Depends(validate_image_upload)] + + +async def validate_document_upload(upload: UploadFile) -> UploadFile: + upload_header_data = await upload.read(FILE_HEADER_SIZE) + document_type = match_document_filetype(upload_header_data) + + if document_type is None: + raise FileFormatResponses.WRONG_FORMAT + + if document_type.mime != upload.content_type: + raise FileFormatResponses.CONTENT_TYPE_MISMATCH + + await upload.seek(0) + return upload + + +ValidatedDocumentUpload = Annotated[UploadFile, Depends(validate_document_upload)] diff --git a/app/storage_v2/models/files_db.py b/app/storage_v2/models/files_db.py index 8dcf46c5..416a3e8a 100644 --- a/app/storage_v2/models/files_db.py +++ b/app/storage_v2/models/files_db.py @@ -14,6 +14,7 @@ class FileKind(StrEnum): UNCATEGORIZED = "uncategorized" IMAGE = "image" + DOCUMENT = "document" ContentDisposition = Literal["inline", "attachment"] @@ -21,14 +22,17 @@ class FileKind(StrEnum): FILE_KIND_TO_FOLDER: dict[FileKind, str] = { FileKind.UNCATEGORIZED: "uncategorized", FileKind.IMAGE: "images", + FileKind.DOCUMENT: "documents", } FILE_KIND_TO_MEDIA_TYPE: dict[FileKind, str | None] = { FileKind.UNCATEGORIZED: None, FileKind.IMAGE: "image/webp", + FileKind.DOCUMENT: "application/pdf", } FILE_KIND_TO_CONTENT_DISPOSITION: dict[FileKind, ContentDisposition] = { FileKind.UNCATEGORIZED: "attachment", FileKind.IMAGE: "inline", + FileKind.DOCUMENT: "inline", } diff --git a/app/storage_v2/routers/files_rst.py b/app/storage_v2/routers/files_rst.py index 548deafa..303f8704 100644 --- a/app/storage_v2/routers/files_rst.py +++ b/app/storage_v2/routers/files_rst.py @@ -17,7 +17,10 @@ StorageTokenResponses, UploadAllowedStorageTokenPayload, ) -from app.storage_v2.dependencies.uploads_dep import ValidatedImageUpload +from app.storage_v2.dependencies.uploads_dep import ( + ValidatedDocumentUpload, + ValidatedImageUpload, +) from app.storage_v2.models.access_groups_db import AccessGroupFile from app.storage_v2.models.files_db import File, FileKind @@ -85,6 +88,24 @@ async def upload_image_file( ) +@router.post( + "/file-kinds/document/files/", + status_code=status.HTTP_201_CREATED, + response_model=File.ResponseSchema, + summary="Upload a new document file", +) +async def upload_document_file( + storage_token_payload: UploadAllowedStorageTokenPayload, + upload: ValidatedDocumentUpload, +) -> File: + return await upload_file( + storage_token_payload=storage_token_payload, + upload_content=await upload.read(), + upload_filename=upload.filename, + file_kind=FileKind.DOCUMENT, + ) + + @router.get( "/files/{file_id}/meta/", response_model=File.ResponseSchema, @@ -107,7 +128,6 @@ def parse_http_datetime(header: str | None) -> datetime | None: @router.get( "/files/{file_id}/", - response_model=File.ResponseSchema, summary="Read any file by id", ) async def read_file( diff --git a/migrations/versions/063_document_files.py b/migrations/versions/063_document_files.py new file mode 100644 index 00000000..f8d9f284 --- /dev/null +++ b/migrations/versions/063_document_files.py @@ -0,0 +1,58 @@ +"""document_files + +Revision ID: 063 +Revises: 062 +Create Date: 2026-07-23 21:32:11.535202 + +""" + +from typing import Sequence, Union + +import sqlalchemy as sa +from alembic import op + +schema_name = "xi_back_2" +table_name = "files" +column_name = "kind" +enum_name = "file_kind" +tmp_enum_name = f"_{enum_name}" + +old_enum = sa.Enum("UNCATEGORIZED", "IMAGE", name=enum_name) + + +# revision identifiers, used by Alembic. +revision: str = "063" +down_revision: Union[str, None] = "061" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + op.execute(f"ALTER TYPE {enum_name} ADD VALUE 'DOCUMENT'") + + +def downgrade() -> None: + conn = op.get_bind() + + # rename new enum + op.execute(f"ALTER TYPE {enum_name} RENAME TO {tmp_enum_name}") + + # update old rows + metadata = sa.MetaData(schema=schema_name) + Files = sa.Table(table_name, metadata, autoload_with=conn) + + conn.execute( + sa.update(Files).where(Files.c.kind == "DOCUMENT").values(kind="UNCATEGORIZED") + ) + + # remove old members by updating to the new enum + old_enum.create(bind=conn) + op.execute( + f"ALTER TABLE {schema_name}.{table_name}" + f" ALTER COLUMN {column_name}" + f" TYPE {old_enum.name}" + f" USING {column_name}::text::{old_enum.name}" + ) + + # remove new enum + op.execute(f"DROP TYPE {tmp_enum_name}") diff --git a/tests/storage_v2/conftest.py b/tests/storage_v2/conftest.py index ab4f4bd3..bc795501 100644 --- a/tests/storage_v2/conftest.py +++ b/tests/storage_v2/conftest.py @@ -7,6 +7,9 @@ import pytest from faker import Faker +from faker_file.providers.pdf_file.generators.pil_generator import ( # type: ignore[import-untyped] + PilPdfGenerator, +) from PIL import Image from pytest_lazy_fixtures import lf from starlette.responses import FileResponse @@ -197,11 +200,33 @@ def png_image_file_input_data( ) +@pytest.fixture() +def pdf_document_file_content(faker: Faker) -> bytes: + return faker.pdf_file( # type: ignore[no-any-return] + pdf_generator_cls=PilPdfGenerator, + raw=True, + ) + + +@pytest.fixture() +def pdf_document_file_input_data( + faker: Faker, pdf_document_file_content: bytes +) -> FileInputData: + return FileInputData( + kind=FileKind.DOCUMENT, + name=faker.file_name(extension="pdf"), + input_content=pdf_document_file_content, + processed_content=pdf_document_file_content, + content_type="application/pdf", + ) + + @pytest.fixture( params=[ pytest.param(lf("uncategorized_file_input_data"), id="uncategorized"), pytest.param(lf("webp_image_file_input_data"), id="webp_image"), pytest.param(lf("png_image_file_input_data"), id="png_image"), + pytest.param(lf("pdf_document_file_input_data"), id="pdf_document"), ], ) def parametrized_file_input_data( diff --git a/tests/storage_v2/functional/test_file_uploads_rst.py b/tests/storage_v2/functional/test_file_uploads_rst.py index 7173678b..61058f95 100644 --- a/tests/storage_v2/functional/test_file_uploads_rst.py +++ b/tests/storage_v2/functional/test_file_uploads_rst.py @@ -116,6 +116,7 @@ async def test_file_uploading( ("image/tiff", "tif"), ("image/tiff", "tiff"), ("image/webp", "webp"), + ("application/pdf", "pdf"), ] @@ -124,6 +125,7 @@ async def test_file_uploading( [ pytest.param(lf("webp_image_file_input_data"), id="webp"), pytest.param(lf("png_image_file_input_data"), id="png"), + pytest.param(lf("pdf_document_file_input_data"), id="pdf"), ], ) async def test_image_file_uploading_content_type_mismatch( @@ -142,7 +144,7 @@ async def test_image_file_uploading_content_type_mismatch( assert_response( authorized_client.post( - "/api/protected/storage-service/v2/file-kinds/image/files/", + f"/api/protected/storage-service/v2/file-kinds/{file_input_data.kind}/files/", headers={"X-Storage-Token": file_upload_storage_token}, files={ "upload": ( @@ -157,21 +159,29 @@ async def test_image_file_uploading_content_type_mismatch( ) -async def test_image_file_uploading_wrong_content_format( +@pytest.mark.parametrize( + "file_input_data", + [ + pytest.param(lf("webp_image_file_input_data"), id="image"), + pytest.param(lf("pdf_document_file_input_data"), id="document"), + ], +) +async def test_file_uploading_wrong_content_format( faker: Faker, authorized_client: TestClient, uncategorized_file_content: bytes, file_upload_storage_token: str, + file_input_data: FileInputData, ) -> None: assert_response( authorized_client.post( - "/api/protected/storage-service/v2/file-kinds/image/files/", + f"/api/protected/storage-service/v2/file-kinds/{file_input_data.kind}/files/", headers={"X-Storage-Token": file_upload_storage_token}, files={ "upload": ( - faker.file_name(extension="webp"), + file_input_data.name, uncategorized_file_content, - "image/webp", + file_input_data.content_type, ) }, ),