Check people quota before import (#381)

Author: DavidMStraub

gramps_webapi/api/media.py

@@ -38,7 +38,7 @@
     get_object_keys_size,
     upload_file_s3,
 )
-from .util import get_db_handle, get_tree_from_jwt
+from .util import abort_with_message, get_db_handle, get_tree_from_jwt
 
 
 PREFIX_S3 = "s3://"
@@ -332,9 +332,10 @@ def update_usage_media() -> int:
     return usage_media
 
 
-def check_quota_media(to_add: int) -> None:
+def check_quota_media(to_add: int, tree: Optional[str] = None) -> None:
     """Check whether the quota allows adding `to_add` bytes and abort if not."""
-    tree = get_tree_from_jwt()
+    if not tree:
+        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()
@@ -344,4 +345,4 @@ def check_quota_media(to_add: int) -> None:
     if quota is None:
         return
     if usage + to_add > quota:
-        abort(405)
+        abort_with_message(405, "Not allowed by media quota")

gramps_webapi/api/resources/importers.py

@@ -25,13 +25,13 @@
 from http import HTTPStatus
 from typing import Any, Dict
 
-from flask import Response, abort, current_app, request
+from flask import Response, current_app, request
 from webargs import fields
 
 from ...auth.const import PERM_IMPORT_FILE
 from ..auth import require_permissions
 from ..tasks import AsyncResult, import_file, make_task_response, run_task
-from ..util import get_db_handle, get_tree_from_jwt, use_args
+from ..util import abort_with_message, get_db_handle, get_tree_from_jwt, use_args
 from . import ProtectedResource
 from .emit import GrampsJSONEncoder
 from .util import get_importers
@@ -56,7 +56,7 @@ def get(self, args: Dict[str, Any], extension: str) -> Response:
         get_db_handle()  # needed to load plugins
         importers = get_importers(extension)
         if not importers:
-            abort(404)
+            abort_with_message(404, f"Importer for extension {extension} not found")
         return self.response(200, importers[0])
 
 
@@ -81,9 +81,13 @@ def post(self, args: Dict, extension: str) -> Response:
         file_path = os.path.join(export_path, file_name)
         with open(file_path, "w+b") as ftmp:
             ftmp.write(request_stream.read())
+        if os.path.getsize(file_path) == 0:
+            abort_with_message(400, "Imported file is empty")
         importers = get_importers(extension.lower())
         if not importers:
-            abort(HTTPStatus.NOT_FOUND)
+            abort_with_message(
+                HTTPStatus.NOT_FOUND, f"Importer for extension {extension} not found"
+            )
         tree = get_tree_from_jwt()
         task = run_task(
             import_file,

gramps_webapi/api/resources/util.py

@@ -29,11 +29,12 @@
 
 import gramps
 import jsonschema
-from flask import abort, current_app
+from flask import current_app
 from gramps.gen.const import GRAMPS_LOCALE as glocale
 from gramps.gen.db import KEY_TO_CLASS_MAP, DbTxn
 from gramps.gen.db.base import DbReadBase, DbWriteBase
 from gramps.gen.db.dbconst import TXNADD, TXNDEL, TXNUPD
+from gramps.gen.db.utils import import_as_dict
 from gramps.gen.display.name import NameDisplay
 from gramps.gen.display.place import PlaceDisplay
 from gramps.gen.errors import HandleError
@@ -72,7 +73,7 @@
 from ...const import DISABLED_IMPORTERS, SEX_FEMALE, SEX_MALE, SEX_UNKNOWN
 from ...types import FilenameOrPath, Handle, TransactionJson
 from ..media import get_media_handler
-from ..util import get_db_handle, get_tree_from_jwt
+from ..util import abort_with_message, get_db_handle, get_tree_from_jwt
 
 pd = PlaceDisplay()
 _ = glocale.translation.gettext
@@ -797,7 +798,7 @@ def add_object(
     """
     if db_handle.readonly:
         # adding objects is forbidden on a read-only db!
-        abort(HTTPStatus.FORBIDDEN)
+        abort_with_message(HTTPStatus.FORBIDDEN, "Forbidden: database is read-only")
     obj_class = obj.__class__.__name__.lower()
     if fail_if_exists:
         if has_handle(db_handle, obj):
@@ -988,7 +989,7 @@ def update_object(
     """
     if db_handle.readonly:
         # updating objects is forbidden on a read-only db!
-        abort(HTTPStatus.FORBIDDEN)
+        abort_with_message(HTTPStatus.FORBIDDEN, "Forbidden: database is read-only")
     obj_class = obj.__class__.__name__.lower()
     if not has_handle(db_handle, obj):
         raise ValueError("Cannot be used for new objects.")
@@ -1185,5 +1186,26 @@ def run_import(
             if delete:
                 os.remove(file_name)
             if not result:
-                abort(500)
+                abort_with_message(500, "Import failed")
             return
+
+
+def dry_run_import(
+    file_name: FilenameOrPath,
+) -> Optional[Dict[str, int]]:
+    """Import a file into an in-memory database and returns object counts."""
+    db_handle: DbReadBase = import_as_dict(filename=file_name, user=User())
+    if db_handle is None:
+        return None
+    return {
+        "people": db_handle.get_number_of_people(),
+        "families": db_handle.get_number_of_families(),
+        "sources": db_handle.get_number_of_sources(),
+        "citations": db_handle.get_number_of_citations(),
+        "events": db_handle.get_number_of_events(),
+        "media": db_handle.get_number_of_media(),
+        "places": db_handle.get_number_of_places(),
+        "repositories": db_handle.get_number_of_repositories(),
+        "notes": db_handle.get_number_of_notes(),
+        "tags": db_handle.get_number_of_tags(),
+    }

gramps_webapi/api/tasks.py

@@ -33,8 +33,9 @@
 from .export import prepare_options, run_export
 from .media import get_media_handler
 from .report import run_report
-from .resources.util import run_import
+from .resources.util import dry_run_import, run_import
 from .util import (
+    check_quota_people,
     get_config,
     get_db_manager,
     get_db_outside_request,
@@ -126,6 +127,10 @@ def search_reindex_incremental(tree) -> None:
 @shared_task()
 def import_file(tree: str, file_name: str, extension: str, delete: bool = True):
     """Import a file."""
+    object_counts = dry_run_import(file_name=file_name)
+    if object_counts is None:
+        raise ValueError(f"Failed importing {file_name}")
+    check_quota_people(to_add=object_counts["people"], tree=tree)
     db_manager = get_db_manager(tree)
     db_handle = db_manager.get_db().db
     run_import(

gramps_webapi/api/util.py

@@ -21,6 +21,7 @@
 
 import hashlib
 import io
+import json
 import os
 import smtplib
 import socket
@@ -29,7 +30,7 @@
 from http import HTTPStatus
 from typing import BinaryIO, List, Optional, Sequence, Tuple
 
-from flask import abort, current_app, g, jsonify, make_response, request
+from flask import Response, abort, current_app, g, jsonify, make_response, request
 from flask_jwt_extended import get_jwt
 from gramps.cli.clidbman import NAME_FILE, CLIDbManager
 from gramps.gen.config import config
@@ -58,8 +59,16 @@ class Parser(FlaskParser):
 
     def handle_error(self, error, req, schema, *, error_status_code, error_headers):
         status_code = error_status_code or self.DEFAULT_VALIDATION_STATUS
+        pretty_message = "".join([c for c in str(error.messages) if c not in "{}[]()'"])
+        payload = {
+            "error": {
+                "code": status_code,
+                "message": pretty_message,
+                "messages": error.messages,
+            }
+        }
         abort(
-            make_response(jsonify(error.messages), status_code),
+            make_response(jsonify(payload), status_code),
             exc=error,
             messages=error.messages,
             schema=schema,
@@ -135,7 +144,9 @@ def get_db_outside_request(tree: str, view_private: bool, readonly: bool) -> DbR
     if not view_private:
         if not readonly:
             # requesting write access on a private proxy DB is impossible & forbidden!
-            abort(HTTPStatus.FORBIDDEN)
+            abort_with_message(
+                HTTPStatus.FORBIDDEN, "Cannot write to a private proxy database"
+            )
         # if we're not authorized to view private records,
         # return a proxy DB instead of the real one
         return ModifiedPrivateProxyDb(dbstate.db)
@@ -168,7 +179,9 @@ def get_db_handle(readonly: bool = True) -> DbReadBase:
     if not view_private:
         if not readonly:
             # requesting write access on a private proxy DB is impossible & forbidden!
-            abort(HTTPStatus.FORBIDDEN)
+            abort_with_message(
+                HTTPStatus.FORBIDDEN, "Cannot write to a private proxy database"
+            )
         return ModifiedPrivateProxyDb(g.db)
 
     if not readonly and "db_write" not in g:
@@ -218,7 +231,7 @@ def get_buffer_for_file(filename: str, delete=True, not_found=False) -> BinaryIO
     except FileNotFoundError:
         if not_found:
             raise FileNotFoundError
-        abort(500)
+        abort_with_message(500, "File not found")
     if delete:
         os.remove(filename)
     return buffer
@@ -282,7 +295,7 @@ def make_cache_key_thumbnails(*args, **kwargs):
     try:
         obj = db_handle.get_media_from_handle(handle)
     except HandleError:
-        abort(404)
+        abort_with_message(404, f"Handle {handle} not found")
     # checksum in the DB
     checksum = obj.checksum
 
@@ -319,7 +332,7 @@ def get_tree_id(guid: str) -> str:
     if not tree_id:
         if current_app.config["TREE"] == TREE_MULTI:
             # multi-tree support enabled but user has no tree ID: forbidden!
-            abort(403)
+            abort_with_message(403, "Forbidden")
         # needed for backwards compatibility: single-tree mode but user without tree ID
         dbmgr = WebDbManager(name=current_app.config["TREE"], create_if_missing=False)
         tree_id = dbmgr.dirname
@@ -343,25 +356,42 @@ def tree_exists(tree_id: str) -> bool:
     return True
 
 
-def update_usage_people() -> int:
+def update_usage_people(tree: Optional[str] = None) -> int:
     """Update the usage of people."""
-    tree = get_tree_from_jwt()
-    db_handle = get_db_handle()
+    if not tree:
+        tree = get_tree_from_jwt()
+    db_handle = get_db_outside_request(
+        tree=tree,
+        view_private=True,
+        readonly=True,
+    )
     usage_people = db_handle.get_number_of_people()
     set_tree_usage(tree, usage_people=usage_people)
     return usage_people
 
 
-def check_quota_people(to_add: int) -> None:
+def check_quota_people(to_add: int, tree: Optional[str] = None) -> None:
     """Check whether the quota allows adding `to_add` people and abort if not."""
-    tree = get_tree_from_jwt()
+    if not tree:
+        tree = get_tree_from_jwt()
     usage_dict = get_tree_usage(tree)
     if not usage_dict or usage_dict.get("usage_people") is None:
-        update_usage_people()
+        update_usage_people(tree=tree)
     usage_dict = get_tree_usage(tree)
     usage = usage_dict["usage_people"]
     quota = usage_dict.get("quota_people")
     if quota is None:
         return
     if usage + to_add > quota:
-        abort(405)
+        abort_with_message(405, "Not allowed by people quota")
+
+
+def abort_with_message(status: int, message: str):
+    """Abort with a JSON response."""
+    payload = {"error": {"code": status, "message": message}}
+    response = Response(
+        response=json.dumps(payload),
+        status=status,
+        mimetype="application/json",
+    )
+    abort(response)

tests/test_endpoints/test_importers.py

@@ -21,14 +21,15 @@
 """Tests for the /api/importers endpoints."""
 
 import io
+import os
 import unittest
 from unittest.mock import patch
 
 from gramps.cli.clidbman import CLIDbManager
 from gramps.gen.dbstate import DbState
 
 from gramps_webapi.app import create_app
-from gramps_webapi.auth import user_db, add_user
+from gramps_webapi.auth import add_user, set_tree_quota, user_db
 from gramps_webapi.auth.const import ROLE_EDITOR, ROLE_OWNER
 from gramps_webapi.const import ENV_CONFIG_FILE, TEST_EMPTY_GRAMPS_AUTH_CONFIG
 
@@ -100,7 +101,7 @@ def setUpClass(cls):
         """Test class setup."""
         cls.name = "empty"
         cls.dbman = CLIDbManager(DbState())
-        _, _name = cls.dbman.create_new_db_cli(cls.name, dbid="sqlite")
+        cls.dbpath, _name = cls.dbman.create_new_db_cli(cls.name, dbid="sqlite")
         cls.dbman.create_new_db_cli(cls.name, dbid="sqlite")
         with patch.dict(
             "os.environ",
@@ -112,6 +113,7 @@ def setUpClass(cls):
             cls.test_app = create_app()
         cls.test_app.config["TESTING"] = True
         cls.client = cls.test_app.test_client()
+        cls.tree = os.path.basename(cls.dbpath)
         with cls.test_app.app_context():
             user_db.create_all()
             for role in TEST_USERS:
@@ -148,10 +150,13 @@ def test_importers_no_data(self):
             data=None,
             headers=headers,
         )
-        assert rv.status_code == 500
+        assert rv.status_code == 400
+        assert "error" in rv.json
+        assert "empty" in rv.json["error"]["message"]
 
     def test_importers_example_data(self):
         """Test importing example.gramps."""
+        os.remove(os.path.join(self.dbpath, "sqlite.db"))
         example_db = ExampleDbInMemory()
         file_obj = io.BytesIO()
         with open(example_db.path, "rb") as f:
@@ -184,3 +189,28 @@ def test_importers_example_data(self):
         # everything doubled
         rv = check_success(self, f"{BASE_URL}/people/")
         assert len(rv) == 2 * 2157
+
+    def test_importers_example_data_quota(self):
+        """Test importing example.gramps with a quota."""
+        os.remove(os.path.join(self.dbpath, "sqlite.db"))
+        with self.test_app.app_context():
+            set_tree_quota(self.tree, quota_people=2000)
+        example_db = ExampleDbInMemory()
+        file_obj = io.BytesIO()
+        with open(example_db.path, "rb") as f:
+            file_obj.write(f.read())
+        file_obj.seek(0)
+        headers = fetch_header(self.client, role=ROLE_OWNER)
+        # database has no people
+        rv = check_success(self, f"{BASE_URL}/people/")
+        assert len(rv) == 0
+        rv = self.client.post(
+            f"{TEST_URL}gramps/file",
+            data=file_obj,
+            headers=headers,
+        )
+        assert rv.status_code == 405
+        rv = check_success(self, f"{BASE_URL}/people/")
+        assert len(rv) == 0
+        with self.test_app.app_context():
+            set_tree_quota(self.tree, quota_people=None)