fastapi-rtk 1.0.19__py3-none-any.whl → 1.0.20__py3-none-any.whl

fastapi_rtk/__init__.py

@@ -130,7 +130,9 @@ __all__ = [
     # .bases
     "DBQueryParams",
     "AbstractQueryBuilder",
+    "BaseFileException",
     "FileNotAllowedException",
+    "FileTooLargeException",
     "AbstractFileManager",
     "AbstractImageManager",
     "AbstractBaseFilter",
@@ -231,13 +233,14 @@ __all__ = [
     "ExtenderMixin",
     "uuid_namegen",
     "secure_filename",
+    "prettify_dict",
+    "format_file_size",
     "hooks",
     "lazy",
     "lazy_import",
     "lazy_self",
     "merge_schema",
     "multiple_async_contexts",
-    "prettify_dict",
     "generate_schema_from_typed_dict",
     "get_pydantic_model_field",
     "smart_run",

fastapi_rtk/_version.py

@@ -1 +1 @@
-__version__ = "1.0.19"
+__version__ = "1.0.20"

fastapi_rtk/api/model_rest_api.py

@@ -47,6 +47,7 @@ from ..utils import (
     SelfType,
     T,
     deep_merge,
+    format_file_size,
     lazy_self,
     merge_schema,
     smart_run,
@@ -1681,29 +1682,36 @@ class ModelRestApi(BaseApi):
             If you are overriding this method, make sure to copy all the decorators too.
         """
         async with AsyncTaskRunner():
-            async with AsyncTaskRunner():
-                body_json = await smart_run(
-                    self._process_body,
-                    session,
-                    body,
-                    self.add_query_rel_fields,
-                    self.add_schema_extra_fields.keys()
-                    if self.add_schema_extra_fields
-                    else None,
-                )
-                item = self.datamodel.obj(**body_json)
-                pre_add = await smart_run(
-                    self.pre_add,
-                    item,
-                    PARAM_BODY_SESSION(body=body, session=session),
-                )
-                if pre_add is not None:
-                    if isinstance(pre_add, Model):
-                        item = pre_add
-                    else:
-                        AsyncTaskRunner.remove_tasks_by_tag("file")
-                        return pre_add
+            async with AsyncTaskRunner(
+                "after_commit", run_tasks_even_if_exception=True
+            ) as after_commit_runner:
+                async with AsyncTaskRunner("before_commit") as before_commit_runner:
+                    body_json = await smart_run(
+                        self._process_body,
+                        session,
+                        body,
+                        self.add_query_rel_fields,
+                        self.add_schema_extra_fields.keys()
+                        if self.add_schema_extra_fields
+                        else None,
+                    )
+                    item = self.datamodel.obj(**body_json)
+                    pre_add = await smart_run(
+                        self.pre_add,
+                        item,
+                        PARAM_BODY_SESSION(body=body, session=session),
+                    )
+                    if pre_add is not None:
+                        if isinstance(pre_add, Model):
+                            item = pre_add
+                        else:
+                            before_commit_runner.remove_tasks_by_tag("file")
+                            after_commit_runner.remove_tasks_by_tag("file")
+                            return pre_add
                 item = await smart_run(self.datamodel.add, session, item)
+                after_commit_runner.remove_tasks_by_tag(
+                    "file"
+                )  # Delete any file tasks scheduled to revert files on error
             post_add = await smart_run(
                 self.post_add,
                 item,
@@ -1739,50 +1747,57 @@ class ModelRestApi(BaseApi):
             If you are overriding this method, make sure to copy all the decorators too.
         """
         async with AsyncTaskRunner():
-            async with AsyncTaskRunner():
-                item = await smart_run(
-                    self.datamodel.get_one,
-                    session,
-                    params={
-                        "list_columns": self.show_select_columns,
-                        "where_id": id,
-                        "filter_classes": self.base_filters,
-                        "opr_filter_classes": self.base_opr_filters,
-                    },
-                )
-                if not item:
-                    raise HTTPException(
-                        fastapi.status.HTTP_404_NOT_FOUND, ErrorCode.ITEM_NOT_FOUND
+            async with AsyncTaskRunner(
+                "after_commit", run_tasks_even_if_exception=True
+            ) as after_commit_runner:
+                async with AsyncTaskRunner("before_commit") as before_commit_runner:
+                    item = await smart_run(
+                        self.datamodel.get_one,
+                        session,
+                        params={
+                            "list_columns": self.show_select_columns,
+                            "where_id": id,
+                            "filter_classes": self.base_filters,
+                            "opr_filter_classes": self.base_opr_filters,
+                        },
                     )
-                body_json = await smart_run(
-                    self._process_body,
-                    session,
-                    body,
-                    self.edit_query_rel_fields,
-                    self.edit_schema_extra_fields.keys()
-                    if self.edit_schema_extra_fields
-                    else None,
-                    item=item,
-                )
-                await smart_run(
-                    self.pre_update_merge,
-                    item,
-                    body_json,
-                    PARAM_BODY_SESSION(body=body, session=session),
-                )
-                item.update(body_json)
-                pre_update = await smart_run(
-                    self.pre_update,
-                    item,
-                    PARAM_BODY_SESSION(body=body, session=session),
-                )
-                if pre_update is not None:
-                    if isinstance(pre_update, Model):
-                        item = pre_update
-                    else:
-                        AsyncTaskRunner.remove_tasks_by_tag("file")
-                        return pre_update
+                    if not item:
+                        raise HTTPException(
+                            fastapi.status.HTTP_404_NOT_FOUND, ErrorCode.ITEM_NOT_FOUND
+                        )
+                    body_json = await smart_run(
+                        self._process_body,
+                        session,
+                        body,
+                        self.edit_query_rel_fields,
+                        self.edit_schema_extra_fields.keys()
+                        if self.edit_schema_extra_fields
+                        else None,
+                        item=item,
+                    )
+                    await smart_run(
+                        self.pre_update_merge,
+                        item,
+                        body_json,
+                        PARAM_BODY_SESSION(body=body, session=session),
+                    )
+                    item.update(body_json)
+                    pre_update = await smart_run(
+                        self.pre_update,
+                        item,
+                        PARAM_BODY_SESSION(body=body, session=session),
+                    )
+                    if pre_update is not None:
+                        if isinstance(pre_update, Model):
+                            item = pre_update
+                        else:
+                            before_commit_runner.remove_tasks_by_tag("file")
+                            after_commit_runner.remove_tasks_by_tag("file")
+                            return pre_update
                 item = await smart_run(self.datamodel.edit, session, item)
+                after_commit_runner.remove_tasks_by_tag(
+                    "file"
+                )  # Delete any file tasks scheduled to revert files on error
             post_update = await smart_run(
                 self.post_update,
                 item,
@@ -1817,60 +1832,75 @@ class ModelRestApi(BaseApi):
             If you are overriding this method, make sure to copy all the decorators too.
         """
         async with AsyncTaskRunner():
-            async with AsyncTaskRunner():
-                item = await smart_run(
-                    self.datamodel.get_one,
-                    session,
-                    params={
-                        "list_columns": self.show_select_columns,
-                        "where_id": id,
-                        "filter_classes": self.base_filters,
-                        "opr_filter_classes": self.base_opr_filters,
-                    },
-                )
-                if not item:
-                    raise HTTPException(
-                        fastapi.status.HTTP_404_NOT_FOUND, ErrorCode.ITEM_NOT_FOUND
+            async with AsyncTaskRunner(
+                "after_commit", run_tasks_even_if_exception=True
+            ) as after_commit_runner:
+                async with AsyncTaskRunner("before_commit") as before_commit_runner:
+                    item = await smart_run(
+                        self.datamodel.get_one,
+                        session,
+                        params={
+                            "list_columns": self.show_select_columns,
+                            "where_id": id,
+                            "filter_classes": self.base_filters,
+                            "opr_filter_classes": self.base_opr_filters,
+                        },
                     )
-                pre_delete = await smart_run(
-                    self.pre_delete,
-                    item,
-                    PARAM_ID_SESSION(id=id, session=session),
-                )
-                if pre_delete is not None:
-                    if isinstance(pre_delete, Model):
-                        item = pre_delete
-                    else:
-                        return pre_delete
-                # Delete all related files and images
-                file_and_image_columns = [
-                    x
-                    for x in self.datamodel.get_file_column_list()
-                    + self.datamodel.get_image_column_list()
-                ]
-                schema = self.datamodel.generate_schema(
-                    file_and_image_columns,
-                    with_id=False,
-                    with_name=False,
-                    with_property=False,
-                )
-                schema_data = schema.model_validate(item, from_attributes=True)
-                for column in file_and_image_columns:
-                    fm = (
-                        self.datamodel.file_manager
-                        if self.datamodel.is_file(column)
-                        else self.datamodel.image_manager
+                    if not item:
+                        raise HTTPException(
+                            fastapi.status.HTTP_404_NOT_FOUND, ErrorCode.ITEM_NOT_FOUND
+                        )
+                    pre_delete = await smart_run(
+                        self.pre_delete,
+                        item,
+                        PARAM_ID_SESSION(id=id, session=session),
                     )
-                    filenames = getattr(schema_data, column, None) or []
-                    if not isinstance(filenames, list):
-                        filenames = [filenames]
-                    for filename in filenames:
-                        AsyncTaskRunner.add_task(
-                            lambda fm=fm, filename=filename: smart_run(
-                                fm.delete_file, filename
-                            )
+                    if pre_delete is not None:
+                        if isinstance(pre_delete, Model):
+                            item = pre_delete
+                        else:
+                            return pre_delete
+                    # Delete all related files and images
+                    file_and_image_columns = [
+                        x
+                        for x in self.datamodel.get_file_column_list()
+                        + self.datamodel.get_image_column_list()
+                    ]
+                    schema = self.datamodel.generate_schema(
+                        file_and_image_columns,
+                        with_id=False,
+                        with_name=False,
+                        with_property=False,
+                    )
+                    schema_data = schema.model_validate(item, from_attributes=True)
+                    for column in file_and_image_columns:
+                        fm = (
+                            self.datamodel.file_manager
+                            if self.datamodel.is_file(column)
+                            else self.datamodel.image_manager
                         )
+                        filenames = getattr(schema_data, column, None) or []
+                        if not isinstance(filenames, list):
+                            filenames = [filenames]
+                        for filename in filenames:
+                            old_content = await smart_run(fm.get_file, filename)
+                            before_commit_runner.add_task(
+                                lambda fm=fm, filename=filename: smart_run(
+                                    fm.delete_file, filename
+                                )
+                            )
+                            after_commit_runner.add_task(
+                                lambda fm=fm,
+                                content=old_content,
+                                filename=filename: smart_run(
+                                    fm.save_content_to_file, content, filename
+                                ),
+                                tags=["file"],
+                            )
                 await smart_run(self.datamodel.delete, session, item)
+                after_commit_runner.remove_tasks_by_tag(
+                    "file"
+                )  # Delete any file tasks scheduled to revert files on error
             post_delete = await smart_run(
                 self.post_delete,
                 item,
@@ -2371,15 +2401,30 @@ class ModelRestApi(BaseApi):
                         )
                         if item and hasattr(item, key) and getattr(item, key):
                             actual_old_filenames = getattr(item, key)
+                            before_commit_runner = AsyncTaskRunner.get_runner(
+                                "before_commit"
+                            )
+                            after_commit_runner = AsyncTaskRunner.get_runner(
+                                "after_commit"
+                            )
                             # Delete only the files or images that are not in the new old_filenames
                             for filename in actual_old_filenames:
                                 if filename not in old_filenames:
-                                    AsyncTaskRunner.add_task(
+                                    old_content = await smart_run(fm.get_file, filename)
+                                    before_commit_runner.add_task(
                                         lambda fm=fm, old_filename=filename: smart_run(
                                             fm.delete_file, old_filename
                                         ),
                                         tags=["file"],
                                     )
+                                    after_commit_runner.add_task(
+                                        lambda fm=fm,
+                                        content=old_content,
+                                        filename=filename: smart_run(
+                                            fm.save_content_to_file, content, filename
+                                        ),
+                                        tags=["file"],
+                                    )
 
                         new_filenames = []
                         # Loop through value instead of only file values so the order is maintained
@@ -2395,12 +2440,27 @@ class ModelRestApi(BaseApi):
                         # Delete existing file or image if it is being updated
                         if item and hasattr(item, key) and getattr(item, key):
                             filename = getattr(item, key)
-                            AsyncTaskRunner.add_task(
+                            old_content = await smart_run(fm.get_file, filename)
+                            before_commit_runner = AsyncTaskRunner.get_runner(
+                                "before_commit"
+                            )
+                            after_commit_runner = AsyncTaskRunner.get_runner(
+                                "after_commit"
+                            )
+                            before_commit_runner.add_task(
                                 lambda fm=fm, old_filename=filename: smart_run(
                                     fm.delete_file, old_filename
                                 ),
                                 tags=["file"],
                             )
+                            after_commit_runner.add_task(
+                                lambda fm=fm,
+                                content=old_content,
+                                filename=filename: smart_run(
+                                    fm.save_content_to_file, content, filename
+                                ),
+                                tags=["file"],
+                            )
 
                         # Only process if the value exists and is not None
                         if value:
@@ -2501,12 +2561,30 @@ class ModelRestApi(BaseApi):
                 ),
             )
         content = await file.read()
-        AsyncTaskRunner.add_task(
+        if not fm.is_file_size_allowed(content):
+            raise HTTPWithValidationException(
+                fastapi.status.HTTP_400_BAD_REQUEST,
+                "bytes_too_long",
+                "body",
+                key,
+                translate(
+                    "File size from '{filename}' exceeds the allowed limit {file_size_limit}.",
+                    filename=file.filename,
+                    file_size_limit=format_file_size(fm.max_file_size),
+                ),
+            )
+        before_commit_runner = AsyncTaskRunner.get_runner("before_commit")
+        after_commit_runner = AsyncTaskRunner.get_runner("after_commit")
+        before_commit_runner.add_task(
             lambda fm=fm, content=content, new_name=new_name: smart_run(
                 fm.save_content_to_file, content, new_name
             ),
             tags=["file"],
         )
+        after_commit_runner.add_task(
+            lambda fm=fm, new_name=new_name: smart_run(fm.delete_file, new_name),
+            tags=["file"],
+        )
         return new_name
 
     """

fastapi_rtk/bases/__init__.py

@@ -8,7 +8,9 @@ from .session import *
 __all__ = [
     "DBQueryParams",
     "AbstractQueryBuilder",
+    "BaseFileException",
     "FileNotAllowedException",
+    "FileTooLargeException",
     "AbstractFileManager",
     "AbstractImageManager",
     "AbstractBaseFilter",

fastapi_rtk/bases/file_manager.py

@@ -6,10 +6,22 @@ import fastapi
 from ..exceptions import FastAPIReactToolkitException
 from ..utils import hooks, lazy, uuid_namegen
 
-__all__ = ["FileNotAllowedException", "AbstractFileManager", "AbstractImageManager"]
+__all__ = [
+    "BaseFileException",
+    "FileNotAllowedException",
+    "FileTooLargeException",
+    "AbstractFileManager",
+    "AbstractImageManager",
+]
 
 
-class FileNotAllowedException(FastAPIReactToolkitException):
+class BaseFileException(FastAPIReactToolkitException):
+    """
+    Base exception class for file-related errors.
+    """
+
+
+class FileNotAllowedException(BaseFileException):
     """
     Exception raised when a file is not allowed based on its extension.
     """
@@ -18,13 +30,25 @@ class FileNotAllowedException(FastAPIReactToolkitException):
         super().__init__(f"File '{filename}' is not allowed.")
 
 
+class FileTooLargeException(BaseFileException):
+    """
+    Exception raised when a file exceeds the maximum allowed size.
+    """
+
+    def __init__(self, filename: str, max_size: int):
+        super().__init__(
+            f"File '{filename}' exceeds the maximum allowed size of {max_size} bytes."
+        )
+
+
 class AbstractFileManager(abc.ABC):
     """
     Abstract base class for file managers.
     """
 
     base_path: str = None
-    allowed_extensions: list[str] = None
+    allowed_extensions: list[str] | None = None
+    max_file_size: int | None = None
     namegen = lazy(lambda: uuid_namegen)
     permission = lazy(lambda: 0o755)
 
@@ -32,6 +56,7 @@ class AbstractFileManager(abc.ABC):
         self,
         base_path: str | None = None,
         allowed_extensions: list[str] | None = None,
+        max_file_size: int | None = None,
         namegen: typing.Callable[[str], str] | None = None,
         permission: int | None = None,
     ):
@@ -41,6 +66,7 @@ class AbstractFileManager(abc.ABC):
         Args:
             base_path (str | None, optional): Base path for file storage. Defaults to None.
             allowed_extensions (list[str] | None, optional): Allowed file extensions. Defaults to None.
+            max_file_size (int | None, optional): Maximum file size allowed. Defaults to None.
             namegen (typing.Callable[[str], str] | None, optional): Callable for generating file names. Defaults to None.
             permission (int | None, optional): File permission settings. Defaults to None.
 
@@ -51,6 +77,8 @@ class AbstractFileManager(abc.ABC):
             self.base_path = base_path
         if allowed_extensions is not None:
             self.allowed_extensions = allowed_extensions
+        if max_file_size is not None:
+            self.max_file_size = max_file_size
         if namegen is not None:
             self.namegen = namegen
         if permission is not None:
@@ -66,15 +94,26 @@ class AbstractFileManager(abc.ABC):
 
     def __init_subclass__(cls):
         # Add pre-hook to save_file and save_content_to_file to check if the file is allowed
-        def check_is_file_allowed(self, *args, **kwargs):
-            filename = None
-            if "filename" in kwargs:
-                filename = kwargs["filename"]
-            elif len(args) > 1:
-                filename = args[1]
-            if filename and not self.is_filename_allowed(filename):
+        def check_is_file_allowed(self: typing.Self, *args, **kwargs):
+            filename = kwargs.get("filename", args[1] if len(args) > 1 else None)
+            if not filename:
+                raise ValueError("Filename must be provided.")
+            if not self.is_filename_allowed(filename):
                 raise FileNotAllowedException(filename)
 
+        # Add pre-hook to save_file and save_content_to_file to check if the file size is allowed
+        def check_is_file_size_allowed(self: typing.Self, *args, **kwargs):
+            input = kwargs.get(
+                "file_data", args[0] if len(args) > 0 else None
+            ) or kwargs.get("content", args[0] if len(args) > 0 else None)
+            filename = kwargs.get("filename", args[1] if len(args) > 1 else None)
+            if not input or not filename:
+                raise ValueError(
+                    "Both filename and file data/content must be provided."
+                )
+            if not self.is_file_size_allowed(input):
+                raise FileTooLargeException(filename, self.max_file_size)
+
         if cls.save_file is not AbstractFileManager.save_file:
             wrapped_save_file = hooks(pre=check_is_file_allowed)(cls.save_file)
             cls.save_file = wrapped_save_file
@@ -83,6 +122,16 @@ class AbstractFileManager(abc.ABC):
                 cls.save_content_to_file
             )
             cls.save_content_to_file = wrapped_save_content_to_file
+        if cls.save_file is not AbstractFileManager.save_file:
+            wrapped_save_file_size = hooks(pre=check_is_file_size_allowed)(
+                cls.save_file
+            )
+            cls.save_file = wrapped_save_file_size
+        if cls.save_content_to_file is not AbstractFileManager.save_content_to_file:
+            wrapped_save_content_size = hooks(pre=check_is_file_size_allowed)(
+                cls.save_content_to_file
+            )
+            cls.save_content_to_file = wrapped_save_content_size
 
     """
     --------------------------------------------------------------------------------------------------------
@@ -243,6 +292,29 @@ class AbstractFileManager(abc.ABC):
             and filename.rsplit(".", 1)[1].lower() in self.allowed_extensions
         )
 
+    def is_file_size_allowed(self, file_size: int | fastapi.UploadFile | bytes | str):
+        """
+        Check if a file size is allowed based on the maximum file size.
+
+        Args:
+            file_size (int | fastapi.UploadFile | bytes | str): The size of the file in bytes or the file data.
+
+        Returns:
+            bool: True if the file size is allowed, False otherwise.
+        """
+        if self.max_file_size is None:
+            return True
+
+        size = 0
+        if isinstance(file_size, int):
+            size = file_size
+        elif isinstance(file_size, fastapi.UploadFile):
+            size = file_size.size or 0
+        elif isinstance(file_size, (bytes, str)):
+            size = len(file_size)
+
+        return size <= self.max_file_size
+
     def generate_name(self, filename: str) -> str:
         """
         Generates a name for the given file data.

fastapi_rtk/file_managers/file_manager.py

@@ -16,6 +16,7 @@ class FileManager(AbstractFileManager):
 
     base_path = lazy(lambda: Setting.UPLOAD_FOLDER)
     allowed_extensions = lazy(lambda: Setting.FILE_ALLOWED_EXTENSIONS)
+    max_file_size = lazy(lambda: Setting.FILE_MAX_SIZE)
 
     def post_init(self):
         if not self.base_path:

fastapi_rtk/file_managers/image_manager.py

@@ -13,6 +13,7 @@ class ImageManager(FileManager, AbstractImageManager):
 
     base_path = lazy(lambda: Setting.IMG_UPLOAD_FOLDER)
     allowed_extensions = lazy(lambda: Setting.IMG_ALLOWED_EXTENSIONS)
+    max_file_size = lazy(lambda: Setting.IMG_MAX_SIZE)
 
     def post_init(self):
         if not self.base_path: