Implement people and media quotas (fixes #365) (#366)

* Implement people and media quotas

* Fix tests

* Add database migration

* Check people quota in transactions

Author: DavidMStraub

alembic_users/versions/66e56620891a_add_trees_table.py

@@ -0,0 +1,35 @@
+"""Add trees table
+
+Revision ID: 66e56620891a
+Revises: e176543c72a8
+Create Date: 2023-05-05 22:48:14.628117
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = '66e56620891a'
+down_revision = 'e176543c72a8'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table('trees',
+    sa.Column('id', sa.String(), nullable=False),
+    sa.Column('quota_media', sa.Integer(), nullable=True),
+    sa.Column('quota_people', sa.Integer(), nullable=True),
+    sa.Column('usage_media', sa.Integer(), nullable=True),
+    sa.Column('usage_people', sa.Integer(), nullable=True),
+    sa.PrimaryKeyConstraint('id')
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_table('trees')
+    # ### end Alembic commands ###

gramps_webapi/api/file.py

@@ -60,6 +60,10 @@ def get_file_object(self) -> BinaryIO:
         """Return a binary file object."""
         raise NotImplementedError
 
+    def get_file_size(self) -> int:
+        """Return the file size in bytes."""
+        raise NotImplementedError
+
     def file_exists(self) -> bool:
         """Check if the file exists."""
         raise NotImplementedError
@@ -140,6 +144,17 @@ def get_file_object(self) -> BinaryIO:
             stream = BytesIO(f.read())
         return stream
 
+    def get_file_size(self) -> int:
+        """Return the file size in bytes.
+
+        Assumes the file exists!
+        """
+        try:
+            self._check_path()
+        except ValueError:
+            abort(403)
+        return os.path.getsize(self.path_abs)
+
     def send_file(
         self, etag: Optional[str] = None, download: bool = False, filename: str = ""
     ):
@@ -215,13 +230,14 @@ def get_checksum(fp) -> str:
     return md5sum
 
 
-def process_file(stream: Union[Any, BinaryIO]) -> Tuple[str, BinaryIO]:
+def process_file(stream: Union[Any, BinaryIO]) -> Tuple[str, int, BinaryIO]:
     """Process a file from a stream that has a read method."""
     fp = BytesIO()
     fp.write(stream.read())
     fp.seek(0)
     checksum = get_checksum(fp)
+    size = fp.tell()
     if not checksum:
         raise IOError("Unable to process file.")
     fp.seek(0)
-    return checksum, fp
+    return checksum, size, fp

gramps_webapi/api/media.py

@@ -4,15 +4,20 @@
 from pathlib import Path
 from typing import BinaryIO, List, Optional, Set
 
-from flask import current_app
+from flask import abort, current_app
 from gramps.gen.lib import Media
 from gramps.gen.utils.file import expand_media_path
 
+from ..auth import get_tree_usage, set_tree_usage
 from ..types import FilenameOrPath
 from ..util import get_extension
 from .file import FileHandler, LocalFileHandler, upload_file_local
-from .s3 import ObjectStorageFileHandler, list_object_keys, upload_file_s3
-from .util import get_db_handle
+from .s3 import (
+    ObjectStorageFileHandler,
+    get_object_keys_size,
+    upload_file_s3,
+)
+from .util import get_db_handle, get_tree_from_jwt
 
 
 PREFIX_S3 = "s3://"
@@ -60,6 +65,10 @@ def filter_existing_files(self, objects: List[Media]) -> List[Media]:
         """Given a list of media objects, return the ones with existing files."""
         raise NotImplementedError
 
+    def get_media_size(self) -> int:
+        """Return the total disk space used by all existing media objects."""
+        raise NotImplementedError
+
 
 class MediaHandlerLocal(MediaHandlerBase):
     """Handler for local media files."""
@@ -80,7 +89,7 @@ def upload_file(
             if Path(path).is_absolute():
                 # Don't allow absolute paths! This will raise
                 # if path is not relative to base_dir
-                rel_path: FilenameOrPath = Path(path).relative_to(base_dir)
+                rel_path: FilenameOrPath = Path(path).relative_to(self.base_dir)
             else:
                 rel_path = path
             upload_file_local(self.base_dir, rel_path, stream)
@@ -94,6 +103,29 @@ def filter_existing_files(self, objects: List[Media]) -> List[Media]:
             obj for obj in objects if self.get_file_handler(obj.handle).file_exists()
         ]
 
+    def get_media_size(self) -> int:
+        """Return the total disk space used by all existing media objects.
+
+        Only works with a request context.
+        """
+        if not os.path.isdir(self.base_dir):
+            raise ValueError(f"Directory {self.base_dir} does not exist")
+        size = 0
+        paths_seen = set()
+        db_handle = get_db_handle()
+        for obj in db_handle.iter_media():
+            path = obj.path
+            if os.path.isabs(path):
+                if Path(self.base_dir).resolve() not in Path(path).resolve().parents:
+                    continue  # file outside base dir - ignore
+            else:
+                path = os.path.join(self.base_dir, path)
+            if Path(path).is_file() and path not in paths_seen:
+                file_size = os.path.getsize(path)
+                size += file_size
+                paths_seen.add(path)
+        return size
+
 
 class MediaHandlerS3(MediaHandlerBase):
     """Generic handler for object storage media files."""
@@ -124,7 +156,9 @@ def prefix(self) -> Optional[str]:
 
     def get_remote_keys(self) -> Set[str]:
         """Return the set of all object keys that are known to exist on remote."""
-        keys = list_object_keys(self.bucket_name, endpoint_url=self.endpoint_url)
+        keys = get_object_keys_size(
+            self.bucket_name, prefix=self.prefix, endpoint_url=self.endpoint_url
+        )
         return set(removeprefix(key, self.prefix or "").lstrip("/") for key in keys)
 
     def get_file_handler(self, handle) -> ObjectStorageFileHandler:
@@ -160,6 +194,17 @@ def filter_existing_files(self, objects: List[Media]) -> List[Media]:
         remote_keys = self.get_remote_keys()
         return [obj for obj in objects if obj.checksum in remote_keys]
 
+    def get_media_size(self) -> int:
+        """Return the total disk space used by all existing media objects."""
+        db_handle = get_db_handle()
+        keys = set(obj.checksum for obj in db_handle.iter_media())
+        keys_size = get_object_keys_size(
+            bucket_name=self.bucket_name,
+            prefix=self.prefix,
+            endpoint_url=self.endpoint_url,
+        )
+        return sum(keys_size.get(key, 0) for key in keys)
+
 
 def MediaHandler(base_dir: Optional[str]) -> MediaHandlerBase:
     """Return an appropriate media handler."""
@@ -193,3 +238,27 @@ def get_media_handler(tree: Optional[str] = None) -> MediaHandlerBase:
             # construct subdirectory using OS dependent path join
             base_dir = os.path.join(base_dir, prefix)
     return MediaHandler(base_dir)
+
+
+def update_usage_media() -> int:
+    """Update the usage of media."""
+    tree = get_tree_from_jwt()
+    media_handler = get_media_handler(tree=tree)
+    usage_media = media_handler.get_media_size()
+    set_tree_usage(tree, usage_media=usage_media)
+    return usage_media
+
+
+def check_quota_media(to_add: int) -> None:
+    """Check whether the quota allows adding `to_add` bytes and abort if not."""
+    tree = get_tree_from_jwt()
+    usage_dict = get_tree_usage(tree)
+    if not usage_dict or usage_dict.get("usage_media") is None:
+        update_usage_media()
+    usage_dict = get_tree_usage(tree)
+    usage = usage_dict["usage_media"]
+    quota = usage_dict.get("quota_media")
+    if quota is None:
+        return
+    if usage + to_add > quota:
+        abort(405)

gramps_webapi/api/resources/base.py

@@ -33,14 +33,17 @@
 from gramps.gen.utils.grampslocale import GrampsLocale
 from webargs import fields, validate
 
+from ...auth import get_tree_usage, set_tree_usage
 from ...auth.const import PERM_ADD_OBJ, PERM_DEL_OBJ, PERM_EDIT_OBJ
 from ..auth import require_permissions
 from ..search import SearchIndexer
 from ..util import (
+    check_quota_people,
     get_db_handle,
     get_locale_for_language,
     get_search_indexer,
     get_tree_from_jwt,
+    update_usage_people,
     use_args,
 )
 from . import ProtectedResource, Resource
@@ -250,6 +253,9 @@ def delete(self, handle: str) -> Response:
         trans_dict = delete_object(
             self.db_handle_writable, handle, self.gramps_class_name
         )
+        # update usage
+        if self.gramps_class_name == "Person":
+            update_usage_people()
         # update search index
         tree = get_tree_from_jwt()
         indexer: SearchIndexer = get_search_indexer(tree)
@@ -420,6 +426,9 @@ def get(self, args: Dict) -> Response:
     def post(self) -> Response:
         """Post a new object."""
         require_permissions([PERM_ADD_OBJ])
+        # check quota
+        if self.gramps_class_name == "Person":
+            check_quota_people(to_add=1)
         obj = self._parse_object()
         if not obj:
             abort(400)
@@ -430,6 +439,9 @@ def post(self) -> Response:
             except ValueError:
                 abort(400)
             trans_dict = transaction_to_json(trans)
+        # update usage
+        if self.gramps_class_name == "Person":
+            update_usage_people()
         # update search index
         tree = get_tree_from_jwt()
         indexer: SearchIndexer = get_search_indexer(tree)

gramps_webapi/api/resources/file.py

@@ -34,7 +34,7 @@
 from ...auth.const import PERM_EDIT_OBJ
 from ..auth import require_permissions
 from ..file import process_file
-from ..media import get_media_handler
+from ..media import check_quota_media, get_media_handler, update_usage_media
 from ..util import get_db_handle, get_tree_from_jwt, use_args
 from . import ProtectedResource
 from .util import transaction_to_json, update_object
@@ -86,21 +86,28 @@ def put(self, args: Dict, handle: str) -> Response:
         mime = request.content_type
         if not mime:
             abort(HTTPStatus.NOT_ACCEPTABLE)
-        checksum, f = process_file(request.stream)
+        checksum, size, f = process_file(request.stream)
         tree = get_tree_from_jwt()
         media_handler = get_media_handler(tree)
+        file_handler = media_handler.get_file_handler(handle)
         if checksum == obj.checksum:
-            file_handler = media_handler.get_file_handler(handle)
             if not args.get("uploadmissing") or file_handler.file_exists():
                 # don't allow PUTting if the file didn't change
                 abort(HTTPStatus.CONFLICT)
             # we're uploading a missing file!
+            # new size will add to the quota
+            check_quota_media(to_add=size)
             # use existing path
             path = obj.get_path()
             media_handler.upload_file(f, checksum, mime, path=path)
             return Response(status=200)
         if args.get("uploadmissing"):
             abort(HTTPStatus.CONFLICT)
+        # we're updating an existing file
+        size_old = file_handler.get_file_size()
+        size_delta = size - size_old
+        if size_delta > 0:
+            check_quota_media(to_add=size_delta)
         media_handler.upload_file(f, checksum, mime)
         obj.set_checksum(checksum)
         path = media_handler.get_default_filename(checksum, mime)
@@ -114,6 +121,7 @@ def put(self, args: Dict, handle: str) -> Response:
             except ValueError:
                 abort(400)
             trans_dict = transaction_to_json(trans)
+        update_usage_media()
         return Response(
             response=json.dumps(trans_dict), status=200, mimetype="application/json"
         )

gramps_webapi/api/resources/media.py

@@ -33,7 +33,7 @@
 from ...auth.const import PERM_ADD_OBJ
 from ..auth import require_permissions
 from ..file import process_file
-from ..media import get_media_handler
+from ..media import check_quota_media, get_media_handler, update_usage_media
 from ..util import get_tree_from_jwt
 from .base import (
     GrampsObjectProtectedResource,
@@ -79,7 +79,8 @@ def post(self) -> Response:
         mime = request.content_type
         if not mime:
             abort(HTTPStatus.NOT_ACCEPTABLE)
-        checksum, f = process_file(request.stream)
+        checksum, size, f = process_file(request.stream)
+        check_quota_media(to_add=size)
         tree = get_tree_from_jwt()
         media_handler = get_media_handler(tree)
         media_handler.upload_file(f, checksum, mime)
@@ -95,4 +96,5 @@ def post(self) -> Response:
             except ValueError:
                 abort(400)
             trans_dict = transaction_to_json(trans)
+        update_usage_media()
         return self.response(201, trans_dict, total_items=len(trans_dict))

gramps_webapi/api/resources/objects.py

@@ -24,14 +24,19 @@
 
 from flask import Response, abort, current_app, request
 from gramps.gen.db import DbTxn
-from gramps.gen.lib import Family
+from gramps.gen.lib import Family, Person
 from gramps.gen.lib.primaryobj import BasicPrimaryObject as GrampsObject
 from gramps.gen.lib.serialize import from_json
 
 from ...auth.const import PERM_ADD_OBJ, PERM_EDIT_OBJ
 from ..auth import require_permissions
 from ..search import SearchIndexer
-from ..util import get_db_handle, get_search_indexer, get_tree_from_jwt
+from ..util import (
+    check_quota_people,
+    get_db_handle,
+    get_search_indexer,
+    get_tree_from_jwt,
+)
 from . import ProtectedResource
 from .util import add_object, fix_object_dict, transaction_to_json, validate_object_dict
 
@@ -66,6 +71,8 @@ def post(self) -> Response:
             require_permissions([PERM_EDIT_OBJ])
         if not objects:
             abort(400)
+        number_new_people = sum(isinstance(obj, Person) for obj in objects)
+        check_quota_people(to_add=number_new_people)
         db_handle = get_db_handle(readonly=False)
         with DbTxn("Add objects", db_handle) as trans:
             for obj in objects:

gramps_webapi/api/resources/transactions.py

@@ -33,7 +33,12 @@
 from ...auth.const import PERM_ADD_OBJ, PERM_DEL_OBJ, PERM_EDIT_OBJ
 from ..auth import require_permissions
 from ..search import SearchIndexer
-from ..util import get_db_handle, get_search_indexer, get_tree_from_jwt
+from ..util import (
+    check_quota_people,
+    get_db_handle,
+    get_search_indexer,
+    get_tree_from_jwt,
+)
 from . import ProtectedResource
 from .util import transaction_to_json
 
@@ -50,6 +55,10 @@ def post(self) -> Response:
         if not payload:
             abort(400)  # disallow empty payload
         db_handle = get_db_handle(readonly=False)
+        new_people = sum(
+            item["type"] == "add" and item["_class"] == "Person" for item in payload
+        )
+        check_quota_people(to_add=new_people)
         with DbTxn("Raw transaction", db_handle) as trans:
             for item in payload:
                 try: